using System.Security.Cryptography;
using System.Text;
using System.Text.Json;
using Mono.Cecil;
using Mono.Cecil.Cil;

namespace XFESystem.Tools.TlsEntropyPatch;

public static class Program
{
    public const string OriginalSha256 = "ef92fc661e8d7ba8ec4d39a7cdfcdba41f14cb327f0d24fe79d0b5bb428899b5";
    private const string AdapterType = "XFESystem.NativeEntropy.NativeEntropy";
    private const string BclType = "System.Security.Cryptography.RandomNumberGenerator";
    public static readonly string[] SourcePaths = ["Tools/TlsEntropyPatch/Program.cs", "Tools/TlsEntropyPatch/TlsEntropyPatch.csproj",
        "Tools/TlsEntropyPatch/Packages.lock.json", "Tools/TlsEntropyPatch/NativeEntropy.targets",
        "Tools/TlsEntropyPatch/Adapter/NativeEntropy.cs", "Tools/TlsEntropyPatch/Adapter/XFESystem.NativeEntropy.csproj",
        "Tools/TlsEntropyPatch/Adapter/Packages.lock.json"];
    public static int Main(string[] args)
    {
        if (args.Length != 3) throw new ArgumentException("TlsEntropyPatch <repository root> <private BC dll> <private native entropy adapter dll>");
        Apply(args[0], args[1], args[2]); return 0;
    }
    public static string Hash(string path) => Convert.ToHexStringLower(SHA256.HashData(File.ReadAllBytes(path)));
    public static IEnumerable<TypeDefinition> Types(IEnumerable<TypeDefinition> source)
    { foreach (var type in source) { yield return type; foreach (var nested in Types(type.NestedTypes)) yield return nested; } }
    public static void Apply(string repository, string path, string adapterPath)
    {
        repository = Path.GetFullPath(repository); path = Path.GetFullPath(path); adapterPath = Path.GetFullPath(adapterPath);
        VerifyPrivate(repository, path);
        if (Path.GetDirectoryName(path) != Path.GetDirectoryName(adapterPath) || Path.GetFileName(adapterPath) != "XFESystem.NativeEntropy.dll")
            throw new InvalidOperationException("The adapter must be in the same private application directory.");
        var sourceHashes = SourcePaths.ToDictionary(item => item, item => Hash(Path.Combine(repository, item)));
        string adapterHash = Hash(adapterPath), before = Hash(path), receiptPath = path + ".XfeEntropy.json";
        using var adapter = AssemblyDefinition.ReadAssembly(adapterPath, new ReaderParameters { InMemory = true });
        if (adapter.Name.Name != "XFESystem.NativeEntropy") throw new InvalidOperationException("Unexpected entropy adapter assembly.");
        var feeder = adapter.MainModule.GetType(AdapterType) ?? throw new InvalidOperationException("Missing explicit entropy feeder.");
        var imports = Types(adapter.MainModule.Types).SelectMany(type => type.Methods).Where(method => method.HasPInvokeInfo).ToArray();
        if (imports.Length != 1 || imports[0].PInvokeInfo.Module.Name != "xfe_pal" || imports[0].PInvokeInfo.EntryPoint != "xfe_runtime_random_fill" ||
            imports[0].ReturnType.FullName != "System.Int32" || !imports[0].Parameters.Select(parameter => parameter.ParameterType.FullName).SequenceEqual(["System.Byte*", "System.UInt32"]))
            throw new InvalidOperationException("The entropy adapter must expose exactly the checked native hardware ABI.");
        using var assembly = AssemblyDefinition.ReadAssembly(path, new ReaderParameters { InMemory = true });
        if (assembly.Name.Name != "BouncyCastle.Cryptography" || assembly.Name.Version != new Version(2, 0, 0, 0))
            throw new InvalidOperationException("Only the pinned BC 2.7.0 net6.0 payload is supported.");
        var methods = Types(assembly.MainModule.Types).SelectMany(type => type.Methods).Where(method => method.HasBody).ToArray();
        var calls = methods.SelectMany(method => method.Body.Instructions.Select(instruction => (method, instruction)))
            .Where(item => item.instruction.Operand is MethodReference target && target.DeclaringType.FullName == BclType).ToArray();
        if (before != OriginalSha256)
        {
            if (!File.Exists(receiptPath)) throw new InvalidOperationException("Unknown BC bytes or missing native entropy receipt.");
            using var previous = JsonDocument.Parse(File.ReadAllBytes(receiptPath)); var record = previous.RootElement;
            if (record.GetProperty("Kind").GetString() != "XfeNativeBouncyCastleEntropyPortV1" || record.GetProperty("BeforeSha256").GetString() != OriginalSha256 ||
                record.GetProperty("AfterSha256").GetString() != before || record.GetProperty("AdapterSha256").GetString() != adapterHash ||
                record.GetProperty("ToolSha256").GetString() != Hash(typeof(Program).Assembly.Location) || calls.Length != 0 ||
                methods.Sum(method => method.Body.Instructions.Count(instruction => instruction.Operand is MethodReference reference && reference.DeclaringType.FullName == AdapterType)) != 12)
                throw new InvalidOperationException("The private BC entropy port differs from its verification receipt.");
            foreach (var item in sourceHashes)
                if (record.GetProperty("SourceSha256").GetProperty(item.Key).GetString() != item.Value)
                    throw new InvalidOperationException("The entropy port source changed after patching.");
            Console.WriteLine("XFE_NATIVE_BC_ENTROPY_ALREADY_READY|" + before); return;
        }
        if (calls.Length != 12 || calls.Count(item => ((MethodReference)item.instruction.Operand).Name == "Fill") != 4 ||
            calls.Count(item => ((MethodReference)item.instruction.Operand).Name == "Create") != 2 ||
            calls.Count(item => ((MethodReference)item.instruction.Operand).Name == "GetBytes") != 5 ||
            calls.Count(item => ((MethodReference)item.instruction.Operand).Name == "Dispose") != 1)
            throw new InvalidOperationException("The pinned BC RNG call shape differs.");
        var unchangedAudit = methods.ToDictionary(method => method.FullName, Fingerprint);
        var reports = new List<object>();
        foreach (var (method, instruction) in calls)
        {
            var original = (MethodReference)instruction.Operand;
            if (instruction.OpCode != (original.HasThis ? OpCodes.Callvirt : OpCodes.Call))
                throw new InvalidOperationException("Unexpected BC RNG dispatch shape.");
            var parameters = original.Parameters.Select(parameter => parameter.ParameterType.FullName).ToList();
            if (original.HasThis) parameters.Insert(0, BclType);
            var target = feeder.Methods.Single(candidate => candidate.IsStatic && candidate.Name == original.Name &&
                candidate.ReturnType.FullName == original.ReturnType.FullName && candidate.Parameters.Select(parameter => parameter.ParameterType.FullName).SequenceEqual(parameters));
            reports.Add(new { Method = method.FullName, Offset = instruction.Offset, OriginalCall = original.FullName, ReplacementCall = target.FullName });
            // Both call forms have exactly the same evaluation-stack inputs and
            // result. No instruction offsets/branches/EH, keys or policy change.
            instruction.OpCode = OpCodes.Call; instruction.Operand = assembly.MainModule.ImportReference(target);
        }
        string temporary = path + ".XfeEntropy.tmp";
        assembly.Write(temporary);
        using (var written = AssemblyDefinition.ReadAssembly(temporary))
        {
            var audited = Types(written.MainModule.Types).SelectMany(type => type.Methods).Where(method => method.HasBody).ToArray();
            if (audited.Length != unchangedAudit.Count || audited.Any(method => !unchangedAudit.TryGetValue(method.FullName, out var expected) || expected != Fingerprint(method)))
                throw new InvalidOperationException("A BC method changed outside the exact entropy-call substitution.");
        }
        File.Move(temporary, path, true);
        var receipt = new { Kind = "XfeNativeBouncyCastleEntropyPortV1", BouncyCastleVersion = "2.7.0",
            PackageSha256 = "f091ffccab4d03993e660bace277659a79dee0972f54d7f1f4bd46d680966241",
            UpstreamCommit = "4007498b13582d90ee1eda5d9920c324428b98b3", BeforeSha256 = before, AfterSha256 = Hash(path),
            AdapterSha256 = adapterHash, ToolSha256 = Hash(typeof(Program).Assembly.Location), SourceSha256 = sourceHashes,
            CallCount = reports.Count, Calls = reports, NativeEntropyAbi = "xfe_pal::xfe_runtime_random_fill; 1 success, all other results clear and throw",
            MethodsAudited = unchangedAudit.Count, OnlyEntropyCallsChanged = true,
            MaximumEntropyRequestBytes = 1048576, CryptoPolicyModified = false, GlobalPackageCacheModified = false,
            OpenSslInitializationEmulated = false, HostRandomFallback = false, GuestAcceptancePassed = false,
            DerivedPrivateNativeOnly = true, OriginalStrongNameSignatureNotAsserted = true };
        File.WriteAllText(receiptPath, JsonSerializer.Serialize(receipt, new JsonSerializerOptions { WriteIndented = true }));
        Console.WriteLine("XFE_NATIVE_BC_ENTROPY_READY|12|" + receipt.AfterSha256);
    }
    private static string Fingerprint(MethodDefinition method)
    {
        var text = new StringBuilder().Append(method.Body.InitLocals).Append('|').Append(method.Body.MaxStackSize);
        foreach (var variable in method.Body.Variables) text.Append('|').Append(variable.VariableType.FullName);
        foreach (var instruction in method.Body.Instructions)
        {
            text.Append('|').Append(instruction.Offset).Append(':');
            if (instruction.Operand is MethodReference target && target.DeclaringType.FullName is BclType or AdapterType)
            {
                text.Append("Entropy:").Append(target.Name).Append(':').Append(target.ReturnType.FullName);
                if (target.HasThis) text.Append(',').Append(BclType);
                foreach (var parameter in target.Parameters) text.Append(',').Append(parameter.ParameterType.FullName);
            }
            else
            {
                text.Append(instruction.OpCode.Code).Append(':');
                if (instruction.Operand is Instruction branch) text.Append(branch.Offset);
                else if (instruction.Operand is Instruction[] branches) foreach (var branchTarget in branches) text.Append(branchTarget.Offset).Append(',');
                else text.Append(instruction.Operand);
            }
        }
        foreach (var handler in method.Body.ExceptionHandlers)
            text.Append('|').Append(handler.HandlerType).Append(':').Append(handler.CatchType?.FullName).Append(':')
                .Append(handler.TryStart?.Offset).Append(':').Append(handler.TryEnd?.Offset).Append(':')
                .Append(handler.HandlerStart?.Offset).Append(':').Append(handler.HandlerEnd?.Offset).Append(':').Append(handler.FilterStart?.Offset);
        return Convert.ToHexStringLower(SHA256.HashData(Encoding.UTF8.GetBytes(text.ToString())));
    }
    private static void VerifyPrivate(string repository, string path)
    {
        if (!File.Exists(Path.Combine(repository, "Source/XFESystem.NativeRuntime/Application/NativeApplication.csproj")) ||
            Path.GetFileName(path) != "BouncyCastle.Cryptography.dll" ||
            !(path.StartsWith(Path.Combine(repository, "build") + Path.DirectorySeparatorChar, StringComparison.OrdinalIgnoreCase) ||
              path.StartsWith(Path.Combine(repository, "Source/XFESystem.NativeRuntime/Application/bin") + Path.DirectorySeparatorChar, StringComparison.OrdinalIgnoreCase)))
            throw new InvalidOperationException("Only a repository-private native application BC output can be patched.");
        for (string? current = path; current is not null && !current.Equals(repository, StringComparison.OrdinalIgnoreCase); current = Path.GetDirectoryName(current))
            if ((File.GetAttributes(current) & FileAttributes.ReparsePoint) != 0)
                throw new InvalidOperationException("A private output cannot traverse a symbolic link or junction.");
    }
}