path: root/src/Agent.Windows/Tools/WindowsToolBackend.cs

using System.Collections.Concurrent;
using System.IO.Enumeration;
using System.Text;
using System.Text.Json;
using System.Text.RegularExpressions;
using Agent.Core.Tools;
using Agent.Admin.Client;
using Agent.Windows.Network;
using Agent.Windows.Recovery;
using Agent.Windows.Security;

namespace Agent.Windows.Tools;

public sealed record EffectiveSandboxPolicy(
    IReadOnlyList<FilePolicyEntry> Files,
    IReadOnlyList<ApprovedEntryExecutable> EntryExecutables,
    IReadOnlyList<string> RuntimeRoots,
    IReadOnlyList<string> ArtifactRoots,
    SandboxJobLimits JobLimits,
    string NetworkAccess,
    IReadOnlyList<(string Host, IReadOnlyList<int> Ports)> NetworkAllowances,
    string JournalDirectory,
    string BackupDirectory,
    long MemoryBackupLimitBytes,
    long WholeReadLimitBytes,
    int SearchResultLimit,
    int TailLineLimit,
    int ProcessTimeoutSeconds = 300,
    string ResourceProfile = "normal",
    IReadOnlyList<string>? BaselineCapabilities = null);

public sealed class WindowsToolBackend(EffectiveSandboxPolicy settings, ISecurityAuditSink audit, ITaskUpdateHandler? taskUpdates = null,
    IElevatedNetworkSessionFactory? elevatedNetwork = null) : IInternalToolBackend
{
    private readonly ConcurrentDictionary<Guid, TaskContext> _contexts = new();
    private readonly PathPolicy _paths = new(settings.Files);
    private readonly EntryExecutablePolicy _entryExecutables = new(settings.EntryExecutables, settings.RuntimeRoots, settings.ArtifactRoots);
    private readonly ITaskUpdateHandler _taskUpdates = taskUpdates ?? new TaskGraphToolHandler();
    private readonly IElevatedNetworkSessionFactory _elevatedNetwork = elevatedNetwork ??
        new ElevatedNetworkSessionFactory(Path.Combine(AppContext.BaseDirectory, "admin", "agent-mcp-admin.exe"));
    private bool _cleanupFailed;

    public async Task BeginTaskAsync(Guid taskId, Guid sessionId, CancellationToken cancellationToken)
    {
        if (settings.BaselineCapabilities is not null &&
            !settings.BaselineCapabilities.SequenceEqual(["registryRead", "lpacInstrumentation"], StringComparer.Ordinal))
            throw new InvalidOperationException("Effective Sandbox Policy baseline capabilities do not match the fixed LPAC baseline.");
        if (_cleanupFailed) throw new InvalidOperationException("A prior Task Resource cleanup failed; restart is required for Recovery.");
        if (_contexts.Count != 0) throw new InvalidOperationException("Only one Task Resource scope may be active.");
        AppContainerProfile profile = AppContainerProfile.EnsureFixed();
        TaskResourceJournal journal = TaskResourceJournal.Create(settings.JournalDirectory, taskId);
        var acl = new TaskAclManager(profile.Sid, audit, journal);
        var context = new TaskContext(journal, acl, profile, sessionId);
        _taskUpdates.BeginTask(taskId, sessionId);
        try
        {
            if (!_contexts.TryAdd(taskId, context)) throw new InvalidOperationException("Task context already exists.");
            audit.Write(new(DateTimeOffset.UtcNow, taskId, "task.start", "running", new Dictionary<string, string>(), SessionId: sessionId));
        }
        catch
        {
            try { await CleanupAsync(taskId, context).ConfigureAwait(false); }
            finally { _taskUpdates.EndTask(taskId); }
            throw;
        }
    }

    public async Task<ToolExecutionResult> ExecuteAsync(Guid taskId, string toolName, JsonElement arguments, CancellationToken cancellationToken)
    {
        TaskContext context = Get(taskId);
        audit.Write(new(DateTimeOffset.UtcNow, taskId, "tool.call", "attempted",
            new Dictionary<string, string> { ["tool"] = toolName }, SessionId: context.SessionId));
        try
        {
            return toolName switch
            {
                "file_search" => Search(arguments),
                "file_read" => Read(arguments),
                "file_edit" => Edit(context, arguments),
                "process_run" => await RunProcessAsync(taskId, context, arguments, cancellationToken).ConfigureAwait(false),
                "task_update" => _taskUpdates.Execute(taskId, arguments),
                "dir" => DirectoryOp(context, arguments),
                _ => throw new ToolValidationException($"Unknown internal tool '{toolName}'.")
            };
        }
        catch (SandboxPolicyException ex)
        {
            audit.Write(new(DateTimeOffset.UtcNow, taskId, "policy.deny", "denied",
                new Dictionary<string, string> { ["tool"] = toolName, ["message"] = ex.Message },
                SessionId: context.SessionId, ErrorCode: ex.Code));
            return new ToolExecutionResult(false, ex.Message, false, $"{toolName}:{ex.Code}:{ex.Message}", RetryAllowed: false,
                ErrorCode: ex.Code, ErrorCategory: ex.Category, TargetPath: ex.TargetPath,
                RequiredPermission: ex.RequiredPermission, Resolution: ex.Resolution);
        }
        catch (AdminOperationException ex)
        {
            audit.Write(new(DateTimeOffset.UtcNow, taskId, ex.Code == "sandbox.ipc_rejected" ? "sandbox.ipc.reject" : "sandbox.elevation.deny",
                "failed", new Dictionary<string, string> { ["message"] = ex.Message }, SessionId: context.SessionId, ErrorCode: ex.Code));
            return new ToolExecutionResult(false, ex.Message, false, $"{toolName}:{ex.Code}", RetryAllowed: false,
                ErrorCode: ex.Code, ErrorCategory: "platform", Resolution: "Approve the Windows elevation prompt and retry the Agent Task.");
        }
        catch (Exception ex) when (ex is ToolValidationException or JsonException or InvalidOperationException)
        {
            return new ToolExecutionResult(false, ex.Message, false,
                $"{toolName}:tool.input_invalid:{ex.Message}", RetryAllowed: true,
                ErrorCode: "tool.input_invalid", ErrorCategory: "tool",
                Resolution: "Correct the Tool arguments using the published schema and retry.");
        }
    }

    public async Task EndTaskAsync(Guid taskId, bool success, CancellationToken cancellationToken)
    {
        if (!_contexts.TryRemove(taskId, out TaskContext? context)) return;
        try
        {
            await CleanupAsync(taskId, context).ConfigureAwait(false);
            audit.Write(new(DateTimeOffset.UtcNow, taskId, "task.end", success ? "completed" : "stopped",
                new Dictionary<string, string>(), SessionId: context.SessionId));
        }
        catch
        {
            audit.Write(new(DateTimeOffset.UtcNow, taskId, "task.end", "cleanup_failed",
                new Dictionary<string, string>(), SessionId: context.SessionId));
            throw;
        }
        finally { _taskUpdates.EndTask(taskId); }
    }

    private ToolExecutionResult Search(JsonElement arguments)
    {
        string root = arguments.GetProperty("root").GetString()!;
        _paths.Demand(root, FileCapability.Read);
        string canonical = ResolvedPath.Resolve(root);
        string mode = arguments.GetProperty("mode").GetString()!;
        string query = arguments.GetProperty("query").GetString()!;
        int max = arguments.TryGetProperty("max_results", out JsonElement maxElement) ? maxElement.GetInt32() : 100;
        max = Math.Min(max, settings.SearchResultLimit);
        if (mode == "tail")
        {
            int lines = arguments.TryGetProperty("tail_lines", out JsonElement lineElement) ? lineElement.GetInt32() : 100;
            lines = Math.Min(lines, settings.TailLineLimit);
            if (new FileInfo(canonical).Length > settings.WholeReadLimitBytes * 16) throw new IOException("Tail input is too large.");
            if (!TextFileDetector.IsText(canonical)) throw new ToolValidationException("File appears to be binary; tail mode is text-only.");
            string[] content = File.ReadLines(canonical).TakeLast(lines).ToArray();
            return Success(JsonSerializer.Serialize(content));
        }
        IEnumerable<string> files = File.Exists(canonical) ? [canonical] : Directory.EnumerateFiles(canonical, "*", new EnumerationOptions
        {
            RecurseSubdirectories = true,
            IgnoreInaccessible = true,
            ReturnSpecialDirectories = false,
            AttributesToSkip = FileAttributes.ReparsePoint
        });
        IEnumerable<string> results = mode switch
        {
            "glob" => files.Where(path => FileSystemName.MatchesSimpleExpression(query, Path.GetRelativePath(canonical, path), ignoreCase: true)),
            "name" => files.Where(path => Path.GetFileName(path).Contains(query, StringComparison.OrdinalIgnoreCase)),
            "text" => files.Where(path => IsTextMatch(path, text => text.Contains(query, StringComparison.Ordinal))),
            "regex" => SearchRegex(files, query),
            _ => throw new ToolValidationException($"Unsupported search mode '{mode}'.")
        };
        return Success(JsonSerializer.Serialize(results.Take(max).Select(Path.GetFullPath).ToArray()));
    }

    private ToolExecutionResult Read(JsonElement arguments)
    {
        string path = arguments.GetProperty("path").GetString()!;
        _paths.Demand(path, FileCapability.Read);
        string canonical = ResolvedPath.Resolve(path);
        string mode = arguments.GetProperty("mode").GetString()!;
        FileInfo info = new(canonical);
        if (mode == "metadata") return Success(JsonSerializer.Serialize(new { path = canonical, info.Length, info.LastWriteTimeUtc, sha256 = EntryExecutablePolicy.ComputeSha256(canonical) }));
        if (mode == "whole")
        {
            if (info.Length > settings.WholeReadLimitBytes) throw new IOException("Whole read exceeds configured limit; use range mode.");
            return Success(File.ReadAllText(canonical));
        }
        long start = arguments.TryGetProperty("start", out JsonElement startElement) ? startElement.GetInt64() : 0;
        int length = arguments.GetProperty("length").GetInt32();
        if (length > settings.WholeReadLimitBytes) throw new IOException("Range length exceeds configured limit.");
        using FileStream stream = File.OpenRead(canonical);
        stream.Position = start;
        byte[] buffer = new byte[Math.Min(length, checked((int)Math.Max(0, stream.Length - start)))];
        _ = stream.Read(buffer);
        return Success(Encoding.UTF8.GetString(buffer));
    }

    private ToolExecutionResult Edit(TaskContext context, JsonElement arguments)
    {
        string path = arguments.GetProperty("path").GetString()!;
        string operation = arguments.GetProperty("operation").GetString()!;
        string? expected = arguments.TryGetProperty("expected_sha256", out JsonElement expectedElement) && expectedElement.ValueKind == JsonValueKind.String ? expectedElement.GetString() : null;
        var editor = new TransactionalFileEditor(context.Journal, settings.BackupDirectory, settings.MemoryBackupLimitBytes);
        if (operation == "delete")
        {
            _paths.Demand(path, FileCapability.Delete);
            FileDeleteTransaction transaction = editor.Delete(path, expected);
            transaction.Commit();
        }
        else if (operation == "rename")
        {
            string newPath = arguments.GetProperty("new_path").GetString()!;
            _paths.Demand(path, FileCapability.Write);
            _paths.Demand(newPath, FileCapability.Write);
            FileRenameTransaction transaction = editor.Rename(path, newPath);
            transaction.Commit();
        }
        else
        {
            _paths.Demand(path, FileCapability.Write);
            string replacement = operation == "patch"
                ? ApplyPatch(path, arguments.GetProperty("patch").GetString()!)
                : arguments.GetProperty("content").GetString()!;
            FileEditTransaction transaction = editor.Replace(path, replacement, expected);
            transaction.Commit();
        }
        string auditPath = operation == "rename"
            ? $"{Path.GetFullPath(path)} -> {Path.GetFullPath(arguments.GetProperty("new_path").GetString()!)}"
            : Path.GetFullPath(path);
        string eventName = operation == "delete" ? "file.delete" : "file.write";
        audit.Write(new(DateTimeOffset.UtcNow, context.Journal.TaskId, eventName, "completed",
            new Dictionary<string, string> { ["path"] = auditPath, ["operation"] = operation }, SessionId: context.SessionId));
        return Success($"{operation} completed for {auditPath}", externalEffect: true);
    }

    private ToolExecutionResult DirectoryOp(TaskContext context, JsonElement arguments)
    {
        string path = arguments.GetProperty("path").GetString()!;
        string operation = arguments.GetProperty("operation").GetString()!;
        var editor = new TransactionalFileEditor(context.Journal, settings.BackupDirectory, settings.MemoryBackupLimitBytes);
        if (operation == "mkdir")
        {
            _paths.Demand(path, FileCapability.Write);
            DirectoryMkdirTransaction transaction = editor.Mkdir(path);
            transaction.Commit();
            audit.Write(new(DateTimeOffset.UtcNow, context.Journal.TaskId, "dir.mkdir", "completed", new Dictionary<string, string> { ["path"] = Path.GetFullPath(path) }));
            return Success($"mkdir completed for {Path.GetFullPath(path)}", externalEffect: true);
        }
        if (operation == "rmdir")
        {
            _paths.Demand(path, FileCapability.Delete);
            DirectoryRmdirTransaction transaction = editor.Rmdir(path);
            transaction.Commit();
            audit.Write(new(DateTimeOffset.UtcNow, context.Journal.TaskId, "dir.rmdir", "completed", new Dictionary<string, string> { ["path"] = Path.GetFullPath(path) }));
            return Success($"rmdir completed for {Path.GetFullPath(path)}", externalEffect: true);
        }
        if (operation == "list")
        {
            _paths.Demand(path, FileCapability.Read);
            string canonical = ResolvedPath.Resolve(path);
            var entries = Directory.EnumerateFileSystemEntries(canonical, "*", SearchOption.TopDirectoryOnly)
                .Select(Path.GetFullPath)
                .ToArray();
            return Success(JsonSerializer.Serialize(entries));
        }
        throw new ToolValidationException($"Unsupported dir operation '{operation}'.");
    }

    private async Task<ToolExecutionResult> RunProcessAsync(Guid taskId, TaskContext context, JsonElement arguments, CancellationToken cancellationToken)
    {
        string executable = arguments.GetProperty("executable").GetString()!;
        string cwd = arguments.GetProperty("cwd").GetString()!;
        _paths.Demand(cwd, FileCapability.Read);
        string[] args = arguments.GetProperty("args").EnumerateArray().Select(static item => item.GetString()!).ToArray();
        _ = _entryExecutables.VerifyEntry(executable, cwd);
        string? managedEntryPoint = _entryExecutables.ManagedEntryPoint(executable, args, cwd);
        if (managedEntryPoint is not null)
        {
            _paths.Demand(managedEntryPoint, FileCapability.Execute);
            _entryExecutables.VerifyManagedEntryPoint(managedEntryPoint);
        }
        await EnsureProcessResourcesAsync(taskId, context, cancellationToken).ConfigureAwait(false);
        int seconds = ResolveProcessTimeout(arguments, settings.ProcessTimeoutSeconds);
        string? stdin = arguments.TryGetProperty("stdin", out JsonElement stdinElement) && stdinElement.ValueKind == JsonValueKind.String ? stdinElement.GetString() : null;
        RestrictedProcessResult result = new RestrictedProcessLauncher(_entryExecutables, audit, context.Profile.Sid, settings.JobLimits,
            allowNetwork: settings.NetworkAccess != "none")
            .Run(taskId, executable, args, cwd, TimeSpan.FromSeconds(seconds), stdin, context.SessionId, cancellationToken);
        string output = $"Process exited with code {result.ExitCode} (PID {result.ProcessId}).\nstdout:\n{result.StandardOutput}\nstderr:\n{result.StandardError}";
        return new(result.ExitCode == 0, output, true,
            result.ExitCode == 0 ? null : $"process:{executable}:{result.ExitCode}");
    }

    private async Task EnsureProcessResourcesAsync(Guid taskId, TaskContext context, CancellationToken cancellationToken)
    {
        if (context.ProcessResourcesApplied) return;
        await context.ResourceGate.WaitAsync(cancellationToken).ConfigureAwait(false);
        try
        {
            if (context.ProcessResourcesApplied) return;
            IReadOnlyList<NetworkEndpoint> allowances = settings.NetworkAccess != "none"
                ? await NetworkAllowanceSnapshot.ResolveAsync(settings.NetworkAllowances, cancellationToken).ConfigureAwait(false)
                : [];
            IReadOnlyList<(string Path, FileCapability Capabilities)> aclTargets = ProcessAclTargets();
            foreach ((string path, _) in aclTargets) TaskAclManager.DemandWriteDac(path);
            audit.Write(new(DateTimeOffset.UtcNow, taskId, "sandbox.elevation.request", "requested",
                new Dictionary<string, string> { ["reason"] = "network_task_resource" }, SessionId: context.SessionId));
            IElevatedNetworkSession? elevated = null;
            try
            {
                elevated = await _elevatedNetwork.StartAsync(taskId, cancellationToken).ConfigureAwait(false);
                audit.Write(new(DateTimeOffset.UtcNow, taskId, "sandbox.elevation.success", "connected",
                    new Dictionary<string, string> { ["protocol"] = AdminProtocol.Version.ToString() }, SessionId: context.SessionId));
                ApplyProcessAcls(taskId, context, aclTargets);
                context.Network = await TaskNetworkPolicy.ApplyAsync(elevated, context.Journal, context.Profile.Sid,
                    settings.NetworkAccess, allowances, audit, context.SessionId, cancellationToken).ConfigureAwait(false);
                elevated = null; // ownership transferred to TaskNetworkPolicy
                context.ProcessResourcesApplied = true;
            }
            finally
            {
                if (elevated is not null) await elevated.DisposeAsync().ConfigureAwait(false);
            }
        }
        finally { context.ResourceGate.Release(); }
    }

    private IReadOnlyList<(string Path, FileCapability Capabilities)> ProcessAclTargets()
    {
        var targets = new List<(string Path, FileCapability Capabilities)>();
        foreach ((string path, FileCapability capabilities, bool osRuntimeRoot) in settings.Files
            .GroupBy(static entry => ResolvedPath.Resolve(entry.Path, false), StringComparer.OrdinalIgnoreCase)
            .Select(static group => (group.Key,
                group.Aggregate(FileCapability.None, static (value, item) => value | item.Capabilities),
                group.All(static item => item.Origin == "os_runtime_root"))))
        {
            if (osRuntimeRoot) continue;
            // Entry files are identity gates, not ACL roots. Runtime/Workspace directory ACLs
            // provide inherited LPAC access where configured; OS/Program Files entries retain
            // their existing protected ACLs and are validated by the real LPAC launch test.
            if (File.Exists(path)) continue;
            targets.Add((path, capabilities));
        }
        return targets;
    }

    private static void ApplyProcessAcls(Guid taskId, TaskContext context,
        IReadOnlyList<(string Path, FileCapability Capabilities)> targets)
    {
        foreach ((string path, FileCapability capabilities) in targets)
            context.Acls.Add(context.Acl.Grant(taskId, path, capabilities));
    }

    internal static int ResolveProcessTimeout(JsonElement arguments, int configuredTimeoutSeconds)
    {
        int seconds = arguments.TryGetProperty("timeout_seconds", out JsonElement timeout) ? timeout.GetInt32() : configuredTimeoutSeconds;
        if (seconds <= 0) throw new ToolValidationException("timeout_seconds must be positive.");
        if (seconds > configuredTimeoutSeconds)
            throw new ToolValidationException($"timeout_seconds exceeds configured sandbox.resources.execution_timeout_seconds ({configuredTimeoutSeconds}).");
        return seconds;
    }

    private async Task CleanupAsync(Guid taskId, TaskContext context)
    {
        List<Exception> failures = [];
        if (context.Network is not null)
            try { await context.Network.DisposeAsync().ConfigureAwait(false); } catch (Exception ex) { failures.Add(ex); }
        foreach (AclSnapshot snapshot in context.Acls.AsEnumerable().Reverse()) try { context.Acl.Restore(taskId, snapshot); } catch (Exception ex) { failures.Add(ex); }
        if (failures.Count == 0) context.Journal.CompleteAndDelete();
        if (failures.Count > 0)
        {
            _cleanupFailed = true;
            throw new AggregateException("Task resource cleanup failed.", failures);
        }
    }

    private TaskContext Get(Guid taskId) => _contexts.TryGetValue(taskId, out TaskContext? context) ? context : throw new InvalidOperationException("Task Resource scope is not active.");
    private static ToolExecutionResult Success(string output, bool externalEffect = false) => new(true, output, externalEffect);
    private static bool IsTextMatch(string path, Func<string, bool> predicate)
    {
        try { return new FileInfo(path).Length <= 1024 * 1024 && predicate(File.ReadAllText(path)); }
        catch (Exception ex) when (ex is IOException or UnauthorizedAccessException or DecoderFallbackException) { return false; }
    }
    private static IEnumerable<string> SearchRegex(IEnumerable<string> files, string query)
    {
        var regex = new Regex(query, RegexOptions.CultureInvariant, TimeSpan.FromSeconds(1));
        return files.Where(path => IsTextMatch(path, regex.IsMatch));
    }
    private static string ApplyPatch(string path, string patch)
    {
        string current = File.Exists(path) ? File.ReadAllText(path) : string.Empty;
        JsonDocument document;
        try { document = JsonDocument.Parse(patch); }
        catch (JsonException) { throw new ToolValidationException("Patch must be a JSON object string containing old_text and new_text; unified diff is unsupported."); }
        using (document)
        {
            if (document.RootElement.ValueKind != JsonValueKind.Object ||
                !document.RootElement.TryGetProperty("old_text", out JsonElement oldElement) || oldElement.ValueKind != JsonValueKind.String ||
                !document.RootElement.TryGetProperty("new_text", out JsonElement newElement) || newElement.ValueKind != JsonValueKind.String)
                throw new ToolValidationException("Patch must be a JSON object string containing string old_text and new_text fields.");
            string oldText = oldElement.GetString()!;
            string newText = newElement.GetString()!;
            int first = current.IndexOf(oldText, StringComparison.Ordinal);
            if (first < 0 || current.IndexOf(oldText, first + oldText.Length, StringComparison.Ordinal) >= 0) throw new InvalidOperationException("Patch old_text must match exactly once.");
            return current[..first] + newText + current[(first + oldText.Length)..];
        }
    }

    private sealed class TaskContext(TaskResourceJournal journal, TaskAclManager acl, AppContainerProfile profile, Guid sessionId)
    {
        public TaskResourceJournal Journal { get; } = journal;
        public Guid SessionId { get; } = sessionId;
        public AppContainerProfile Profile { get; } = profile;
        public TaskAclManager Acl { get; } = acl;
        public List<AclSnapshot> Acls { get; } = [];
        public TaskNetworkPolicy? Network { get; set; }
        public SemaphoreSlim ResourceGate { get; } = new(1, 1);
        public bool ProcessResourcesApplied { get; set; }
    }
}