whitemagic/WhiteMagic/Discovery/PeHeaderParser.cs

using System.ComponentModel;
using System.Runtime.InteropServices;
using WhiteMagic.Native;

namespace WhiteMagic.Discovery;

/// <summary>
/// Represents a section in a PE file.
/// </summary>
public readonly record struct PeSection
{
    /// <summary>
    /// The 8-byte null-terminated section name (e.g., ".text", ".data").
    /// </summary>
    public string Name { get; init; }

    /// <summary>
    /// The virtual address of the section when loaded into memory (RVA).
    /// </summary>
    public IntPtr VirtualAddress { get; init; }

    /// <summary>
    /// The size of the section in memory.
    /// </summary>
    public int VirtualSize { get; init; }
}

/// <summary>
/// Parses PE headers to expose section information and entry points.
/// </summary>
public sealed class PeHeaderParser
{
    private readonly MemoryBase _memory;
    private readonly IntPtr _baseAddress;

    /// <summary>
    /// Creates a new PE header parser for the module at the specified base address.
    /// </summary>
    /// <param name="memory">The memory accessor.</param>
    /// <param name="baseAddress">The base address of the module.</param>
    public PeHeaderParser(MemoryBase memory, IntPtr baseAddress)
    {
        ArgumentNullException.ThrowIfNull(memory);
        if (baseAddress == IntPtr.Zero)
            throw new ArgumentException("Base address cannot be zero.", nameof(baseAddress));

        _memory = memory;
        _baseAddress = baseAddress;
    }

    /// <summary>
    /// Gets the entry point RVA (Relative Virtual Address) of the PE file.
    /// </summary>
    /// <returns>The entry point RVA, or <see cref="IntPtr.Zero"/> if unavailable.</returns>
    /// <exception cref="Win32Exception">Reading memory fails.</exception>
    /// <exception cref="InvalidDataException">The PE headers are invalid.</exception>
    public IntPtr EntryPoint
    {
        get
        {
            // Read and parse PE headers
            var (optionalHeader, _) = ParseOptionalHeader();

            if (optionalHeader is null)
                return IntPtr.Zero;

            // Entry point is at different offsets for PE32 vs PE32+
            bool isPe32Plus = IsPe32Plus();

            if (isPe32Plus)
            {
                // PE32+: AddressOfEntryPoint is at offset 16 in OPTIONAL_HEADER (64-bit)
                return (IntPtr)BitConverter.ToUInt32(
                    optionalHeader.AsSpan(16, 4));
            }
            else
            {
                // PE32: AddressOfEntryPoint is at offset 16 in OPTIONAL_HEADER (32-bit)
                return (IntPtr)BitConverter.ToUInt32(
                    optionalHeader.AsSpan(16, 4));
            }
        }
    }

    /// <summary>
    /// Enumerates all sections in the PE file.
    /// </summary>
    /// <returns>An enumerable of PE sections.</returns>
    /// <exception cref="Win32Exception">Reading memory fails.</exception>
    /// <exception cref="InvalidDataException">The PE headers are invalid.</exception>
    public IEnumerable<PeSection> Sections
    {
        get
        {
            var (optionalHeader, sectionHeaders) = ParseOptionalHeaderAndSectionHeaders();

            if (sectionHeaders is null || sectionHeaders.Length == 0)
                yield break;

            foreach (var sectionHeader in sectionHeaders)
            {
                // Parse section name (8-byte, null-terminated)
                string name = ParseSectionName(sectionHeader);

                // VirtualAddress and VirtualSize
                uint virtualAddress = BitConverter.ToUInt32(sectionHeader, 12);
                uint virtualSize = BitConverter.ToUInt32(sectionHeader, 8);

                yield return new PeSection
                {
                    Name = name,
                    VirtualAddress = (IntPtr)virtualAddress,
                    VirtualSize = (int)virtualSize
                };
            }
        }
    }

