/* ###
 * IP: GHIDRA
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *      http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package ghidra.app.plugin.core.disassembler;

import java.util.List;

import ghidra.app.plugin.core.analysis.AutoAnalysisManager;
import ghidra.app.services.*;
import ghidra.app.util.PseudoDisassembler;
import ghidra.app.util.importer.MessageLog;
import ghidra.framework.options.Options;
import ghidra.program.model.address.*;
import ghidra.program.model.data.*;
import ghidra.program.model.listing.*;
import ghidra.program.model.mem.*;
import ghidra.util.exception.CancelledException;
import ghidra.util.task.TaskMonitor;

/**
 * Check operand references to memory locations looking for Data
 */
public class AddressTableAnalyzer extends AbstractAnalyzer {
	private static String DESCRIPTION = "Analyzes undefined data for address tables.";

	private static final String OPTION_NAME_MIN_TABLE_SIZE = "Minimum Table Size";
	private static final String OPTION_NAME_TABLE_ALIGNMENT = "Table Alignment";
	private static final String OPTION_NAME_PTR_ALIGNMENT = "Pointer Alignment";
	private static final String OPTION_NAME_AUTO_LABEL_TABLE = "Auto Label Table";
	private static final String OPTION_NAME_MIN_POINTER_ADDR = "Minimum Pointer Address";
	private static final String OPTION_NAME_MAX_POINTER_DIFF = "Maxmimum Pointer Distance";
	private static final String OPTION_NAME_RELOCATION_GUIDE = "Relocation Table Guide";
	private static final String OPTION_NAME_ALLOW_OFFCUT_REFERENCES = "Allow Offcut References";

	protected static final String OPTION_NAME_CREATE_BOOKMARKS = "Create Analysis Bookmarks";
	private static final String OPTION_DESCRIPTION_CREATE_BOOKMARKS =
		"If checked, an analysis bookmark will be created at each location where an address table is constructed.";
	private static final boolean OPTION_DEFAULT_CREATE_BOOKMARKS_ENABLED = true;

	private static final String OPTION_DESCRIPTION_MIN_TABLE_SIZE =
		"The minimum number of consecutive addresses that constitute an address table";
	private static final String OPTION_DESCRIPTION_TABLE_ALIGNMENT =
		"Only check for tables aligned to this number of bytes.";
	private static final String OPTION_DESCRIPTION_PTR_ALIGNMENT =
		"Only check for ptr table entries aligned to this number of bytes.";
	private static final String OPTION_DESCRIPTION_AUTO_LABEL_TABLE =
		"Label the start of the table and each entry as part of the table.";
	private static final String OPTION_DESCRIPTION_MIN_POINTER_ADDR =
		"Minimum Address that any value is considered a pointer";
	private static final String OPTION_DESCRIPTION_MAX_POINTER_DIFF =
		"Maximum distance in bytes between pointers before the table is broken up.";
	private static final String OPTION_DESCRIPTION_RELOCATION_GUIDE =
		"Select this check box to use relocation table entries to guide pointer analysis.";
	private static final String OPTION_DESCRIPTION_ALLOW_OFFCUT_REFERENCES =
		"Allow table entries that are offcut references to defined data or instructions.";

	private static final int OPTION_DEFAULT_TABLE_ALIGNMENT = 4;
	private static final int OPTION_DEFAULT_PTR_ALIGNMENT = 1;
	private static final boolean OPTION_DEFAULT_AUTO_LABEL_TABLE = false;
	private static final boolean OPTION_DEFAULT_RELOCATION_GUIDE_ENABLED = true;
	private static final boolean OPTION_DEFAULT_ALLOW_OFFCUT_REFERENCES = false;
	private static final int OPTION_DEFAULT_MIN_POINTER_ADDR = 0x1024;
	private static final int OPTION_DEFAULT_MAX_POINTER_DIFF = 0xffffff;

	private int minimumTableSize = -1;
	private int tableAlignment = OPTION_DEFAULT_TABLE_ALIGNMENT;
	private int ptrAlignment = OPTION_DEFAULT_PTR_ALIGNMENT;
	private boolean autoLabelTable = OPTION_DEFAULT_AUTO_LABEL_TABLE;
	private boolean createBookmarksEnabled = OPTION_DEFAULT_CREATE_BOOKMARKS_ENABLED;

	private long minPointerAddress = OPTION_DEFAULT_MIN_POINTER_ADDR;
	private long maxPointerDistance = OPTION_DEFAULT_MAX_POINTER_DIFF;

	private boolean relocationGuideEnabled = OPTION_DEFAULT_RELOCATION_GUIDE_ENABLED;

