// -----------------------------------------------------------------------------
// MachoHardeningExtractor.cs
// Sprint: SPRINT_3500_0004_0001_smart_diff_binary_output
// Task: SDIFF-BIN-013a - Implement MachO hardening extractor (bonus)
// Description: Extracts security hardening flags from macOS Mach-O binaries
// -----------------------------------------------------------------------------

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

namespace StellaOps.Scanner.Analyzers.Native.Hardening;

/// <summary>
/// Extracts hardening flags from macOS Mach-O binaries.
/// Detects PIE, code signing, RESTRICT, hardened runtime, and more.
/// Per Sprint 3500.4 - Smart-Diff Binary Analysis.
/// </summary>
public sealed class MachoHardeningExtractor : IHardeningExtractor
{
    private readonly TimeProvider _timeProvider;

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

    // Mach-O magic numbers
    private const uint MH_MAGIC = 0xFEEDFACE;    // 32-bit
    private const uint MH_CIGAM = 0xCEFAEDFE;    // 32-bit (reversed)
    private const uint MH_MAGIC_64 = 0xFEEDFACF; // 64-bit
    private const uint MH_CIGAM_64 = 0xCFFAEDFE; // 64-bit (reversed)
    private const uint FAT_MAGIC = 0xCAFEBABE;   // Universal binary
    private const uint FAT_CIGAM = 0xBEBAFECA;   // Universal (reversed)

    // Mach-O header flags (from mach/loader.h)
    private const uint MH_PIE = 0x00200000;                 // Position Independent Executable
    private const uint MH_NO_HEAP_EXECUTION = 0x01000000;   // No heap execution
    private const uint MH_ALLOW_STACK_EXECUTION = 0x00020000; // Allow stack execution (bad!)
    private const uint MH_NOFIXPREBINDING = 0x00000400;
    private const uint MH_TWOLEVEL = 0x00000080;            // Two-level namespace

    // Load command types
    private const uint LC_SEGMENT = 0x01;
    private const uint LC_SEGMENT_64 = 0x19;
    private const uint LC_CODE_SIGNATURE = 0x1D;
    private const uint LC_ENCRYPTION_INFO = 0x21;
    private const uint LC_ENCRYPTION_INFO_64 = 0x2C;
    private const uint LC_DYLD_INFO = 0x22;
    private const uint LC_DYLD_INFO_ONLY = 0x80000022;
    private const uint LC_DYLIB_CODE_SIGN_DRS = 0x2F;
    private const uint LC_BUILD_VERSION = 0x32;
    private const uint LC_RPATH = 0x8000001C;

    // Segment flags
    private const uint SG_PROTECTED_VERSION_1 = 0x08;

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

    /// <inheritdoc />
    public bool CanExtract(string path)
    {
        var ext = Path.GetExtension(path).ToLowerInvariant();
        // Mach-O can be .dylib, .bundle, or extensionless executables
        return ext is ".dylib" or ".bundle" or ".framework" or ""
            || Path.GetFileName(path).StartsWith("lib", StringComparison.OrdinalIgnoreCase);
    }

    /// <inheritdoc />
    public bool CanExtract(ReadOnlySpan<byte> header)
    {
        if (header.Length < 4) return false;
        var magic = BinaryPrimitives.ReadUInt32BigEndian(header);
        return magic is MH_MAGIC or MH_CIGAM or MH_MAGIC_64 or MH_CIGAM_64 or FAT_MAGIC or FAT_CIGAM;
    }

    /// <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
        using var ms = new MemoryStream();
        await stream.CopyToAsync(ms, ct);
        var data = ms.ToArray();

        if (data.Length < 28) // Minimum Mach-O header
        {
            return CreateResult(path, digest, [], ["Invalid Mach-O: too small"]);
        }

        // Check magic and determine endianness
        var magic = BinaryPrimitives.ReadUInt32BigEndian(data.AsSpan(0, 4));
        var isLittleEndian = magic is MH_CIGAM or MH_CIGAM_64;
        var is64Bit = magic is MH_MAGIC_64 or MH_CIGAM_64;

        // Handle universal binaries - just extract first architecture for now
        if (magic is FAT_MAGIC or FAT_CIGAM)
        {
            var fatResult = ExtractFromFat(data, path, digest);
            if (fatResult is not null)
                return fatResult;
            return CreateResult(path, digest, [], ["Universal binary: no supported architectures"]);
        }

        // Normalize magic
        magic = isLittleEndian
            ? BinaryPrimitives.ReadUInt32LittleEndian(data.AsSpan(0, 4))
            : BinaryPrimitives.ReadUInt32BigEndian(data.AsSpan(0, 4));

        if (magic is not (MH_MAGIC or MH_MAGIC_64))
        {
            return CreateResult(path, digest, [], ["Invalid Mach-O magic"]);
        }

        // Parse header
        var headerSize = is64Bit ? 32 : 28;
        if (data.Length < headerSize)
        {
            return CreateResult(path, digest, [], ["Invalid Mach-O: truncated header"]);
        }

        var headerFlags = ReadUInt32(data, is64Bit ? 24 : 24, isLittleEndian);
        var ncmds = ReadUInt32(data, is64Bit ? 16 : 16, isLittleEndian);
        var sizeofcmds = ReadUInt32(data, is64Bit ? 20 : 20, isLittleEndian);

        // === Check PIE flag ===
        var hasPie = (headerFlags & MH_PIE) != 0;
        flags.Add(new HardeningFlag(HardeningFlagType.Pie, hasPie, hasPie ? "enabled" : null, "MH_FLAGS"));
        if (!hasPie) missing.Add("PIE");

