VitaLoader Ghidra not finished

//@author Sergi Granell
//@category _NEW_
//@keybinding
//@menupath
//@toolbar

import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.nio.IntBuffer;
import java.util.HashMap;

import generic.continues.RethrowContinuesFactory;
import ghidra.app.script.GhidraScript;
import ghidra.app.util.bin.BinaryReader;
import ghidra.app.util.bin.ByteArrayProvider;
import ghidra.app.util.bin.ByteProvider;
import ghidra.app.util.bin.StructConverter;
import ghidra.app.util.bin.StructConverterUtil;
import ghidra.app.util.bin.format.elf.ElfException;
import ghidra.app.util.bin.format.elf.ElfHeader;
import ghidra.app.util.opinion.ElfLoader;
import ghidra.program.model.address.Address;
import ghidra.program.model.data.DataType;
import ghidra.program.model.listing.Program;
import ghidra.program.model.mem.MemoryAccessException;
import ghidra.program.model.mem.MemoryBlock;
import ghidra.util.Msg;
import ghidra.util.exception.DuplicateNameException;
import org.yaml.snakeyaml.Yaml;

public class VitaLoader extends GhidraScript {
    private final short ET_SCE_RELEXEC = (short)0xFE04;

        /*sce_module_info_t = new StructureDataType(new CategoryPath("/Vita"), "sce_module_info_t", 0);
        sce_module_info_t.add(StructConverter.WORD, "attributes", null);
        sce_module_info_t.add(StructConverter.WORD, "version", null);
        sce_module_info_t.add(StructConverter.STRING, 27, "name", null);
        sce_module_info_t.add(StructConverter.BYTE, "type", null);
        sce_module_info_t.add(Pointer32DataType.dataType, "gp_value", null);
        sce_module_info_t.add(StructConverter.DWORD, "export_top", null);
        sce_module_info_t.add(StructConverter.DWORD, "export_end", null);
        sce_module_info_t.add(StructConverter.DWORD, "import_top", null);
        sce_module_info_t.add(StructConverter.DWORD, "import_end", null);
        sce_module_info_t.add(StructConverter.DWORD, "module_nid", null);
        sce_module_info_t.add(StructConverter.DWORD, "field_38", null);
        sce_module_info_t.add(StructConverter.DWORD, "field_3C", null);
        sce_module_info_t.add(StructConverter.DWORD, "field_40", null);
        sce_module_info_t.add(StructConverter.DWORD, "module_start", null);
        sce_module_info_t.add(StructConverter.DWORD, "module_stop", null);
        sce_module_info_t.add(StructConverter.DWORD, "exidx_top", null);
        sce_module_info_t.add(StructConverter.DWORD, "exidx_end", null);
        sce_module_info_t.add(StructConverter.DWORD, "extab_top", null);
        sce_module_info_t.add(StructConverter.DWORD, "extab_end", null);*/

    public class SceModuleInfo implements StructConverter {
        public short attributes;
        public short version;
        public String name;
        public byte type;
        public int gp_value;
        public int export_top;
        public int export_end;
        public int import_top;
        public int import_end;
        public int module_nid;
        public int field_38;
        public int field_3C;
        public int field_40;
        public int module_start;
        public int module_stop;
        public int exidx_top;
        public int exidx_end;
        public int extab_top;
        public int extab_end;
        public static final int SIZE = 0x5c;

        SceModuleInfo(BinaryReader reader) throws IOException {
            attributes = reader.readNextShort();
            version = reader.readNextShort();
            name = reader.readNextAsciiString(27);
            type = reader.readNextByte();
            gp_value = reader.readNextInt();
            export_top = reader.readNextInt();
            export_end = reader.readNextInt();
            import_top = reader.readNextInt();
            import_end = reader.readNextInt();
            module_nid = reader.readNextInt();
            field_38 = reader.readNextInt();
            field_3C = reader.readNextInt();
            field_40 = reader.readNextInt();
            module_start = reader.readNextInt();
            module_stop = reader.readNextInt();
            exidx_top = reader.readNextInt();
            exidx_end = reader.readNextInt();
            extab_top = reader.readNextInt();
            extab_end = reader.readNextInt();
        }

        public DataType toDataType() throws DuplicateNameException, IOException {
            return StructConverterUtil.toDataType(this);
        }
    }

