content: /* ###
   * 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.
   */
  //UEFI helper script
  
  import firmware.common.UUIDUtils;
  import firmware.uefi_te.TELoader;
  import firmware.uefi_te.TerseExecutableHeader;
  import ghidra.app.decompiler.*;
  import ghidra.app.plugin.core.datamgr.archive.DuplicateIdException;
  import ghidra.app.plugin.core.analysis.AutoAnalysisManager;
  import ghidra.app.script.GhidraScript;
  import ghidra.app.util.bin.BinaryReader;
  import ghidra.app.util.bin.ByteArrayProvider;
  import ghidra.app.util.bin.format.pe.MachineConstants;
  import ghidra.app.util.bin.format.pe.NTHeader;
  import ghidra.app.util.bin.format.pe.PeSubsystem;
  import ghidra.app.util.bin.format.pe.PortableExecutable;
  import ghidra.app.util.opinion.PeLoader;
  import ghidra.app.services.DataTypeManagerService;
  import ghidra.program.model.address.Address;
  import ghidra.program.model.data.*;
  import ghidra.program.model.listing.*;
  import ghidra.program.model.listing.Function.FunctionUpdateType;
  import ghidra.program.model.mem.MemoryBlock;
  import ghidra.program.model.pcode.*;
  import ghidra.program.model.symbol.SourceType;
  import ghidra.util.Msg;
  import ghidra.util.exception.DuplicateNameException;
  import ghidra.util.exception.InvalidInputException;
  
  import java.io.File;
  import java.io.IOException;
  import java.util.*;
  
  /**
   * UEFI helper script
   */
  public class UEFIHelper extends GhidraScript {
  	private DataTypeManager uefiTypeManager;
  
  	/**
  	 * Loads the specified file from the plugin's data directory.
  	 *
  	 * @param name the specified filename
  	 */
  	private File loadDataFile(String name) throws IOException {
  		return new File(new File(
  				sourceFile.getParentFile().getParentFile().getFile(true), "data"), name);
  	}
  
  	/**
  	 * Loads the specified data type library and returns a corresponding DataTypeManager.
  	 *
  	 * @param name the name of the data type library
  	 * @return     the DataTypeManager for the specified data type library
  	 */
  	private DataTypeManager loadDataTypeLibrary(String name) throws DuplicateIdException,
  			IOException {
  		// Check if the data type library was already loaded.
  		AutoAnalysisManager aam = AutoAnalysisManager.getAnalysisManager(currentProgram);
  		DataTypeManagerService service = aam.getDataTypeManagerService();
  		DataTypeManager[] managers = service.getDataTypeManagers();
  		for (DataTypeManager manager : managers) {
  			if (manager.getName().equals(name)) {
  				return manager;
  			}
  		}
  
  		// Load the data type library from the plugin's data directory.
  		return service.openDataTypeArchive(name);
  	}
  
  	/**
  	 * Updates a specified Function with the signature from a specified FunctionDefinition.
  	 *
  	 * @param function   the specified Function
  	 * @param definition the specified FunctionDefinition
  	 */
  	private void updateFunctionSignature(Function function,
  			FunctionDefinition definition) throws DuplicateNameException, InvalidInputException {
  		// Build the list of parameters.
  		ArrayList<ParameterImpl> parameters = new ArrayList<>();
  		ParameterDefinition[] parameterDefinitions = definition.getArguments();
  		for (ParameterDefinition parameterDefinition : parameterDefinitions) {
  			parameters.add(new ParameterImpl(parameterDefinition.getName(),
  					parameterDefinition.getDataType(), currentProgram));
  		}
  
  		// Build the return type.
  		ReturnParameterImpl returnType = new ReturnParameterImpl(definition.getReturnType(),
  				currentProgram);
  
  		// Update the function with the generated parameter list and return type.
  		function.setName(definition.getName(), SourceType.DEFAULT);
  		function.updateFunction(null, returnType, parameters,
  				FunctionUpdateType.DYNAMIC_STORAGE_FORMAL_PARAMS, false, SourceType.DEFAULT);
  	}
  
  	/**
  	 * Defines a specified DataType at a specified Address. Existing data or instructions that
  	 * would conflict with the new DataType will be removed.
  	 *
  	 * @param address  the specified Address
  	 * @param dataType the specified DataType
  	 * @param name     an optional name to use for the data type's label
  	 * @param comment  an optional comment
  	 */
  	private void defineData(Address address, DataType dataType, String name,
  			String comment) throws Exception {
  		// Remove any existing data or instructions that would overlap with this definition.
  		for (int i = 0; i < dataType.getLength(); i++) {
  			Address currentAddress = address.add(i);
  			Data existingData = getDataAt(currentAddress);
  			if (existingData != null) {
  				removeData(existingData);
  			} else {
  				Instruction existingInstruction = getInstructionAt(currentAddress);
  				if (existingInstruction != null) {
  					removeInstruction(existingInstruction);
  				}
  			}
  		}
  
  		// Apply the data type definition and create a label.
  		createData(address, dataType);
  		if (name != null) {
  			createLabel(address, name, true);
  		}
  
  		// Add a comment (if specified).
  		if (comment != null) {
  			setPlateComment(address, comment);
  		}
  	}
  
  	/**
  	 * Searches for assignment statements to global variables and propagates the source data type
  	 * to the global variable.
  	 *
  	 * @param root the root ClangTokenGroup to search in
  	 */
  	private void propagateGlobalVariables(ClangTokenGroup root) throws Exception {
  		for (int i = 0; i < root.numChildren(); i++) {
  			ClangNode childNode = root.Child(i);
  			if (childNode instanceof ClangTokenGroup) {
  				// Assignment statements begin with a destination variable.
  				if (childNode.numChildren() > 3 &&
  					childNode.Child(0) instanceof ClangVariableToken) {
  					HighVariable destination =
  							((ClangVariableToken) childNode.Child(0)).getHighVariable();
  					ClangNode sourceNode = childNode.Child(childNode.numChildren() - 1);
  
  					// Verify that there is a source variable and that the destination is a global
  					// variable.
  					if (destination instanceof HighGlobal &&
  						(sourceNode instanceof ClangVariableToken ||
  						sourceNode instanceof ClangFieldToken)) {
  						Address globalAddress = destination.getRepresentative().getAddress();
  
  						// Retrieve the source data type.
  						DataType sourceDataType = null;
  						if (sourceNode instanceof ClangVariableToken) {
  							ClangVariableToken sourceToken = (ClangVariableToken) sourceNode;
  							if (sourceToken.getHighVariable() != null) {
  								sourceDataType = sourceToken.getHighVariable().getDataType();
  							}
  						} else {
  							ClangFieldToken sourceToken = (ClangFieldToken) sourceNode;
  							if (sourceToken instanceof Structure) {
  								Structure structureType = (Structure) sourceToken.getDataType();
  								sourceDataType = structureType.getDataTypeAt(
  										sourceToken.getOffset()).getDataType();
  							}
  						}
  
  						// Apply label names for certain global variables.
  						// These are defined by the UefiBootServicesTableLib and
  						// UefiRuntimeServicesTableLib libraries in EDK2.
  						if (sourceDataType != null) {
  							String name = null;
  							switch (sourceDataType.getName()) {
  								case "EFI_BOOT_SERVICES *":
  									name = "gBS";
  									break;
  								case "EFI_HANDLE":
  									name = "gImageHandle";
  									break;
  								case "EFI_RUNTIME_SERVICES *":
  									name = "gRT";
  									break;
  								case "EFI_SYSTEM_TABLE *":
  									name = "gST";
  									break;
  								default:
  									// Avoid propagating other data types.
  									continue;
  							}
  
  							// Update the global variable with the source data type.
  							println("Found global assignment: " + childNode.toString());
  							defineData(globalAddress, sourceDataType, name, null);
  							printf("%s> - Applied %s data type to 0x%s\n", getScriptName(),
  									destination.getDataType().getName(),
  									globalAddress.toString().toUpperCase());
  						}
  					}
  				} else {
  					propagateGlobalVariables((ClangTokenGroup) childNode);
  				}
  			}
  		}
  	}
  
  	/**
  	 * Propagates data types to functions called from a specified ClangTokenGroup.
  	 *
  	 * @param decompiler the DecompInterface
  	 * @param options    the specified options for the previously specified DecompInterface
  	 * @param root       the specified ClangTokenGroup
  	 */
  	private void propagateFunctionParameters(DecompInterface decompiler, DecompileOptions options,
  			ClangTokenGroup root) throws Exception {
  		for (int i = 0; i < root.numChildren(); i++) {
  			ClangNode childNode = root.Child(i);
  			if (childNode instanceof ClangTokenGroup && childNode.numChildren() > 1) {
  				// Look for a function call in this token group.
  				ClangFuncNameToken funcToken = null;
  				int funcTokenIndex = 0;
  				while (funcTokenIndex < childNode.numChildren()) {
  					if (childNode.Child(funcTokenIndex) instanceof ClangFuncNameToken) {
  						funcToken = (ClangFuncNameToken) childNode.Child(funcTokenIndex);
  						if (funcToken.getPcodeOp() != null &&
  							funcToken.getPcodeOp().getOpcode() == PcodeOp.CALL) {
  							break;
  						} else {
  							funcToken = null;
  						}
  					}
  
  					funcTokenIndex++;
  				}
  
  				if (funcToken != null) {
  					Msg.debug(this, "Found function name token: " + funcToken.toString());
  					Function function = getFunctionAt(
  							funcToken.getPcodeOp().getInput(0).getAddress());
  					ArrayList<ParameterImpl> parameters = new ArrayList<>();
  					Parameter[] originalParameters = function.getParameters();
  					int parameterNumber = 0;
  					// Parse each parameter token and update the function signature with each
  					// parameter's data type.
  					for (int j = 1; j < childNode.numChildren() &&
  							parameterNumber < originalParameters.length; j++) {
  						if (childNode.Child(j) instanceof ClangVariableToken) {
  							ClangVariableToken variableToken =
  									(ClangVariableToken) childNode.Child(j);
  							if (variableToken.getHighVariable() != null) {
  								DataType parameterType =
  										variableToken.getHighVariable().getDataType();
  								Parameter originalParameter =
  										originalParameters[parameterNumber++];
  								Msg.debug(this, "Found parameter: " +
  										variableToken.getHighVariable().getName() + " (" +
  										parameterType.getName() + ')');
  								switch (parameterType.getName()) {
  									case "EFI_BOOT_SERVICES *":
  									case "EFI_HANDLE":
  									case "EFI_RUNTIME_SERVICES *":
  									case "EFI_SYSTEM_TABLE *":
  										// Replace the data type for the current parameter.
  										parameters.add(new ParameterImpl(originalParameter.getName(),
  												parameterType, getCurrentProgram()));
  										break;
  									default:
  										// Reuse the original data type for the current parameter.
  										parameters.add(new ParameterImpl(originalParameter.getName(),
  												originalParameter.getDataType(),
  												originalParameter.getVariableStorage(),
  												getCurrentProgram()));
  								}
  							}
  						}
  					}
  
  					if (parameters.size() == originalParameters.length) {
  						function.updateFunction(null, null, parameters,
  								FunctionUpdateType.DYNAMIC_STORAGE_FORMAL_PARAMS, false,
  								SourceType.DEFAULT);
  						println("Updated " + function.getName() + " function signature");
  					} else {
  						Msg.error(this, "Failed to parse " + function.getName() + " parameters");
  					}
  				} else {
  					propagateFunctionParameters(decompiler, options, (ClangTokenGroup) childNode);
  				}
  			}
  		}
  	}
  
  	/**
  	 * Defines the signature for the current program's entry point function.
  	 *
  	 * @param entryPointAddress the address of the entry point function
  	 */
  	private void defineEntryPoint(Address entryPointAddress) throws Exception {
  		Function entryPoint = getFunctionAt(entryPointAddress);
  		if (entryPoint == null) {
  			// Create the entry point function if it wasn't already defined.
  			entryPoint = createFunction(entryPointAddress, "_ModuleEntryPoint");
  		}
  
  		println("Found entry point function at 0x" + entryPointAddress.toString().toUpperCase());
  
  		// TODO: Pick correct entry point signature based off module type (PEI vs DXE, etc)
  		FunctionDefinition entryPointDefinition = (FunctionDefinition) uefiTypeManager.getDataType(
  				"/UefiApplicationEntryPoint.h/functions/_ModuleEntryPoint");
  		updateFunctionSignature(entryPoint, entryPointDefinition);
  
  		// Decompile the entry point function.
  		DecompInterface decompiler = new DecompInterface();
  		DecompileOptions options = new DecompileOptions();
  		options.grabFromProgram(currentProgram);
  		decompiler.setOptions(options);
  		decompiler.setSimplificationStyle("decompile");
  		decompiler.openProgram(currentProgram);
  		DecompileResults results = decompiler.decompileFunction(entryPoint,
  				options.getDefaultTimeout(), getMonitor());
  
  		// Propagate global types in the entry point (e.g. gBS/etc).
  		ClangTokenGroup tokenGroup = results.getCCodeMarkup();
  		println("Searching for global assignments...");
  		propagateGlobalVariables(tokenGroup);
  
  		// Propagate entry point parameters to called functions.
  		println("Propagating types to called functions...");
  		propagateFunctionParameters(decompiler, options, tokenGroup);
  
  		// Propagate global types in functions called by the entry point.
  		Set<Function> functions = entryPoint.getCalledFunctions(getMonitor());
  		for (Function function : functions) {
  			println("Searching for global assignments in " + function.getName() + "...");
  			results = decompiler.decompileFunction(function,
  					options.getDefaultTimeout(), getMonitor());
  			propagateGlobalVariables(results.getCCodeMarkup());
  		}
  	}
  
  	/**
  	 * Searches for known GUIDs in the specified MemoryBlock and applies the EFI_GUID data type.
  	 *
  	 * @param block the specified MemoryBlock
  	 */
  	private void defineGUIDs(MemoryBlock block) throws Exception {
  		println("Searching for GUIDs...");
  
  		// Read the contents .data segment.
  		byte[] blockBytes = new byte[(int) block.getSize()];
  		block.getBytes(block.getStart(), blockBytes);
  		BinaryReader reader = new BinaryReader(new ByteArrayProvider(blockBytes), true);
  
  		// Find the EFI_GUID data type.
  		DataType efiGuidType = uefiTypeManager.getDataType("/UefiBaseType.h/EFI_GUID");
  
  		// Search for known GUIDs in the GUID database.
  		long index = 0;
  		while (index < block.getSize() - efiGuidType.getLength()) {
  			UUID uuid = UUIDUtils.fromBinaryReader(reader);
  			if (uuid.equals(UUID.fromString("00000000-0000-0000-0000-000000000000"))) {
  				index += efiGuidType.getLength();
  			} else {
  				if (UUIDUtils.dbContains(uuid)) {
  					Address guidAddress = block.getStart().add(index);
  					println("GUID: Found " + UUIDUtils.getName(uuid) + " (" + uuid.toString() +
  							") at 0x" + guidAddress.toString().toUpperCase());
  
  					// Apply the EFI_GUID data type for the GUID we found.
  					defineData(guidAddress, efiGuidType, UUIDUtils.getName(uuid),
  							uuid.toString());
  					index += efiGuidType.getLength();
  				} else {
  					index += 1;
  				}
  			}
  
  			reader.setPointerIndex(index);
  		}
  	}
  
  	@Override
  	public void run() throws Exception {
  		println("UEFIHelper - UEFI helper script");
  
  		// Make sure the current program is a UEFI executable.
  		short machine;
  		int subsystem;
  		Address entryPointAddress = currentProgram.getImageBase();
  		if (currentProgram.getExecutableFormat().equals(PeLoader.PE_NAME)) {
  			// Parse the PE headers.
  			MemoryBlock peBlock = currentProgram.getMemory().getBlock(PeLoader.HEADERS);
  			byte[] blockBytes = new byte[(int) peBlock.getSize()];
  			peBlock.getBytes(peBlock.getStart(), blockBytes);
  			BinaryReader reader = new BinaryReader(new ByteArrayProvider(blockBytes), true);
  			int ntHeaderOffset = reader.readInt(0x3C);
  			NTHeader ntHeader = new NTHeader(reader, ntHeaderOffset,
  					PortableExecutable.SectionLayout.FILE, false, false);
  			println("Loaded Portable Executable");
  
  			machine = ntHeader.getFileHeader().getMachine();
  			subsystem = ntHeader.getOptionalHeader().getSubsystem();
  			entryPointAddress = entryPointAddress.add(
  					ntHeader.getOptionalHeader().getAddressOfEntryPoint());
  		} else if (currentProgram.getExecutableFormat().equals(TELoader.TE_NAME)) {
  			// Parse the TE header.
  			MemoryBlock teBlock = currentProgram.getMemory().getBlock(TELoader.HEADERS);
  			byte[] blockBytes = new byte[(int) teBlock.getSize()];
  			teBlock.getBytes(teBlock.getStart(), blockBytes);
  			BinaryReader reader = new BinaryReader(new ByteArrayProvider(blockBytes), true);
  			TerseExecutableHeader teHeader = new TerseExecutableHeader(reader);
  			println("Loaded Terse Executable");
  
  			machine = teHeader.getMachineType();
  			subsystem = teHeader.getSubsystem();
  			entryPointAddress = entryPointAddress.add(teHeader.getEntryPointAddress());
  		} else {
  			Msg.showError(this, null, "UEFIHelper", "Current program is not a PE/TE binary (" +
  					currentProgram.getExecutableFormat() + ')');
  			return;
  		}
  
  		// TODO: When https://github.com/NationalSecurityAgency/ghidra/pull/501 is merged, switch
  		// to checking for the existence of the UEFI property in the current program's options.
  		if (subsystem < PeSubsystem.IMAGE_SUBSYSTEM_EFI_APPLICATION.ordinal() &&
  			subsystem > PeSubsystem.IMAGE_SUBSYSTEM_EFI_ROM.ordinal()) {
  			Msg.showError(this, null, "UEFIHelper",
  					"Current program is not a UEFI binary (subsystem = " + subsystem + ')');
  			return;
  		}
  
  		String libraryName;
  		switch (machine) {
  			case MachineConstants.IMAGE_FILE_MACHINE_AMD64:
  				libraryName = "uefi_x64.gdt";
  				break;
  			case MachineConstants.IMAGE_FILE_MACHINE_ARM:
  				libraryName = "uefi_arm.gdt";
  				break;
  			case MachineConstants.IMAGE_FILE_MACHINE_ARM64:
  				libraryName = "uefi_aarch64.gdt";
  				break;
  			case MachineConstants.IMAGE_FILE_MACHINE_I386:
  				libraryName = "uefi_ia32.gdt";
  				break;
  			default:
  				Msg.showError(this, null, "UEFIHelper",
  						"No data type library present for current program's machine type (" +
  						machine + ')');
  				return;
  		}
  
  		// Load the UEFI data type library.
  		uefiTypeManager = loadDataTypeLibrary(libraryName);
  
  		// Fix the entry point function signature and propagate the parameter types.
  		defineEntryPoint(entryPointAddress);
  
  		// Search for known GUIDs in the program's .data and .text segments.
  		MemoryBlock dataBlock = currentProgram.getMemory().getBlock(".data");
  		if (dataBlock != null) {
  			defineGUIDs(dataBlock);
  		}
  
  		MemoryBlock textBlock = currentProgram.getMemory().getBlock(".text");
  		if (textBlock != null) {
  			defineGUIDs(textBlock);
  		}
  
  		// Potential additions / TODO(?)
  		//  - Locate GUIDs in calls to EFI boot/runtime services and apply the EFI_GUID data type
  		//  - Maintain a database for protocol GUID->data type name mappings and apply data types
  		//    for protocol interfaces (in calls to LocateProtocol/etc)
  	}
  }
  