// -----------------------------------------------------------------------------
// PeHardeningExtractor.cs
// Sprint: SPRINT_3500_0004_0001_smart_diff_binary_output
// Task: SDIFF-BIN-010 - Implement PeHardeningExtractor
// Task: SDIFF-BIN-011 - Implement PE DllCharacteristics parsing
// Task: SDIFF-BIN-012 - Implement PE Authenticode detection
// Description: Extracts security hardening flags from Windows PE binaries
// -----------------------------------------------------------------------------

using System.Buffers.Binary;
using System.Collections.Immutable;

namespace StellaOps.Scanner.Analyzers.Native.Hardening;

/// <summary>
/// Extracts hardening flags from Windows PE (Portable Executable) binaries.
/// Detects ASLR, DEP, CFG, Authenticode, Safe SEH, and other security features.
/// Per Sprint 3500.4 - Smart-Diff Binary Analysis.
/// </summary>
public sealed class PeHardeningExtractor : IHardeningExtractor
{
    private readonly TimeProvider _timeProvider;

    public PeHardeningExtractor(TimeProvider? timeProvider = null)
    {
        _timeProvider = timeProvider ?? TimeProvider.System;
    }

    // PE magic bytes: MZ (DOS header)
    private const ushort DOS_MAGIC = 0x5A4D; // "MZ"
    private const uint PE_SIGNATURE = 0x00004550; // "PE\0\0"

    // PE Optional Header magic values
    private const ushort PE32_MAGIC = 0x10B;
    private const ushort PE32PLUS_MAGIC = 0x20B;

    // DllCharacteristics flags (PE32/PE32+)
    private const ushort IMAGE_DLLCHARACTERISTICS_HIGH_ENTROPY_VA = 0x0020;
    private const ushort IMAGE_DLLCHARACTERISTICS_DYNAMIC_BASE = 0x0040;    // ASLR
    private const ushort IMAGE_DLLCHARACTERISTICS_FORCE_INTEGRITY = 0x0080;
    private const ushort IMAGE_DLLCHARACTERISTICS_NX_COMPAT = 0x0100;       // DEP
    private const ushort IMAGE_DLLCHARACTERISTICS_NO_SEH = 0x0400;
    private const ushort IMAGE_DLLCHARACTERISTICS_GUARD_CF = 0x4000;        // CFG

    // Data Directory indices
    private const int IMAGE_DIRECTORY_ENTRY_SECURITY = 4;   // Authenticode certificate
    private const int IMAGE_DIRECTORY_ENTRY_LOAD_CONFIG = 10;

    /// <inheritdoc />
    public BinaryFormat SupportedFormat => BinaryFormat.Pe;

    /// <inheritdoc />
    public bool CanExtract(string path)
    {
        var ext = Path.GetExtension(path).ToLowerInvariant();
        return ext is ".exe" or ".dll" or ".sys" or ".ocx" or ".scr";
    }

    /// <inheritdoc />
    public bool CanExtract(ReadOnlySpan<byte> header)
    {
        if (header.Length < 2) return false;
        var magic = BinaryPrimitives.ReadUInt16LittleEndian(header);
        return magic == DOS_MAGIC;
    }

    /// <inheritdoc />
    public async Task<BinaryHardeningFlags> ExtractAsync(string path, string digest, CancellationToken ct = default)
    {
        await using var stream = File.OpenRead(path);
        return await ExtractAsync(stream, path, digest, ct);
    }