	private boolean allowOffcutReferences = OPTION_DEFAULT_ALLOW_OFFCUT_REFERENCES;
	
	private static final String ADDRESS_TABLE_BOOKMARK_TYPENAME = "Address Table";

	// true if the processor uses Address low bit to refer to code
	private boolean processorHasLowBitCode = false;

	private long lastID;


	public AddressTableAnalyzer() {
		super("Create Address Tables", DESCRIPTION, AnalyzerType.BYTE_ANALYZER);
		setPriority(AnalysisPriority.DATA_TYPE_PROPOGATION.before());
		setSupportsOneTimeAnalysis();

		// The analyzer should be off by default (as stated in its description)
		setDefaultEnablement(false);
	}

	@Override
	public boolean canAnalyze(Program program) {
		// only analyze programs with address spaces > 16 bits
		int addrSize = program.getAddressFactory().getDefaultAddressSpace().getSize();
		// if processor uses the low bit to change instruction modes, then allow offcut refs
		processorHasLowBitCode = PseudoDisassembler.hasLowBitCodeModeInAddrValues(program);
		// don't align based on instruction start rules, data is important too.
		//   if do this, will miss, runs of ptrs to data/code/data/code/....
		// ptrAlignment = program.getLanguage().getInstructionAlignment();
		// if (processorHasLowBitCode) {
		// 	ptrAlignment = 1;
		// }
		return (addrSize == 32 || addrSize == 64);
	}

	@Override
	public boolean added(Program program, AddressSetView addrSet, TaskMonitor monitor,
			MessageLog log) throws CancelledException {
		AutoAnalysisManager mgr = AutoAnalysisManager.getAnalysisManager(program);

		// remove memory blocks that should not be searched
		addrSet = removeNonSearchableMemory(program, addrSet);

		// remove defined locations that can't be part of an address table
		// only do this the first time entering
		long id = mgr.getProgram().getCurrentTransactionInfo().getID();
		if (id != lastID) {
			lastID = id;
			addrSet = removeDefined(program, addrSet);
		}
		
		if (addrSet.isEmpty()) {
			return true;
		}

		long addrCount = program.getMemory().getNumAddresses();
		monitor.initialize(addrCount);
		monitor.setMessage("Analyze Address Tables");

		// iterate over addresses in the selectedmodule
		addrCount -= addrSet.getNumAddresses();

		Address minAddr = addrSet.getMinAddress();
		Address maxAddr = minAddr;

		// Iterate over all references within the new address set
		//   Evaluate each reference
		//

		AddressIterator addrIter = addrSet.getAddresses(true);

		boolean didTable = false;
		while (!didTable && addrIter.hasNext()) {
			monitor.checkCancelled();
			
			addrCount++;
			monitor.setProgress(addrCount);
			Address start = addrIter.next();
			maxAddr = start;

			if (start.getOffset() % tableAlignment != 0) {
				continue;
			}

			if ((addrCount % 2048) == 1) {
				monitor.setMessage("Analyze Tables " + start);
			}

			AddressTable addressTable =
				AddressTable.getEntry(program, start, monitor, true, minimumTableSize, ptrAlignment,
					0, AddressTable.MINIMUM_SAFE_ADDRESS, relocationGuideEnabled);
			
			if (addressTable == null) {
				continue;
			}
			
			didTable = processAddressTable(addressTable, program, mgr, monitor);
			
			// Assumes that the entire table found was processed
			long tableByteLen = addressTable.getByteLength();
			addrCount += tableByteLen;
			
			// update maximum addresses consumed by table
			try {
				maxAddr = start.addNoWrap(tableByteLen-1);
				addrIter = addrSet.getAddresses(maxAddr.addNoWrap(1), true);
			}
			catch (AddressOverflowException e) {
				// ignore, just keep using current iterator
			}
		} // end of while (addrIt.hasNext())

		// Set up a one time analysis to get us back into here if
		//   there are still addresses on the set
		//
		AddressSet set =
			addrSet.subtract(program.getAddressFactory().getAddressSet(minAddr, maxAddr));
		if (!set.isEmpty()) {
			mgr.scheduleOneTimeAnalysis(this, set);
		}
		return true;
	}

