using System.Security.Principal;
using Agent.Windows.Recovery;
using Agent.Windows.Security;

namespace Agent.Windows.Network;

public sealed class TaskNetworkPolicy : IDisposable
{
    private readonly WfpDynamicSession _session;
    private readonly TaskResourceJournal _journal;
    private readonly ISecurityAuditSink _audit;
    private readonly Guid _entryId;
    private readonly SecurityIdentifier _appContainerSid;
    private readonly Guid? _loopbackEntryId;
    private readonly bool _loopbackAdded;
    private readonly Guid? _sessionId;
    private bool _disposed;

    private TaskNetworkPolicy(WfpDynamicSession session, TaskResourceJournal journal, ISecurityAuditSink audit, Guid entryId,
        SecurityIdentifier appContainerSid, Guid? loopbackEntryId, bool loopbackAdded, Guid? sessionId)
    {
        _session = session; _journal = journal; _audit = audit; _entryId = entryId;
        _appContainerSid = appContainerSid; _loopbackEntryId = loopbackEntryId; _loopbackAdded = loopbackAdded;
        _sessionId = sessionId;
    }

    public static TaskNetworkPolicy Apply(TaskResourceJournal journal, SecurityIdentifier appContainerSid, IReadOnlyList<NetworkEndpoint> allowances,
        ISecurityAuditSink audit, Guid? sessionId = null, bool enableLoopback = true)
    {
        bool originallyExempt = enableLoopback && LoopbackExemption.Contains(appContainerSid);
        Guid? loopbackEntryId = enableLoopback ? journal.Prepare("network.loopback_exemption", appContainerSid.Value,
            new Dictionary<string, string> { ["originally_present"] = originallyExempt.ToString() }) : null;
        bool loopbackAdded = false;
        Guid entryId = journal.Prepare("wfp.dynamic", appContainerSid.Value, new Dictionary<string, string>
        {
            ["session"] = "dynamic", ["allowances"] = string.Join(',', allowances.Select(static item => $"{item.Address}:{item.Port}"))
        });
        WfpDynamicSession session = WfpDynamicSession.Open();
        try
        {
            if (enableLoopback)
            {
                loopbackAdded = LoopbackExemption.Ensure(appContainerSid);
                journal.MarkApplied(loopbackEntryId!.Value, new Dictionary<string, string> { ["added"] = loopbackAdded.ToString() });
            }
            IReadOnlyList<ulong> ids = session.ApplyDefaultBlock(appContainerSid, allowances);
            journal.MarkApplied(entryId, new Dictionary<string, string> { ["filter_ids"] = string.Join(',', ids) });
            audit.Write(new(DateTimeOffset.UtcNow, journal.TaskId, "network.rule.apply", "applied", new Dictionary<string, string>
            {
                ["ipv4_ipv6_default_block"] = "true", ["filter_count"] = ids.Count.ToString(), ["session"] = "dynamic"
            }, SessionId: sessionId));
            if (enableLoopback) audit.Write(new(DateTimeOffset.UtcNow, journal.TaskId, "network.loopback.apply", loopbackAdded ? "added" : "preserved",
                new Dictionary<string, string> { ["app_container_sid"] = appContainerSid.Value }, SessionId: sessionId));
            return new(session, journal, audit, entryId, appContainerSid, loopbackEntryId, loopbackAdded, sessionId);
        }
        catch
        {
            session.Dispose();
            journal.MarkCleaned(entryId);
            if (loopbackAdded) LoopbackExemption.Remove(appContainerSid);
            if (loopbackEntryId is not null) journal.MarkCleaned(loopbackEntryId.Value);
            throw;
        }
    }

    public void Dispose()
    {
        if (_disposed) return;
        _session.RemovePolicy();
        _session.Dispose();
        _journal.MarkCleaned(_entryId);
        if (_loopbackAdded) LoopbackExemption.Remove(_appContainerSid);
        if (_loopbackEntryId is not null) _journal.MarkCleaned(_loopbackEntryId.Value);
        _audit.Write(new(DateTimeOffset.UtcNow, _journal.TaskId, "network.rule.remove", "removed", new Dictionary<string, string> { ["residual_filters"] = "0" }, SessionId: _sessionId));
        if (_loopbackEntryId is not null) _audit.Write(new(DateTimeOffset.UtcNow, _journal.TaskId, "network.loopback.remove", "removed",
            new Dictionary<string, string> { ["residual_exemptions"] = LoopbackExemption.Contains(_appContainerSid) ? "preexisting" : "0" }, SessionId: _sessionId));
        _disposed = true;
    }
}