1 files changed, 75 insertions, 0 deletions
diff --git a/tests/Agent.Windows.Tests/DirectElevatedNetworkSession.cs b/tests/Agent.Windows.Tests/DirectElevatedNetworkSession.cs
new file mode 100644
index 0000000..b194ed1
--- /dev/null
+++ b/tests/Agent.Windows.Tests/DirectElevatedNetworkSession.cs
@@ -0,0 +1,75 @@
+using System.Net;
+using System.Security.Principal;
+using Agent.Admin.Client;
+using Agent.Admin.Host;
+
+namespace Agent.Windows.Tests;
+
+internal sealed class DirectElevatedNetworkSession : IElevatedNetworkSession
+{
+    private WfpDynamicSession? _wfp;
+    private SecurityIdentifier? _sid;
+    private bool _loopbackAdded;
+
+    public Task<bool> InspectLoopbackAsync(string applicationSid, CancellationToken cancellationToken) =>
+        Task.FromResult(LoopbackExemption.Contains(new SecurityIdentifier(applicationSid)));
+
+    public Task ApplyAsync(string applicationSid, string networkAccess, IReadOnlyList<AdminNetworkEndpoint> endpoints,
+        bool originallyLoopbackExempt, CancellationToken cancellationToken)
+    {
+        _sid = new(applicationSid);
+        NetworkEndpoint[] values = endpoints.Select(static item => new NetworkEndpoint(IPAddress.Parse(item.Address), (ushort)item.Port)).ToArray();
+        _wfp = WfpDynamicSession.Open();
+        _wfp.ApplyDefaultBlock(_sid, values);
+        if (values.Any(static item => IPAddress.IsLoopback(item.Address)) && !originallyLoopbackExempt)
+            _loopbackAdded = LoopbackExemption.Ensure(_sid);
+        return Task.CompletedTask;
+    }
+
+    public Task RecoverLoopbackAsync(string applicationSid, bool originallyLoopbackExempt, CancellationToken cancellationToken)
+    {
+        var sid = new SecurityIdentifier(applicationSid);
+        if (originallyLoopbackExempt) _ = LoopbackExemption.Ensure(sid); else LoopbackExemption.Remove(sid);
+        return Task.CompletedTask;
+    }
+
+    public Task CleanupAsync(CancellationToken cancellationToken)
+    {
+        if (_loopbackAdded && _sid is not null) LoopbackExemption.Remove(_sid);
+        _loopbackAdded = false;
+        if (_wfp is not null) { _wfp.RemovePolicy(); _wfp.Dispose(); _wfp = null; }
+        return Task.CompletedTask;
+    }
+
+    public async ValueTask DisposeAsync() => await CleanupAsync(CancellationToken.None).ConfigureAwait(false);
+}
+
+internal sealed class DirectElevatedNetworkSessionFactory : IElevatedNetworkSessionFactory
+{
+    public int Starts { get; private set; }
+    public Task<IElevatedNetworkSession> StartAsync(Guid taskId, CancellationToken cancellationToken)
+    {
+        Starts++;
+        return Task.FromResult<IElevatedNetworkSession>(new DirectElevatedNetworkSession());
+    }
+}
+
+internal sealed class NoOpElevatedNetworkSessionFactory : IElevatedNetworkSessionFactory
+{
+    public int Starts { get; private set; }
+    public Task<IElevatedNetworkSession> StartAsync(Guid taskId, CancellationToken cancellationToken)
+    {
+        Starts++;
+        return Task.FromResult<IElevatedNetworkSession>(new NoOpElevatedNetworkSession());
+    }
+
+    private sealed class NoOpElevatedNetworkSession : IElevatedNetworkSession
+    {
+        public Task<bool> InspectLoopbackAsync(string applicationSid, CancellationToken cancellationToken) => Task.FromResult(false);
+        public Task ApplyAsync(string applicationSid, string networkAccess, IReadOnlyList<AdminNetworkEndpoint> endpoints,
+            bool originallyLoopbackExempt, CancellationToken cancellationToken) => Task.CompletedTask;
+        public Task RecoverLoopbackAsync(string applicationSid, bool originallyLoopbackExempt, CancellationToken cancellationToken) => Task.CompletedTask;
+        public Task CleanupAsync(CancellationToken cancellationToken) => Task.CompletedTask;
+        public ValueTask DisposeAsync() => ValueTask.CompletedTask;
+    }
+}