	/**
	 * Process the table of addresses.  This could create multiple address tables if the
	 * table has an entry that breaks the table.
	 * 
	 * Note: current algorithm processes the entire table, if this changes, code that
	 *       assumes this must be re-factored to know how many bytes were consumed
	 * 
	 * @param addressTable address table
	 * @param program program
	 * @param mgr autoAnalysis manager
	 * @param monitor monitor for checking if table creation canceled
	 * @return true if a table was actually made, false otherwise
	 * @throws CancelledException if use cancels
	 */
	private boolean processAddressTable(AddressTable addressTable, Program program, AutoAnalysisManager mgr,
			TaskMonitor monitor) throws CancelledException {
		boolean didTable = false;
		
		while (addressTable != null) {
			Address startTableAddr = addressTable.getTopAddress();

			int tableLen = checkTable(addressTable, program, monitor);
			
			// check if there is already an address table bookmark here
			Bookmark bookmark = program.getBookmarkManager()
					.getBookmark(addressTable.getTopAddress(), BookmarkType.ANALYSIS,
						ADDRESS_TABLE_BOOKMARK_TYPENAME);

			// if table too small, or bookmark already here
			// - skip the bad tableLen entries
			// - continue with any additional entries in the table
			if (bookmark != null || tableLen < minimumTableSize) {
				// There might be more to the table. Get a new smaller table.
				// This will skip one element in the table, that is assumed to have
				// broken the table up, thus (tableLen+1)
				//
				// table entry will be null if no more table entries beyond tableLen+1
				addressTable = addressTable.newRemainingAddressTable(tableLen+1);
				continue;
			}

			// make the table
			addressTable.makeTable(program, 0, tableLen - 1, autoLabelTable, false);

			int tableByteLen = addressTable.getByteLength(0, tableLen - 1, false);
			Address endTableAddr = startTableAddr.add(tableByteLen-1);
			mgr.codeDefined(new AddressSet(startTableAddr, endTableAddr));

			// put info bookmark in
			if (createBookmarksEnabled) {
				program.getBookmarkManager()
						.setBookmark(addressTable.getTopAddress(), BookmarkType.ANALYSIS,
							ADDRESS_TABLE_BOOKMARK_TYPENAME, "Address table[" + tableLen + "] created");
			}

			// if all are valid code, disassemble
			List<Address> validCodeList = addressTable.getFunctionEntries(program, 0);
			if (validCodeList != null &&
				validCodeList.size() >= addressTable.getNumberAddressEntries()) {
				AddressSet validCodeSet = new AddressSet();
				AddressSet validFuncSet = new AddressSet();
				for (Address addr : validCodeList) {
					// set target context correctly. Target address will get
					// aligned in DisassembleCmd
					PseudoDisassembler.setTargetContextForDisassembly(program, addr);

					// even though they are valid code, don't do them if
					// there is already code there.
					if (program.getListing().getCodeUnitContaining(addr) == null) {
						validCodeSet.addRange(addr, addr);
					}
					// For Now, Never make functions from address tables,
					// Could later be an option
					// if (pdis.isValidSubroutine(addr, true)) {
					// validFuncSet.addRange(addr, addr);
					// }
				}
				// disassemble valid code
				if (!validCodeSet.isEmpty()) {
					mgr.disassemble(validCodeSet,
						AnalysisPriority.DATA_TYPE_PROPOGATION.before());
				}
				// if valid functions, schedule much later, so switch table analysis can pick it up.
				if (!validFuncSet.isEmpty()) {
					//  For Now, Never make functions from address tables, Could later be an option
					//		mgr.createFunction(validFuncSet, true,
					//		AnalysisPriority.DATA_TYPE_PROPOGATION.getNext());
				}
			}
			
			// There might be more to the table. Get a new smaller table.
			// This will skip one element in the table, that is assumed to have
			// broken the table up, thus (tableLen+1)
			//
			// table entry will be null if no more table entries beyond tableLen+1
			addressTable = addressTable.newRemainingAddressTable(tableLen+1);
			didTable = true;
		}
		
		return didTable;
	}

	/**
	 * Remove any addresses from set where instructions or defined data exists that would
	 * get filtered during address table creation.
	 */
	private AddressSetView removeDefined(Program program, AddressSetView addrSet) {
		AddressSet subSet = new AddressSet();

		// find defined data that should be removed
		DataIterator definedData = program.getListing().getDefinedData(addrSet, true);
		for (Data data : definedData) {
			DataType dataType = data.getDataType();			
			if (dataType instanceof Undefined) {
				continue;
			}
			
			if (data.isPointer()) {
				continue;
			}

			// anything else, add to subset to be removed
			subSet.add(data.getMinAddress(), data.getMaxAddress());
		}
		addrSet = addrSet.subtract(subSet);
		
		// Find defined instructions that should be removed
		subSet = new AddressSet();
		InstructionIterator instructions = program.getListing().getInstructions(addrSet, true);
		for (Instruction instruction : instructions) {
			subSet.add(instruction.getMinAddress(), instruction.getMaxAddress());
		}
		addrSet = addrSet.subtract(subSet);
		
		return addrSet;
	}

