diff --git a/src/macho/MachO.go b/src/macho/MachO.go index c463694..17d9564 100644 --- a/src/macho/MachO.go +++ b/src/macho/MachO.go @@ -116,7 +116,6 @@ func Write(writer io.Writer, code []byte, data []byte) { binary.Write(writer, binary.LittleEndian, &m.CodeHeader) binary.Write(writer, binary.LittleEndian, &m.DataHeader) binary.Write(writer, binary.LittleEndian, &m.UnixThread) - writer.Write(bytes.Repeat([]byte{0x00}, codePadding)) writer.Write(code) writer.Write(bytes.Repeat([]byte{0x00}, dataPadding)) diff --git a/src/pe/EXE.go b/src/pe/EXE.go index 238b3a7..bfbaef9 100644 --- a/src/pe/EXE.go +++ b/src/pe/EXE.go @@ -25,73 +25,22 @@ type EXE struct { // Write writes the EXE file to the given writer. func Write(writer io.Writer, code []byte, data []byte, dlls dll.List) { - codeStart, codePadding := fs.Align(HeaderEnd, config.Align) - dataStart, dataPadding := fs.Align(codeStart+len(code), config.Align) - importsStart, importsPadding := fs.Align(dataStart+len(data), config.Align) - - subSystem := IMAGE_SUBSYSTEM_WINDOWS_CUI + var ( + codeStart, codePadding = fs.Align(HeaderEnd, config.Align) + dataStart, dataPadding = fs.Align(codeStart+len(code), config.Align) + importsStart, importsPadding = fs.Align(dataStart+len(data), config.Align) + subSystem = IMAGE_SUBSYSTEM_WINDOWS_CUI + imports, dllData, dllImports, dllDataStart = importLibraries(dlls, importsStart) + importDirectoryStart = dllDataStart + len(dllData) + importDirectorySize = DLLImportSize * len(dllImports) + importSectionSize = len(imports)*8 + len(dllData) + importDirectorySize + imageSize, _ = fs.Align(importsStart+importSectionSize, config.Align) + ) if dlls.Contains("user32") { subSystem = IMAGE_SUBSYSTEM_WINDOWS_GUI } - imports := make([]uint64, 0) - dllData := make([]byte, 0) - dllImports := []DLLImport{} - - for _, library := range dlls { - functionsStart := len(imports) * 8 - dllNamePos := len(dllData) - dllData = append(dllData, library.Name...) - dllData = append(dllData, ".dll"...) - dllData = append(dllData, 0x00) - - dllImports = append(dllImports, DLLImport{ - RvaFunctionNameList: uint32(importsStart + functionsStart), - TimeDateStamp: 0, - ForwarderChain: 0, - RvaModuleName: uint32(dllNamePos), - RvaFunctionAddressList: uint32(importsStart + functionsStart), - }) - - for _, fn := range library.Functions { - if len(dllData)&1 != 0 { - dllData = append(dllData, 0x00) // align the next entry on an even boundary - } - - offset := len(dllData) - dllData = append(dllData, 0x00, 0x00) - dllData = append(dllData, fn...) - dllData = append(dllData, 0x00) - - imports = append(imports, uint64(offset)) - } - - imports = append(imports, 0) - } - - dllDataStart := importsStart + len(imports)*8 - - for i := range imports { - if imports[i] == 0 { - continue - } - - imports[i] += uint64(dllDataStart) - } - - for i := range dllImports { - dllImports[i].RvaModuleName += uint32(dllDataStart) - } - - dllImports = append(dllImports, DLLImport{}) // a zeroed structure marks the end of the list - importDirectoryStart := dllDataStart + len(dllData) - importDirectorySize := DLLImportSize * len(dllImports) - - importSectionSize := len(imports)*8 + len(dllData) + importDirectorySize - imageSize := importsStart + importSectionSize - imageSize, _ = fs.Align(imageSize, config.Align) - pe := &EXE{ DOSHeader: DOSHeader{ Magic: [4]byte{'M', 'Z', 0, 0}, @@ -188,7 +137,6 @@ func Write(writer io.Writer, code []byte, data []byte, dlls dll.List) { binary.Write(writer, binary.LittleEndian, &pe.NTHeader) binary.Write(writer, binary.LittleEndian, &pe.OptionalHeader64) binary.Write(writer, binary.LittleEndian, &pe.Sections) - writer.Write(bytes.Repeat([]byte{0x00}, codePadding)) writer.Write(code) writer.Write(bytes.Repeat([]byte{0x00}, dataPadding)) diff --git a/src/pe/importLibraries.go b/src/pe/importLibraries.go new file mode 100644 index 0000000..d7c8d32 --- /dev/null +++ b/src/pe/importLibraries.go @@ -0,0 +1,60 @@ +package pe + +import "git.akyoto.dev/cli/q/src/dll" + +// importLibraries generates the import address table which contains the addresses of functions imported from DLLs. +func importLibraries(dlls dll.List, importsStart int) ([]uint64, []byte, []DLLImport, int) { + imports := make([]uint64, 0) + dllData := make([]byte, 0) + dllImports := []DLLImport{} + + for _, library := range dlls { + functionsStart := len(imports) * 8 + dllNamePos := len(dllData) + dllData = append(dllData, library.Name...) + dllData = append(dllData, ".dll"...) + dllData = append(dllData, 0x00) + + dllImports = append(dllImports, DLLImport{ + RvaFunctionNameList: uint32(importsStart + functionsStart), + TimeDateStamp: 0, + ForwarderChain: 0, + RvaModuleName: uint32(dllNamePos), + RvaFunctionAddressList: uint32(importsStart + functionsStart), + }) + + for _, fn := range library.Functions { + if len(dllData)&1 != 0 { + dllData = append(dllData, 0x00) // align the next entry on an even boundary + } + + offset := len(dllData) + dllData = append(dllData, 0x00, 0x00) + dllData = append(dllData, fn...) + dllData = append(dllData, 0x00) + + imports = append(imports, uint64(offset)) + } + + imports = append(imports, 0) + } + + dllDataStart := importsStart + len(imports)*8 + + for i := range imports { + if imports[i] == 0 { + continue + } + + imports[i] += uint64(dllDataStart) + } + + for i := range dllImports { + dllImports[i].RvaModuleName += uint32(dllDataStart) + } + + // a zeroed structure marks the end of the list + dllImports = append(dllImports, DLLImport{}) + + return imports, dllData, dllImports, dllDataStart +}