    /// <summary>
    /// Parses the DOS header, PE signature, and optional header.
    /// </summary>
    private (byte[]? OptionalHeader, byte[][]? SectionHeaders) ParseOptionalHeaderAndSectionHeaders()
    {
        // Read DOS header (first 64 bytes)
        byte[] dosHeader = _memory.ReadBytes(_baseAddress, 64);
        if (dosHeader.Length < 64)
            throw new InvalidDataException("Failed to read DOS header.");

        // Verify DOS signature "MZ"
        if (dosHeader[0] != 0x4D || dosHeader[1] != 0x5A)
            throw new InvalidDataException("Invalid DOS signature (not a PE file).");

        // PE header offset is at 0x3C in DOS header
        int peOffset = BitConverter.ToInt32(dosHeader, 0x3C);
        if (peOffset < 0 || peOffset > 0x1000) // Sanity check
            throw new InvalidDataException($"Invalid PE offset: {peOffset}");

        // Read PE signature (4 bytes: "PE\0\0")
        IntPtr peSigAddr = _baseAddress + peOffset;
        byte[] peSignature = _memory.ReadBytes(peSigAddr, 4);
        if (peSignature.Length < 4)
            throw new InvalidDataException("Failed to read PE signature.");

        if (peSignature[0] != 0x50 || peSignature[1] != 0x45 ||
            peSignature[2] != 0x00 || peSignature[3] != 0x00)
            throw new InvalidDataException("Invalid PE signature.");

        // COFF header follows PE signature (20 bytes)
        IntPtr coffAddr = peSigAddr + 4;
        byte[] coffHeader = _memory.ReadBytes(coffAddr, 20);
        if (coffHeader.Length < 20)
            throw new InvalidDataException("Failed to read COFF header.");

        // SizeOfOptionalHeader is at offset 16 in COFF header
        ushort sizeOfOptionalHeader = BitConverter.ToUInt16(coffHeader, 16);
        // NumberOfSections is at offset 2 in COFF header
        ushort numberOfSections = BitConverter.ToUInt16(coffHeader, 2);

        if (numberOfSections == 0 || numberOfSections > 96)
            return (null, null); // No sections or unreasonable number
        // Optional header follows COFF header
        IntPtr optAddr = coffAddr + 20;
        byte[] optionalHeader = _memory.ReadBytes(optAddr, sizeOfOptionalHeader);
        if (optionalHeader.Length < sizeOfOptionalHeader)
            throw new InvalidDataException("Failed to read optional header.");

        // Section headers follow optional header
        IntPtr sectionAddr = optAddr + sizeOfOptionalHeader;
        int sectionHeaderSize = 40; // IMAGE_SECTION_HEADER is 40 bytes

        byte[][] sectionHeaders = new byte[numberOfSections][];
        for (int i = 0; i < numberOfSections; i++)
        {
            byte[] section = _memory.ReadBytes(sectionAddr + (i * sectionHeaderSize), sectionHeaderSize);
            if (section.Length < sectionHeaderSize)
                throw new InvalidDataException($"Failed to read section header {i}.");
            sectionHeaders[i] = section;
        }

        return (optionalHeader, sectionHeaders);
    }

    /// <summary>
    /// Parses just the optional header (for entry point).
    /// </summary>
    private (byte[]? OptionalHeader, byte[][]? SectionHeaders) ParseOptionalHeader()
    {
        return ParseOptionalHeaderAndSectionHeaders();
    }

    /// <summary>
    /// Determines whether the PE file is PE32+ (64-bit) or PE32 (32-bit).
    /// </summary>
    private bool IsPe32Plus()
    {
        var (optionalHeader, _) = ParseOptionalHeaderAndSectionHeaders();

        if (optionalHeader is null || optionalHeader.Length < 2)
            throw new InvalidDataException("Optional header too short.");

        // Magic is at offset 0 in optional header
        // 0x10b = PE32 (32-bit), 0x20b = PE32+ (64-bit)
        ushort magic = BitConverter.ToUInt16(optionalHeader, 0);
        return magic == 0x20b;
    }

    /// <summary>
    /// Parses a null-terminated 8-byte section name.
    /// </summary>
    private static string ParseSectionName(byte[] sectionHeader)
    {
        // Name is first 8 bytes
        var nameBytes = new Span<byte>(sectionHeader, 0, 8);

        // Find null terminator
        int len = 0;
        for (; len < 8; len++)
        {
            if (nameBytes[len] == 0)
                break;
        }

        return System.Text.Encoding.ASCII.GetString(nameBytes[..len]);
    }
}