	private AddressSetView removeNonSearchableMemory(Program program, AddressSetView addrSet) {
		// get rid of any non-initialized blocks

		addrSet = addrSet.intersect(program.getMemory().getLoadedAndInitializedAddressSet());

		MemoryBlock[] blocks = program.getMemory().getBlocks();

		// get rid of any blocks that have empty attributes.
		//   Would be better to get rid of any blocks in a special RAWBytes address space, not overlays.
		//
		AddressSet badBlocks = new AddressSet();
		for (MemoryBlock memoryBlock : blocks) {
			if (memoryBlock.isWrite() || memoryBlock.isRead() || memoryBlock.isExecute()) {
				continue;
			}

			badBlocks.addRange(memoryBlock.getStart(), memoryBlock.getEnd());
		}

		addrSet = addrSet.subtract(badBlocks);

		return addrSet;
	}

	/**
	 * Check the table for consistency and return the number of good entries before a bad
	 * entry is found.
	 * <ul>
	 * <li> check possible strings</li>
	 * <li> pointers that jump all over memory</li>
	 * <li> pointers to offcut instructions</li>
	 * <li> allow offcut strings</li>
	 * </ul>
	 * 
	 * @param tableEntry address table to check
	 * @param program program
	 * @return number of entries in table before hitting an inconsistent entry
	 * @throws CancelledException if user cancels
	 */
	private int checkTable(AddressTable tableEntry, Program program, TaskMonitor monitor) throws CancelledException {
		// search for unicode strings first.
		//   don't create an address table that overlaps a unicode string.
		AddressSetView addrSet = tableEntry.getTableBody();
		AddressSet possibleStrings = findPossibleStrings(program, addrSet, monitor);

		// trim from the first offcut entry on
		Address start = tableEntry.getTopAddress();
		int tableLen = tableEntry.getNumberAddressEntries();
		Address addrs[] = tableEntry.getTableElements();
		for (int i = 0; i < tableLen; i++) {
			monitor.checkCancelled();
			Address tableEntryAddr = start.add(i * 4);
			Address targetAddr = addrs[i];
			if (possibleStrings.contains(tableEntryAddr)) {
				return i;
			}
			if (possibleStrings.contains(tableEntryAddr.add(3))) {
				return i;
			}
			if ((tableEntryAddr.getOffset() > 0) &&
				tableEntryAddr.getOffset() < minPointerAddress) {
				return i;
			}
			// check that the table entries are not all over the place
			if (i > 0) {
				long diff = addrs[i - 1].subtract(addrs[i]);
				diff = Math.abs(diff);
				if (diff > maxPointerDistance) {
					return i;
				}
			}
			CodeUnit cu = program.getListing().getCodeUnitContaining(targetAddr);
			if (cu == null) {
				continue;
			}

			boolean atStartOfCU = cu.getMinAddress().equals(targetAddr);
			if (!allowOffcutReferences && !atStartOfCU) {
				// if the processor uses low bit to reference instructions
				//  allow offcut to an instruction by 1
				if (!processorHasLowBitCode || !(cu instanceof Instruction)) {
					return i;
				}
			}
		}
		return tableLen;
	}

	private AddressSet findPossibleStrings(Program program, AddressSetView addrSet, TaskMonitor monitor) throws CancelledException {
		AddressSet possibleStrSet = new AddressSet();
		Memory memory = program.getMemory();

		AddressIterator addrIter = addrSet.getAddresses(true);
		long maxBytes = addrSet.getNumAddresses();

		MemoryBufferImpl buffer = new MemoryBufferImpl(memory, addrSet.getMinAddress(),
			(int) (maxBytes > 1024 ? 1024 : maxBytes));

		while (addrIter.hasNext()) {
			Address start = addrIter.next();
			
			monitor.checkCancelled();

			// skip over anything that smells like a unicode string
			//
			int strLen = getWStrLen(buffer, start, (int) (maxBytes / 2));
			if (strLen > 4) {
				int numBytes = strLen * 2;
				addrIter = skipBytes(addrIter, addrSet, start, numBytes);
				possibleStrSet.addRange(start, start.add(numBytes));
				continue;
			}
		}
		return possibleStrSet;
	}

	private AddressIterator skipBytes(AddressIterator iter, AddressSetView addrSet, Address start,
			int numBytes) {
		try {
			start = start.addNoWrap(numBytes);
		}
		catch (AddressOverflowException e) {
			return iter;
		}
		iter = addrSet.getAddresses(start, true);
		return iter;
	}

