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.
   */
  
  package firmware.uefi_te;
  
  import ghidra.app.util.Option;
  import ghidra.app.util.bin.BinaryReader;
  import ghidra.app.util.bin.ByteProvider;
  import ghidra.app.util.bin.format.pe.MachineConstants;
  import ghidra.app.util.bin.format.pe.PeSubsystem;
  import ghidra.app.util.importer.MessageLog;
  import ghidra.app.util.opinion.AbstractLibrarySupportLoader;
  import ghidra.app.util.opinion.LoadSpec;
  import ghidra.program.flatapi.FlatProgramAPI;
  import ghidra.program.model.address.Address;
  import ghidra.program.model.address.AddressOutOfBoundsException;
  import ghidra.program.model.address.AddressOverflowException;
  import ghidra.program.model.lang.LanguageCompilerSpecPair;
  import ghidra.program.model.listing.Program;
  import ghidra.program.model.mem.MemoryBlock;
  import ghidra.util.Msg;
  import ghidra.util.task.TaskMonitor;
  
  import java.io.IOException;
  import java.io.InputStream;
  import java.util.ArrayList;
  import java.util.Collection;
  import java.util.List;
  
  /**
   * Loader for Terse Executable (TE) binaries. Terse Executables are simplified PE/COFF executables,
   * used in the UEFI PI stage.
   */
  public class TELoader extends AbstractLibrarySupportLoader {
  	public static final String HEADERS = "Headers";
  	public static final String TE_NAME = "Terse Executable (TE)";
  
  	@Override
  	public Collection<LoadSpec> findSupportedLoadSpecs(ByteProvider provider) {
  		ArrayList<LoadSpec> loadSpecs = new ArrayList<>();
  		BinaryReader reader = new BinaryReader(provider, true);
  		TerseExecutableHeader header = null;
  		try {
  			header = new TerseExecutableHeader(reader);
  		} catch (IOException e) {
  			return loadSpecs;
  		}
  
  		switch (header.getMachineType()) {
  			case MachineConstants.IMAGE_FILE_MACHINE_I386:
  				loadSpecs.add(new LoadSpec(this, header.getImageBase(),
  						new LanguageCompilerSpecPair("x86:LE:32:default", "windows"), true));
  				break;
  			case MachineConstants.IMAGE_FILE_MACHINE_AMD64:
  				loadSpecs.add(new LoadSpec(this, header.getImageBase(),
  						new LanguageCompilerSpecPair("x86:LE:64:default", "windows"), true));
  				break;
  			case MachineConstants.IMAGE_FILE_MACHINE_ARM64:
  				loadSpecs.add(new LoadSpec(this, header.getImageBase(),
  						new LanguageCompilerSpecPair("AARCH64:LE:64:v8A", "windows"), true));
  				break;
  		}
  
  		return loadSpecs;
  	}
  
  	@Override
  	protected void load(ByteProvider provider, LoadSpec loadSpec, List<Option> options,
  						Program program, TaskMonitor monitor,
  						MessageLog log) throws IOException {
  		BinaryReader reader = new BinaryReader(provider, true);
  		TerseExecutableHeader teHeader = new TerseExecutableHeader(reader);
  		EFIImageSectionHeader[] sectionHeaders = teHeader.getSections();
  		FlatProgramAPI api = new FlatProgramAPI(program, monitor);
  		InputStream inputStream = provider.getInputStream(0);
  
  		try {
  			// Create a segment for the TE header and section headers.
  			api.createMemoryBlock(HEADERS, api.toAddr(
  					teHeader.getImageBase() + teHeader.getHeaderOffset()), inputStream,
  					TerseExecutableConstants.TE_HEADER_SIZE + teHeader.getNumSections() *
  					TerseExecutableConstants.SECTION_HEADER_SIZE, false);
  
  			// Mark the TE header and section headers as data.
  			api.createData(api.toAddr(teHeader.getImageBase() + teHeader.getHeaderOffset()),
  					teHeader.toDataType());
  			long sectionHeaderOffset = teHeader.getImageBase() + teHeader.getHeaderOffset() +
  					TerseExecutableConstants.TE_HEADER_SIZE;
  			for (EFIImageSectionHeader sectionHeader : sectionHeaders) {
  				api.createData(api.toAddr(sectionHeaderOffset), sectionHeader.toDataType());
  				sectionHeaderOffset += TerseExecutableConstants.SECTION_HEADER_SIZE;
  			}
  
  			// Create a segment for each section.
  			for (EFIImageSectionHeader sectionHeader : sectionHeaders) {
  				// Skip empty sections.
  				if (sectionHeader.getVirtualSize() == 0) {
  					Msg.debug(this, "Skipping empty section: " + sectionHeader.getName());
  					continue;
  				}
  
  				inputStream = provider.getInputStream(sectionHeader.getVirtualAddress() -
  						teHeader.getHeaderOffset());
  				long startAddress = teHeader.getImageBase() + sectionHeader.getVirtualAddress();
  				try {
  					MemoryBlock section = api.createMemoryBlock(sectionHeader.getName(),
  							api.toAddr(startAddress), inputStream, sectionHeader.getVirtualSize(),
  							false);
  
  					// Set the appropriate permissions for this segment.
  					section.setRead(sectionHeader.isReadable());
  					section.setWrite(sectionHeader.isWritable());
  					section.setExecute(sectionHeader.isExecutable());
  
  					Msg.debug(this, String.format("Added %s section: 0x%X-0x%X",
  							sectionHeader.getName(), startAddress, startAddress +
  							sectionHeader.getVirtualSize()));
  				} catch (AddressOverflowException e) {
  					Msg.warn(this,
  							"Skipping overflowing " + sectionHeader.getName() + " section:", e);
  				} catch (AddressOutOfBoundsException e) {
  					// This is thrown immediately after AddressOverflowException. Ignore this since
  					// we should have already shown the warning message.
  				}
  			}
  
  			// Set the UEFI property if this a UEFI binary.
  			if (teHeader.getSubsystem() >= PeSubsystem.IMAGE_SUBSYSTEM_EFI_APPLICATION.ordinal() &&
  				teHeader.getSubsystem() <= PeSubsystem.IMAGE_SUBSYSTEM_EFI_ROM.ordinal()) {
  				program.getOptions(Program.PROGRAM_INFO).setBoolean("UEFI", true);
  			}
  
  			// Define the entry point function.
  			Address entryPoint = api.toAddr(teHeader.getImageBase() + teHeader.getEntryPointAddress());
  			api.addEntryPoint(entryPoint);
  			api.createFunction(entryPoint, "_ModuleEntryPoint");
  		} catch (Exception e) {
  			Msg.showError(this, null, getName() + " Loader", e.getMessage(), e);
  		}
  	}
  
  	@Override
  	public String getName() {
  		return TE_NAME;
  	}
  }
  