    public class SceModuleImports implements StructConverter {
        public short size;
        public short version;
        public short flags;
        public short num_syms_funcs;
        public short num_syms_vars;
        public short num_syms_unk;
        public int reserved1;
        public int library_nid;
        public int library_name;
        public int reserved2;
        public int func_nid_table;
        public int func_entry_table;
        public int var_nid_table;
        public int var_entry_table;
        public int unk_nid_table;
        public int unk_entry_table;
        public static final int SIZE = 0x34;

        SceModuleImports(BinaryReader reader) throws IOException {
            size = reader.readNextShort();
            version = reader.readNextShort();
            flags = reader.readNextShort();
            num_syms_funcs = reader.readNextShort();
            num_syms_vars = reader.readNextShort();
            num_syms_unk = reader.readNextShort();
            reserved1 = reader.readNextInt();
            library_nid = reader.readNextInt();
            library_name = reader.readNextInt();
            reserved2 = reader.readNextInt();
            func_nid_table = reader.readNextInt();
            func_entry_table = reader.readNextInt();
            var_nid_table = reader.readNextInt();
            var_entry_table = reader.readNextInt();
            unk_nid_table = reader.readNextInt();
            unk_entry_table = reader.readNextInt();
        }

        public DataType toDataType() throws DuplicateNameException, IOException {
            return StructConverterUtil.toDataType(this);
        }
    }

    public class NidDb {
        public class NidDbLibrary {
            public int nid;
            public HashMap<Integer, String> functions;

            public NidDbLibrary(int nid) {
                this.nid = nid;
                functions = new HashMap<Integer, String>();
            }

            public boolean functionExists(int nid) {
                return functions.containsKey(nid);
            }

            public void insertFunction(int nid, String name) {
                functions.put(nid, name);
            }

            public String getFunctionName(int nid) {
                return functions.get(nid);
            }
        }

        public class NidDbModule {
            private int nid;
            private HashMap<Integer, NidDbLibrary> libraries;

            public NidDbModule(int nid) {
                this.nid = nid;
                libraries = new HashMap<Integer, NidDbLibrary>();
            }

            public boolean libraryExists(int nid) {
                return libraries.containsKey(nid);
            }

            public void insertLibrary(int nid, NidDbLibrary library) {
                libraries.put(nid, library);
            }

            public NidDbLibrary getLibrary(int nid) {
                return libraries.get(nid);
            }

            public String getFunctionName(int libraryNid, int functionNid) {
                return getLibrary(libraryNid).getFunctionName(functionNid);
            }
        }

        private HashMap<Integer, NidDbModule> modules;

        public NidDb() {
            modules = new HashMap<Integer, NidDbModule>();
        }

        public boolean moduleExists(int nid) {
            return modules.containsKey(nid);
        }

        public void insertModule(int nid, NidDbModule module) {
            modules.put(nid, module);
        }

        public NidDbModule getModule(int nid) {
            return modules.get(nid);
        }

        public String getModuleLibraryFunctionName(int moduleNid, int libraryNid, int functionNid) {
            return getModule(moduleNid).getFunctionName(libraryNid, functionNid);
        }

    }


    private ElfHeader getElfHeader(Program program) throws MemoryAccessException, ElfException {
        MemoryBlock elfHeaderMemoryBlock = currentProgram.getMemory().getBlock("_elfHeader");

        int elfHeaderSize = (int)elfHeaderMemoryBlock.getSize();
        Address elfHeaderStart = elfHeaderMemoryBlock.getStart();
        byte[] elfHeaderBytes = new byte[elfHeaderSize];

        elfHeaderMemoryBlock.getBytes(elfHeaderStart, elfHeaderBytes);

        ByteArrayProvider elfHeaderProvider = new ByteArrayProvider(elfHeaderBytes);
        return ElfHeader.createElfHeader(RethrowContinuesFactory.INSTANCE, elfHeaderProvider);
    }

    private MemoryBlock getExecutableMemoryBlock(Program program) {
        for (MemoryBlock block : program.getMemory().getBlocks()) {
            if ((block.getPermissions() & MemoryBlock.EXECUTE) != 0)
                return block;
        }
        return null;
    }