	/**
	 * getWStrLen
	 * @param ad = address where unicode string is supposed to begin
	 * @return number of unicode chars in string, -1 if not
	 * a unicode string.  NOTE: Only English strings are considered.
	 *
	 */
	private int getWStrLen(MemoryBufferImpl memory, Address ad, int max) {
		int i = 0;
		memory.setPosition(ad);
		try {
			for (i = 0; i < max; i++) {
				short value = memory.getShort(2 * i);
				if (value == 0) {
					return i + 1;
				}
				// allow tab, carriage return, and linefeed
				if (value != 0x09 && value != 0x0a && value != 0x0d &&
					(value < 0x20 || value >= 0x7f)) {
					return i;
				}
			}
		}
		catch (MemoryAccessException e) {
			return i;
		}
		return i;
	}

	@Override
	public boolean removed(Program program, AddressSetView set, TaskMonitor monitor,
			MessageLog log) {
		return false;
	}

	@Override
	public boolean getDefaultEnablement(Program program) {

		// a probability calculation to decide to enable based on memory block bytes

		if (minimumTableSize == -1) {
			calculateMinimumTableSize(program);
		}

		return minimumTableSize != AddressTable.TOO_MANY_ENTRIES;
	}

	private void calculateMinimumTableSize(Program program) {
		minimumTableSize =
			AddressTable.getThresholdRunOfValidPointers(program, AddressTable.BILLION_CASES);

		if (minimumTableSize < 2) {
			minimumTableSize = 2;
		}
	}

	@Override
	public void registerOptions(Options options, Program program) {
		options.registerOption(OPTION_NAME_MIN_TABLE_SIZE, minimumTableSize, null,
			OPTION_DESCRIPTION_MIN_TABLE_SIZE);

		options.registerOption(OPTION_NAME_TABLE_ALIGNMENT, tableAlignment, null,
			OPTION_DESCRIPTION_TABLE_ALIGNMENT);

		options.registerOption(OPTION_NAME_PTR_ALIGNMENT, ptrAlignment, null,
			OPTION_DESCRIPTION_PTR_ALIGNMENT);

		options.registerOption(OPTION_NAME_AUTO_LABEL_TABLE, autoLabelTable, null,
			OPTION_DESCRIPTION_AUTO_LABEL_TABLE);

		options.registerOption(OPTION_NAME_RELOCATION_GUIDE, relocationGuideEnabled, null,
			OPTION_DESCRIPTION_RELOCATION_GUIDE);

		options.registerOption(OPTION_NAME_ALLOW_OFFCUT_REFERENCES, allowOffcutReferences, null,
			OPTION_DESCRIPTION_ALLOW_OFFCUT_REFERENCES);

		options.registerOption(OPTION_NAME_MIN_POINTER_ADDR, minPointerAddress, null,
			OPTION_DESCRIPTION_MIN_POINTER_ADDR);

		options.registerOption(OPTION_NAME_MAX_POINTER_DIFF, maxPointerDistance, null,
			OPTION_DESCRIPTION_MAX_POINTER_DIFF);

		options.registerOption(OPTION_NAME_CREATE_BOOKMARKS, createBookmarksEnabled, null,
			OPTION_DESCRIPTION_CREATE_BOOKMARKS);

		optionsChanged(options, program);
	}

	@Override
	public void optionsChanged(Options options, Program program) {

		if (minimumTableSize == -1) {
			calculateMinimumTableSize(program);
		}

		minimumTableSize = options.getInt(OPTION_NAME_MIN_TABLE_SIZE, minimumTableSize);

		tableAlignment = options.getInt(OPTION_NAME_TABLE_ALIGNMENT, tableAlignment);

		ptrAlignment = options.getInt(OPTION_NAME_PTR_ALIGNMENT, ptrAlignment);
		autoLabelTable = options.getBoolean(OPTION_NAME_AUTO_LABEL_TABLE, autoLabelTable);

		relocationGuideEnabled =
			options.getBoolean(OPTION_NAME_RELOCATION_GUIDE, relocationGuideEnabled);

		allowOffcutReferences =
			options.getBoolean(OPTION_NAME_ALLOW_OFFCUT_REFERENCES, allowOffcutReferences);

		minPointerAddress = options.getLong(OPTION_NAME_MIN_POINTER_ADDR, minPointerAddress);

		maxPointerDistance = options.getLong(OPTION_NAME_MAX_POINTER_DIFF, maxPointerDistance);

		createBookmarksEnabled =
			options.getBoolean(OPTION_NAME_CREATE_BOOKMARKS, createBookmarksEnabled);

	}
}