    /// <inheritdoc />
    public async Task<BinaryHardeningFlags> ExtractAsync(Stream stream, string path, string digest, CancellationToken ct = default)
    {
        var flags = new List<HardeningFlag>();
        var missing = new List<string>();

        // Read full file into memory for parsing
        using var ms = new MemoryStream();
        await stream.CopyToAsync(ms, ct);
        var peData = ms.ToArray();

        if (peData.Length < 64) // Minimum DOS header size
        {
            return CreateResult(path, digest, [], ["Invalid PE: too small"]);
        }

        // Validate DOS header
        var dosMagic = BinaryPrimitives.ReadUInt16LittleEndian(peData.AsSpan(0, 2));
        if (dosMagic != DOS_MAGIC)
        {
            return CreateResult(path, digest, [], ["Invalid PE: bad DOS magic"]);
        }

        // Get PE header offset from DOS header (e_lfanew at offset 0x3C)
        var peOffset = BinaryPrimitives.ReadInt32LittleEndian(peData.AsSpan(0x3C, 4));
        if (peOffset < 0 || peOffset + 24 > peData.Length)
        {
            return CreateResult(path, digest, [], ["Invalid PE: bad PE offset"]);
        }

        // Validate PE signature
        var peSignature = BinaryPrimitives.ReadUInt32LittleEndian(peData.AsSpan(peOffset, 4));
        if (peSignature != PE_SIGNATURE)
        {
            return CreateResult(path, digest, [], ["Invalid PE: bad PE signature"]);
        }

        // Parse COFF header (20 bytes after PE signature)
        var coffOffset = peOffset + 4;
        var machine = BinaryPrimitives.ReadUInt16LittleEndian(peData.AsSpan(coffOffset, 2));
        var numberOfSections = BinaryPrimitives.ReadUInt16LittleEndian(peData.AsSpan(coffOffset + 2, 2));
        var characteristics = BinaryPrimitives.ReadUInt16LittleEndian(peData.AsSpan(coffOffset + 18, 2));

        // Parse Optional Header
        var optionalHeaderOffset = coffOffset + 20;
        if (optionalHeaderOffset + 2 > peData.Length)
        {
            return CreateResult(path, digest, [], ["Invalid PE: truncated optional header"]);
        }

        var optionalMagic = BinaryPrimitives.ReadUInt16LittleEndian(peData.AsSpan(optionalHeaderOffset, 2));
        var isPe32Plus = optionalMagic == PE32PLUS_MAGIC;

        // DllCharacteristics offset differs between PE32 and PE32+
        var dllCharacteristicsOffset = optionalHeaderOffset + (isPe32Plus ? 70 : 70);
        if (dllCharacteristicsOffset + 2 > peData.Length)
        {
            return CreateResult(path, digest, [], ["Invalid PE: truncated DllCharacteristics"]);
        }

        var dllCharacteristics = BinaryPrimitives.ReadUInt16LittleEndian(peData.AsSpan(dllCharacteristicsOffset, 2));

        // === TASK SDIFF-BIN-011: Parse DllCharacteristics ===

        // ASLR (Dynamic Base)
        var hasAslr = (dllCharacteristics & IMAGE_DLLCHARACTERISTICS_DYNAMIC_BASE) != 0;
        flags.Add(new HardeningFlag(HardeningFlagType.Aslr, hasAslr, hasAslr ? "enabled" : null, "DllCharacteristics"));
        if (!hasAslr) missing.Add("ASLR");

        // High Entropy VA (64-bit ASLR)
        var hasHighEntropyVa = (dllCharacteristics & IMAGE_DLLCHARACTERISTICS_HIGH_ENTROPY_VA) != 0;
        flags.Add(new HardeningFlag(HardeningFlagType.HighEntropyVa, hasHighEntropyVa, hasHighEntropyVa ? "enabled" : null, "DllCharacteristics"));
        if (!hasHighEntropyVa && isPe32Plus) missing.Add("HIGH_ENTROPY_VA");

        // DEP (NX Compatible)
        var hasDep = (dllCharacteristics & IMAGE_DLLCHARACTERISTICS_NX_COMPAT) != 0;
        flags.Add(new HardeningFlag(HardeningFlagType.Dep, hasDep, hasDep ? "enabled" : null, "DllCharacteristics"));
        if (!hasDep) missing.Add("DEP");

        // CFG (Control Flow Guard)
        var hasCfg = (dllCharacteristics & IMAGE_DLLCHARACTERISTICS_GUARD_CF) != 0;
        flags.Add(new HardeningFlag(HardeningFlagType.Cfg, hasCfg, hasCfg ? "enabled" : null, "DllCharacteristics"));
        if (!hasCfg) missing.Add("CFG");

        // Force Integrity
        var hasForceIntegrity = (dllCharacteristics & IMAGE_DLLCHARACTERISTICS_FORCE_INTEGRITY) != 0;
        flags.Add(new HardeningFlag(HardeningFlagType.ForceIntegrity, hasForceIntegrity, hasForceIntegrity ? "enabled" : null, "DllCharacteristics"));

        // NO_SEH flag (indicates SafeSEH is not used, but NO_SEH means no SEH at all which is okay)
        var noSeh = (dllCharacteristics & IMAGE_DLLCHARACTERISTICS_NO_SEH) != 0;
        // SafeSEH is only for 32-bit binaries
        if (!isPe32Plus)
        {
            // For 32-bit, NO_SEH is acceptable (no SEH = can't exploit SEH)
            // If SEH is used, we'd need to check Load Config for SafeSEH
            var safeSehStatus = noSeh ? "no_seh" : "needs_verification";
            flags.Add(new HardeningFlag(HardeningFlagType.SafeSeh, noSeh, safeSehStatus, "DllCharacteristics"));
            if (!noSeh) missing.Add("SAFE_SEH");
        }

        // === TASK SDIFF-BIN-012: Authenticode Detection ===
        var hasAuthenticode = CheckAuthenticode(peData, optionalHeaderOffset, isPe32Plus);
        flags.Add(new HardeningFlag(HardeningFlagType.Authenticode, hasAuthenticode, hasAuthenticode ? "signed" : null, "Security Directory"));
        if (!hasAuthenticode) missing.Add("AUTHENTICODE");

        // GS (/GS buffer security check) - check Load Config for SecurityCookie
        var hasGs = CheckGsBufferSecurity(peData, optionalHeaderOffset, isPe32Plus);
        flags.Add(new HardeningFlag(HardeningFlagType.Gs, hasGs, hasGs ? "enabled" : null, "Load Config"));
        if (!hasGs) missing.Add("GS");

        return CreateResult(path, digest, flags, missing);
    }