    private void processImports(MemoryBlock block, SceModuleImports imports) throws MemoryAccessException {
        println("Imports NID: " + String.format("0x%08X", imports.library_nid));
        println("Imports num funcs: " + String.format("0x%08X", imports.num_syms_funcs));
        println("Imports nid table: " + String.format("0x%08X", imports.func_nid_table));
        println("Imports size: " + String.format("0x%08X", imports.size));

        Address funcNidTableAddress = block.getStart().getNewAddress(imports.func_nid_table);
        Address funcEntryTableAddress = block.getStart().getNewAddress(imports.func_entry_table);

        byte[] funcNidTableBytes = new byte[4 * imports.num_syms_funcs];
        byte[] funcEntryTableBytes = new byte[4 * imports.num_syms_funcs];

        block.getBytes(funcNidTableAddress, funcNidTableBytes, 0, funcNidTableBytes.length);
        block.getBytes(funcEntryTableAddress, funcEntryTableBytes, 0, funcEntryTableBytes.length);
        IntBuffer funcNidTableIntBuffer = ByteBuffer.wrap(funcNidTableBytes).order(ByteOrder.LITTLE_ENDIAN).asIntBuffer();
        IntBuffer funcEntryTableIntBuffer = ByteBuffer.wrap(funcEntryTableBytes).order(ByteOrder.LITTLE_ENDIAN).asIntBuffer();

        for (int i = 0; i < imports.num_syms_funcs; i++) {
            println(" NID[" + i + "]: " + String.format("0x%08X", funcNidTableIntBuffer.get(i)));
        }
    }

    private void processModuleInfo(MemoryBlock block, SceModuleInfo moduleInfo) throws MemoryAccessException, IOException {
        println("Module name: " + moduleInfo.name);

        Address importsTop = block.getStart().add(moduleInfo.import_top);
        Address importsEnd = block.getStart().add(moduleInfo.import_end);
        while (!importsTop.equals(importsEnd)) {
            byte[] importsBytes = new byte[SceModuleImports.SIZE];
            block.getBytes(importsTop, importsBytes, 0, importsBytes.length);

            ByteProvider importsByteProvider = new ByteArrayProvider(importsBytes);
            BinaryReader importsBinaryReader = new BinaryReader(importsByteProvider, true);
            SceModuleImports moduleImports = new SceModuleImports(importsBinaryReader);
            if (moduleImports.size != SceModuleImports.SIZE) {
                println("Imports size mismatch");
                break;
            }

            processImports(block, moduleImports);

            importsTop = importsTop.add(moduleImports.size);
        }
    }

    @Override
    public void run() throws Exception {
        println("VitaLoader by xerpi");

        if (!ElfLoader.ELF_NAME.equals(currentProgram.getExecutableFormat())) {
            Msg.showError(this, null, "VitaLoader",
                    "Current program is not an ELF program!  (" + currentProgram.getExecutableFormat() + ")");
            return;
        }

        ElfHeader elfHeader = getElfHeader(currentProgram);
        if (elfHeader.e_type() != ET_SCE_RELEXEC) {
            Msg.showError(this, null, "VitaLoader",
                    "Current program is not an PSVita ELF program!  (" + String.format("0x%04X", elfHeader.e_type()) + ")");
            return;
        }


        InputStream input = new FileInputStream(new File("/home/xerpi/Desktop/Vita/vita-headers/db.yml"));
        Yaml yaml = new Yaml();
        Object data = yaml.load(input);


        MemoryBlock textMemoryBlock = getExecutableMemoryBlock(currentProgram);

        println("ELF entry:" + String.format("0x%08X", elfHeader.e_entry()));
        println(".text size: " + String.format("0x%08X", textMemoryBlock.getSize()));

        Address moduleInfoAddress = textMemoryBlock.getStart().add(elfHeader.e_entry());
        byte[] moduleInfoBytes = new byte[SceModuleInfo.SIZE];
        textMemoryBlock.getBytes(moduleInfoAddress, moduleInfoBytes, 0, moduleInfoBytes.length);

        ByteProvider provider = new ByteArrayProvider(moduleInfoBytes);
        BinaryReader reader = new BinaryReader(provider, true);
        SceModuleInfo moduleInfo = new SceModuleInfo(reader);

        processModuleInfo(textMemoryBlock, moduleInfo);
    }
}