diff --git a/cstring.go b/cstring.go new file mode 100644 index 0000000..b967cf1 --- /dev/null +++ b/cstring.go @@ -0,0 +1,13 @@ +package pego + +import "bytes" + +// cstring converts ASCII byte sequence b to string. +// It stops once it finds 0 or reaches end of b. +func cstring(b []byte) string { + i := bytes.IndexByte(b, 0) + if i == -1 { + i = len(b) + } + return string(b[:i]) +} diff --git a/file.go b/file.go index 70d8d53..4f3b56e 100644 --- a/file.go +++ b/file.go @@ -15,6 +15,7 @@ type PE struct { COFFHeader *Header[pe.FileHeader] OptionalHeader32 *Header[pe.OptionalHeader32] OptionalHeader64 *Header[pe.OptionalHeader64] + Sections []*Section } // NewPE parses a Portable Executable or plain COFF file from reader. @@ -101,6 +102,17 @@ func NewPE(reader io.ReaderAt) (*PE, error) { } } + // Sections + numSections := int(p.COFFHeader.Data.NumberOfSections) + p.Sections = make([]*Section, numSections) + for i := range numSections { + section, err := NewSection(reader, &offset) + if err != nil { + return nil, err + } + p.Sections[i] = section + } + // TODO: add protections to defend against malicious files (e.g. oversized segments...) return &p, nil diff --git a/file_test.go b/file_test.go index 32d7706..8dc1075 100644 --- a/file_test.go +++ b/file_test.go @@ -37,6 +37,26 @@ func TestFilePianoExe(t *testing.T) { assert.Equal(t, p.OptionalHeader32.Data.SizeOfImage, uint32(0x630)) assert.Equal(t, p.OptionalHeader32.Data.SizeOfHeaders, uint32(0x230)) assert.Assert(t, p.OptionalHeader64 == nil) + + // Sections + assert.Equal(t, len(p.Sections), 3) + s0 := p.Sections[0] + assert.Equal(t, s0.Name(), ".") + assert.Equal(t, s0.Header.Data.VirtualSize, uint32(0x33c)) + assert.Equal(t, s0.Header.Data.VirtualAddress, uint32(0x230)) + assert.Equal(t, s0.Segment.Size(), int64(0x340)) + + s1 := p.Sections[1] + assert.Equal(t, s1.Name(), ".data") + assert.Equal(t, s1.Header.Data.VirtualSize, uint32(0x6a)) + assert.Equal(t, s1.Header.Data.VirtualAddress, uint32(0x570)) + assert.Equal(t, s1.Segment.Size(), int64(0x70)) + + s2 := p.Sections[2] + assert.Equal(t, s2.Name(), ".reloc") + assert.Equal(t, s2.Header.Data.VirtualSize, uint32(0x4c)) + assert.Equal(t, s2.Header.Data.VirtualAddress, uint32(0x5e0)) + assert.Equal(t, s2.Segment.Size(), int64(0x50)) } // TestFileReadTruncated verifies that NewPE returns an error when the input is truncated at various points. diff --git a/section.go b/section.go new file mode 100644 index 0000000..fa9946c --- /dev/null +++ b/section.go @@ -0,0 +1,34 @@ +package pego + +import ( + "debug/pe" + "io" +) + +type Section struct { + Header *Header[pe.SectionHeader32] + Segment *Segment +} + +func NewSection(reader io.ReaderAt, offset *int64) (*Section, error) { + // read the header + header, err := NewHeader[pe.SectionHeader32](reader, offset) + if err != nil { + return nil, err + } + + // create the segment + segmentOffset := int64(header.Data.PointerToRawData) + segmentSize := int64(header.Data.SizeOfRawData) + segment := NewSegment(reader, &segmentOffset, segmentSize) + + return &Section{ + Header: header, + Segment: segment, + }, nil +} + +func (s *Section) Name() string { + // TODO: it seems there's another case to support + return cstring(s.Header.Data.Name[:]) +}