    /// <summary>
    /// Check if PE has Authenticode signature by examining Security Directory.
    /// </summary>
    private static bool CheckAuthenticode(byte[] peData, int optionalHeaderOffset, bool isPe32Plus)
    {
        try
        {
            // Data directories start after the standard optional header fields
            // PE32: offset 96 from optional header start
            // PE32+: offset 112 from optional header start
            var dataDirectoriesOffset = optionalHeaderOffset + (isPe32Plus ? 112 : 96);
            
            // Security directory is index 4 (each entry is 8 bytes: 4 for RVA, 4 for size)
            var securityDirOffset = dataDirectoriesOffset + (IMAGE_DIRECTORY_ENTRY_SECURITY * 8);
            
            if (securityDirOffset + 8 > peData.Length)
                return false;

            var securityRva = BinaryPrimitives.ReadUInt32LittleEndian(peData.AsSpan(securityDirOffset, 4));
            var securitySize = BinaryPrimitives.ReadUInt32LittleEndian(peData.AsSpan(securityDirOffset + 4, 4));

            // If security directory has non-zero size, there's a certificate
            return securitySize > 0 && securityRva > 0;
        }
        catch
        {
            return false;
        }
    }

    /// <summary>
    /// Check for /GS buffer security by examining Load Config Directory.
    /// </summary>
    private static bool CheckGsBufferSecurity(byte[] peData, int optionalHeaderOffset, bool isPe32Plus)
    {
        try
        {
            var dataDirectoriesOffset = optionalHeaderOffset + (isPe32Plus ? 112 : 96);
            var loadConfigDirOffset = dataDirectoriesOffset + (IMAGE_DIRECTORY_ENTRY_LOAD_CONFIG * 8);
            
            if (loadConfigDirOffset + 8 > peData.Length)
                return false;

            var loadConfigRva = BinaryPrimitives.ReadUInt32LittleEndian(peData.AsSpan(loadConfigDirOffset, 4));
            var loadConfigSize = BinaryPrimitives.ReadUInt32LittleEndian(peData.AsSpan(loadConfigDirOffset + 4, 4));

            // If load config exists and has reasonable size, /GS is likely enabled
            // (Full verification would require parsing the Load Config structure)
            return loadConfigSize >= 64 && loadConfigRva > 0;
        }
        catch
        {
            return false;
        }
    }

    private BinaryHardeningFlags CreateResult(
        string path,
        string digest,
        List<HardeningFlag> flags,
        List<string> missing)
    {
        // Calculate score: enabled flags / total expected flags
        var positiveFlags = new[] { 
            HardeningFlagType.Aslr, 
            HardeningFlagType.Dep, 
            HardeningFlagType.Cfg, 
            HardeningFlagType.Authenticode,
            HardeningFlagType.Gs
        };
        
        var enabledCount = flags.Count(f => f.Enabled && positiveFlags.Contains(f.Name));
        var totalExpected = positiveFlags.Length;
        var score = totalExpected > 0 ? (double)enabledCount / totalExpected : 0.0;

        return new BinaryHardeningFlags(
            Format: BinaryFormat.Pe,
            Path: path,
            Digest: digest,
            Flags: [.. flags],
            HardeningScore: Math.Round(score, 2),
            MissingFlags: [.. missing],
            ExtractedAt: _timeProvider.GetUtcNow());
    }
}