From 80ca44b5c297988adb1dc39c88cfa0ee47409f5a Mon Sep 17 00:00:00 2001 From: CaptWake Date: Mon, 14 Sep 2026 18:25:55 +0200 Subject: [PATCH] add arm64 support --- pe/cert.go | 16 ++++++++-------- pe/exports.go | 2 +- pe/file.go | 8 +++++++- pe/imports.go | 4 ++-- pe/reloc.go | 4 ++-- pe/write.go | 12 +++++------- 6 files changed, 25 insertions(+), 21 deletions(-) diff --git a/pe/cert.go b/pe/cert.go index 40cb175..30ad884 100644 --- a/pe/cert.go +++ b/pe/cert.go @@ -17,15 +17,15 @@ func readCertTable(f *File, r io.ReadSeeker) ([]byte, error) { var certTableOffset, certTableSize uint32 - switch f.FileHeader.Machine { - case IMAGE_FILE_MACHINE_I386: - certTableOffset = f.OptionalHeader.(*OptionalHeader32).DataDirectory[CERTIFICATE_TABLE].VirtualAddress - certTableSize = f.OptionalHeader.(*OptionalHeader32).DataDirectory[CERTIFICATE_TABLE].Size - case IMAGE_FILE_MACHINE_AMD64: - certTableOffset = f.OptionalHeader.(*OptionalHeader64).DataDirectory[CERTIFICATE_TABLE].VirtualAddress - certTableSize = f.OptionalHeader.(*OptionalHeader64).DataDirectory[CERTIFICATE_TABLE].Size + switch oh := f.OptionalHeader.(type) { + case *OptionalHeader32: + certTableOffset = oh.DataDirectory[CERTIFICATE_TABLE].VirtualAddress + certTableSize = oh.DataDirectory[CERTIFICATE_TABLE].Size + case *OptionalHeader64: + certTableOffset = oh.DataDirectory[CERTIFICATE_TABLE].VirtualAddress + certTableSize = oh.DataDirectory[CERTIFICATE_TABLE].Size default: - return nil, errors.New("architecture not supported") + return nil, errors.New("unsupported optional header type") } // check if certificate table exists diff --git a/pe/exports.go b/pe/exports.go index eddac10..84af68d 100644 --- a/pe/exports.go +++ b/pe/exports.go @@ -31,7 +31,7 @@ type Export struct { // Exports - gets exports func (f *File) Exports() ([]Export, error) { - pe64 := f.Machine == IMAGE_FILE_MACHINE_AMD64 + pe64 := f.is64() // grab the number of data directory entries var ddLength uint32 diff --git a/pe/file.go b/pe/file.go index 06c46b6..dab2be6 100644 --- a/pe/file.go +++ b/pe/file.go @@ -137,7 +137,7 @@ func newFileInternal(r io.ReaderAt, memoryMode bool) (*File, error) { return nil, err } switch f.FileHeader.Machine { - case IMAGE_FILE_MACHINE_UNKNOWN, IMAGE_FILE_MACHINE_ARMNT, IMAGE_FILE_MACHINE_AMD64, IMAGE_FILE_MACHINE_I386: + case IMAGE_FILE_MACHINE_UNKNOWN, IMAGE_FILE_MACHINE_ARMNT, IMAGE_FILE_MACHINE_ARM64, IMAGE_FILE_MACHINE_AMD64, IMAGE_FILE_MACHINE_I386: default: return nil, fmt.Errorf("Unrecognised COFF file header machine value of 0x%x", f.FileHeader.Machine) } @@ -442,6 +442,12 @@ func (e *FormatError) Error() string { } // RVAToFileOffset Converts a Relative offset to the actual offset in the file. +// is64 reports whether the file has a PE32+ optional header. +func (f *File) is64() bool { + _, ok := f.OptionalHeader.(*OptionalHeader64) + return ok +} + func (f *File) RVAToFileOffset(rva uint32) uint32 { var offset uint32 for _, section := range f.Sections { diff --git a/pe/imports.go b/pe/imports.go index ebcf881..239f2c8 100644 --- a/pe/imports.go +++ b/pe/imports.go @@ -78,7 +78,7 @@ func (f *File) ImportDirectoryTable() ([]ImportDirectory, *Section, *[]byte, err // satisfied by other libraries at dynamic load time. // It does not return weak symbols. func (f *File) ImportedSymbols() ([]string, error) { - pe64 := f.Machine == IMAGE_FILE_MACHINE_AMD64 + pe64 := f.is64() ida, ds, sectionData, err := f.ImportDirectoryTable() if err != nil { @@ -142,7 +142,7 @@ func (f *File) ImportedLibraries() ([]string, error) { } func (f File) sectionFromDirectoryEntry(directory uint32) (*Section, DataDirectory) { - pe64 := f.Machine == IMAGE_FILE_MACHINE_AMD64 + pe64 := f.is64() // grab the number of data directory entries var ddLength uint32 diff --git a/pe/reloc.go b/pe/reloc.go index e29ee8e..ddb99ee 100644 --- a/pe/reloc.go +++ b/pe/reloc.go @@ -62,7 +62,7 @@ func (f *File) readBaseRelocationTable() (*[]RelocationTableEntry, error) { } var dd DataDirectory - if f.Machine == IMAGE_FILE_MACHINE_AMD64 { + if f.is64() { dd = f.OptionalHeader.(*OptionalHeader64).DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC] } else { dd = f.OptionalHeader.(*OptionalHeader32).DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC] @@ -111,7 +111,7 @@ func (f *File) readBaseRelocationTable() (*[]RelocationTableEntry, error) { // Relocate - performs base relocations on this image to the given offset func (f *File) Relocate(baseAddr uint64, image *[]byte) { var imageBase uint64 - pe64 := f.Machine == IMAGE_FILE_MACHINE_AMD64 + pe64 := f.is64() if pe64 { imageBase = f.OptionalHeader.(*OptionalHeader64).ImageBase } else { diff --git a/pe/write.go b/pe/write.go index a067f41..855a65f 100644 --- a/pe/write.go +++ b/pe/write.go @@ -48,25 +48,23 @@ func (peFile *File) Bytes() ([]byte, error) { oldCertTableOffset, oldCertTableSize uint32 ) - switch peFile.FileHeader.Machine { - case IMAGE_FILE_MACHINE_I386: + switch optionalHeader := peFile.OptionalHeader.(type) { + case *OptionalHeader32: is32bit = true - optionalHeader := peFile.OptionalHeader.(*OptionalHeader32) - binary.Write(peBuf, binary.LittleEndian, peFile.OptionalHeader.(*OptionalHeader32)) + binary.Write(peBuf, binary.LittleEndian, optionalHeader) bytesWritten += uint64(binary.Size(optionalHeader)) oldCertTableOffset = optionalHeader.DataDirectory[CERTIFICATE_TABLE].VirtualAddress oldCertTableSize = optionalHeader.DataDirectory[CERTIFICATE_TABLE].Size - case IMAGE_FILE_MACHINE_AMD64: + case *OptionalHeader64: is32bit = false - optionalHeader := peFile.OptionalHeader.(*OptionalHeader64) binary.Write(peBuf, binary.LittleEndian, optionalHeader) bytesWritten += uint64(binary.Size(optionalHeader)) oldCertTableOffset = optionalHeader.DataDirectory[CERTIFICATE_TABLE].VirtualAddress oldCertTableSize = optionalHeader.DataDirectory[CERTIFICATE_TABLE].Size default: - return nil, errors.New("architecture not supported") + return nil, errors.New("unsupported optional header type") } // write section headers