        // === Check for heap execution ===
        var noHeapExec = (headerFlags & MH_NO_HEAP_EXECUTION) != 0;
        flags.Add(new HardeningFlag(HardeningFlagType.Nx, noHeapExec, noHeapExec ? "no_heap_exec" : null, "MH_FLAGS"));

        // === Check for stack execution (inverted - presence is BAD) ===
        var allowsStackExec = (headerFlags & MH_ALLOW_STACK_EXECUTION) != 0;
        if (allowsStackExec)
        {
            flags.Add(new HardeningFlag(HardeningFlagType.Nx, false, "stack_exec_allowed", "MH_FLAGS"));
            missing.Add("NX");
        }

        // === Parse load commands ===
        var hasCodeSignature = false;
        var hasEncryption = false;
        var hasRpath = false;
        var hasHardenedRuntime = false;
        var hasRestrict = false;

        var offset = headerSize;
        for (var i = 0; i < ncmds && offset + 8 <= data.Length; i++)
        {
            var cmd = ReadUInt32(data, offset, isLittleEndian);
            var cmdsize = ReadUInt32(data, offset + 4, isLittleEndian);

            if (cmdsize < 8 || offset + cmdsize > data.Length)
                break;

            switch (cmd)
            {
                case LC_CODE_SIGNATURE:
                    hasCodeSignature = true;
                    break;

                case LC_ENCRYPTION_INFO:
                case LC_ENCRYPTION_INFO_64:
                    // Check if cryptid is non-zero (actually encrypted)
                    var cryptid = ReadUInt32(data, offset + (cmd == LC_ENCRYPTION_INFO_64 ? 16 : 12), isLittleEndian);
                    hasEncryption = cryptid != 0;
                    break;

                case LC_RPATH:
                    hasRpath = true;
                    break;

                case LC_BUILD_VERSION:
                    // Check for hardened runtime flag in build version
                    if (cmdsize >= 24)
                    {
                        var ntools = ReadUInt32(data, offset + 20, isLittleEndian);
                        // Hardened runtime is indicated by certain build flags
                        // This is a simplification - full check requires parsing tool entries
                        hasHardenedRuntime = ntools > 0;
                    }
                    break;

                case LC_SEGMENT:
                case LC_SEGMENT_64:
                    // Check for __RESTRICT segment
                    var nameLen = cmd == LC_SEGMENT_64 ? 16 : 16;
                    if (cmdsize > nameLen + 8)
                    {
                        var segname = System.Text.Encoding.ASCII.GetString(data, offset + 8, nameLen).TrimEnd('\0');
                        if (segname == "__RESTRICT")
                        {
                            hasRestrict = true;
                        }
                    }
                    break;
            }

            offset += (int)cmdsize;
        }

        // Add code signing flag
        flags.Add(new HardeningFlag(HardeningFlagType.Authenticode, hasCodeSignature, hasCodeSignature ? "signed" : null, "LC_CODE_SIGNATURE"));
        if (!hasCodeSignature) missing.Add("CODE_SIGN");

        // Add RESTRICT flag (prevents DYLD_ env vars)
        flags.Add(new HardeningFlag(HardeningFlagType.Restrict, hasRestrict, hasRestrict ? "enabled" : null, "__RESTRICT segment"));

        // Add RPATH flag (presence can be a security concern)
        flags.Add(new HardeningFlag(HardeningFlagType.Rpath, hasRpath, hasRpath ? "present" : null, "LC_RPATH"));

        // Add encryption flag
        if (hasEncryption)
        {
            flags.Add(new HardeningFlag(HardeningFlagType.ForceIntegrity, true, "encrypted", "LC_ENCRYPTION_INFO"));
        }

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

    /// <summary>
    /// Extract hardening info from the first slice of a universal (fat) binary.
    /// </summary>
    private BinaryHardeningFlags? ExtractFromFat(byte[] data, string path, string digest)
    {
        if (data.Length < 8) return null;

        var magic = BinaryPrimitives.ReadUInt32BigEndian(data.AsSpan(0, 4));
        var isLittleEndian = magic == FAT_CIGAM;

        var nfat = ReadUInt32(data, 4, isLittleEndian);
        if (nfat == 0 || data.Length < 8 + nfat * 20)
            return null;

        // Get first architecture offset and size
        var archOffset = ReadUInt32(data, 16, isLittleEndian);
        var archSize = ReadUInt32(data, 20, isLittleEndian);

        if (archOffset + archSize > data.Length)
            return null;

        // Extract first architecture and re-parse
        var sliceData = data.AsSpan((int)archOffset, (int)archSize).ToArray();
        using var sliceStream = new MemoryStream(sliceData);
        return ExtractAsync(sliceStream, path, digest).GetAwaiter().GetResult();
    }

    private static uint ReadUInt32(byte[] data, int offset, bool littleEndian)
    {
        return littleEndian
            ? BinaryPrimitives.ReadUInt32LittleEndian(data.AsSpan(offset, 4))
            : BinaryPrimitives.ReadUInt32BigEndian(data.AsSpan(offset, 4));
    }

    private BinaryHardeningFlags CreateResult(
        string path,
        string digest,
        List<HardeningFlag> flags,
        List<string> missing)
    {
        // Calculate score based on key flags
        var positiveFlags = new[]
        {
            HardeningFlagType.Pie,
            HardeningFlagType.Nx,
            HardeningFlagType.Authenticode, // Code signing
            HardeningFlagType.Restrict
        };

        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.MachO,
            Path: path,
            Digest: digest,
            Flags: [.. flags],
            HardeningScore: Math.Round(score, 2),
            MissingFlags: [.. missing],
            ExtractedAt: _timeProvider.GetUtcNow());
    }
}