NomanNasirMinhas-Ringer/pefile/pefile.go

// Package pefile wraps debug/pe with helpers tailored to kernel driver (.sys)
// analysis: raw-byte access, RVA<->file-offset mapping, and machine/bitness
// detection.
package pefile

import (
	"bytes"
	"debug/pe"
	"fmt"
	"math"
	"os"
)

// Machine is the CPU architecture of a PE image.
type Machine uint16

const (
	MachineUnknown Machine = 0
	MachineI386    Machine = 0x14c  // IMAGE_FILE_MACHINE_I386
	MachineAMD64   Machine = 0x8664 // IMAGE_FILE_MACHINE_AMD64
	MachineARMNT   Machine = 0x1c4  // IMAGE_FILE_MACHINE_ARMNT
	MachineARM64   Machine = 0xaa64 // IMAGE_FILE_MACHINE_ARM64
)

func (m Machine) String() string {
	switch m {
	case MachineI386:
		return "x86 (32-bit)"
	case MachineAMD64:
		return "x64 (64-bit)"
	case MachineARMNT:
		return "ARM (32-bit)"
	case MachineARM64:
		return "ARM64"
	default:
		return fmt.Sprintf("unknown (0x%04x)", uint16(m))
	}
}

// IsX86 reports whether the image uses an x86/x64 instruction set that the
// x86asm disassembler can decode.
func (m Machine) IsX86() bool {
	return m == MachineI386 || m == MachineAMD64
}

// Section is a PE section with its raw (on-disk) bytes.
type Section struct {
	Name            string
	VirtualAddress  uint32
	VirtualSize     uint32
	Offset          uint32 // file offset of raw data
	Size            uint32 // raw data size on disk
	Characteristics uint32
	Data            []byte // raw bytes (Size long)
}

// File is a parsed PE image plus the original file bytes.
type File struct {
	Path     string
	Raw      []byte
	PE       *pe.File
	Machine  Machine
	Is64     bool
	Sections []*Section
}

// Open reads and parses a PE file. It returns the raw bytes alongside the
// parsed image so callers can scan the whole file (e.g. for strings) without a
// second read.
func Open(path string) (*File, error) {
	raw, err := os.ReadFile(path)
	if err != nil {
		return nil, fmt.Errorf("read %s: %w", path, err)
	}
	if len(raw) < 0x40 {
		return nil, fmt.Errorf("%s: file too small to be a PE image (%d bytes)", path, len(raw))
	}

	pf, err := pe.NewFile(bytes.NewReader(raw))
	if err != nil {
		return nil, fmt.Errorf("parse PE %s: %w", path, err)
	}

	f := &File{
		Path:    path,
		Raw:     raw,
		PE:      pf,
		Machine: Machine(pf.FileHeader.Machine),
	}
	f.Is64 = f.Machine == MachineAMD64 || f.Machine == MachineARM64

	for _, s := range pf.Sections {
		data, derr := s.Data()
		if derr != nil {
			// A section whose raw data cannot be read (e.g. truncated file)
			// is still recorded with empty data so callers can report it.
			data = nil
		}
		f.Sections = append(f.Sections, &Section{
			Name:            s.Name,
			VirtualAddress:  s.VirtualAddress,
			VirtualSize:     s.VirtualSize,
			Offset:          s.Offset,
			Size:            s.Size,
			Characteristics: s.Characteristics,
			Data:            data,
		})
	}
	return f, nil
}

// Section returns the section with the given name (case-insensitive), or nil.
func (f *File) Section(name string) *Section {
	for _, s := range f.Sections {
		if len(s.Name) >= len(name) && bytes.EqualFold([]byte(s.Name[:len(name)]), []byte(name)) {
			return s
		}
	}
	return nil
}

// RVAToOffset maps a relative virtual address to a file offset. It returns an
// error if the RVA does not fall inside any section's raw data.
func (f *File) RVAToOffset(rva uint32) (uint32, error) {
	for _, s := range f.Sections {
		if s.Size == 0 {
			continue
		}
		if rva >= s.VirtualAddress && rva < s.VirtualAddress+s.Size {
			return s.Offset + (rva - s.VirtualAddress), nil
		}
	}
	return 0, fmt.Errorf("RVA 0x%x not mapped to any section", rva)
}

// OffsetToRVA maps a file offset back to an RVA, or returns an error if the
// offset is not inside a section's raw data.
func (f *File) OffsetToRVA(off uint32) (uint32, error) {
	for _, s := range f.Sections {
		if s.Size == 0 {
			continue
		}
		if off >= s.Offset && off < s.Offset+s.Size {
			return s.VirtualAddress + (off - s.Offset), nil
		}
	}
	return 0, fmt.Errorf("file offset 0x%x not mapped to any section", off)
}

// ReadRVA returns n bytes at the given RVA, or nil if the RVA is not mapped to
// raw data or the read would run past the end of the file.
func (f *File) ReadRVA(rva uint32, n int) []byte {
	off, err := f.RVAToOffset(rva)
	if err != nil {
		return nil
	}
	if int(off)+n > len(f.Raw) {
		return nil
	}
	return f.Raw[off : off+uint32(n)]
}

// Entropy estimates the Shannon entropy of a byte slice in bits per byte. It is
// used to flag packed/encrypted drivers whose static analysis will be unreliable.
func Entropy(b []byte) float64 {
	if len(b) == 0 {
		return 0
	}
	var freq [256]int
	for _, c := range b {
		freq[c]++
	}
	var h float64
	n := float64(len(b))
	for _, c := range freq {
		if c == 0 {
			continue
		}
		p := float64(c) / n
		h -= p * math.Log2(p)
	}
	return h
}