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.Windows.Network;
using Agent.Windows.Recovery;
using Agent.Windows.Security;

namespace Agent.Windows.Tools;

public sealed record WindowsSandboxSettings(
    IReadOnlyList<FilePolicyEntry> Files,
    IReadOnlyList<ExecutableApproval> Executables,
    bool AllowShell,
    SandboxJobLimits JobLimits,
    string NetworkMode,
    IReadOnlyList<(string Host, IReadOnlyList<int> Ports)> NetworkAllowances,
    string JournalDirectory,
    string BackupDirectory,
    long MemoryBackupLimitBytes,
    long WholeReadLimitBytes,
    int SearchResultLimit,
    int TailLineLimit,
    int ProcessTimeoutSeconds = 300);

public sealed class WindowsToolBackend(WindowsSandboxSettings settings, ISecurityAuditSink audit, ITaskUpdateHandler? taskUpdates = null) : IInternalToolBackend
{
    private readonly ConcurrentDictionary<Guid, TaskContext> _contexts = new();
    private readonly PathPolicy _paths = new(settings.Files);
    private readonly ExecutablePolicy _executables = new(settings.Executables, settings.AllowShell);
    private readonly ITaskUpdateHandler _taskUpdates = taskUpdates ?? new TaskGraphToolHandler();

    public async Task BeginTaskAsync(Guid taskId, Guid sessionId, CancellationToken cancellationToken)
    {
        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, sessionId);
        _taskUpdates.BeginTask(taskId, sessionId);
        try
        {
            foreach ((string path, FileCapability capabilities) in settings.Files
                .GroupBy(static entry => CanonicalPath.Resolve(entry.Path, false), StringComparer.OrdinalIgnoreCase)
                .Select(static group => (group.Key, group.Aggregate(FileCapability.None, static (value, item) => value | item.Capabilities))))
            {
                // OS/runtime executables may already grant AppContainer execute through their protected system ACL.
                // Do not rewrite a file ACL when execute is the only requested capability; CreateProcess remains
                // constrained by AppContainer access plus canonical path/SHA-256 approval.
                if (File.Exists(path) && capabilities == FileCapability.Execute) continue;
                // A single inheritable ACE on an approved directory covers existing and newly-created
                // descendants. Snapshotting and rewriting every child ACL is both redundant and
                // prohibitively expensive for runtime trees such as Python installations.
                context.Acls.Add(acl.Grant(taskId, path, capabilities));
            }
            IReadOnlyList<NetworkEndpoint> allowances = settings.NetworkMode == "allow"
                ? await NetworkAllowanceSnapshot.ResolveAsync(settings.NetworkAllowances, cancellationToken).ConfigureAwait(false)
                : [];
            context.Network = TaskNetworkPolicy.Apply(journal, profile.Sid, allowances, audit, sessionId);
            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 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" => Task.FromResult(Search(arguments)),
                "file_read" => Task.FromResult(Read(arguments)),
                "file_edit" => Task.FromResult(Edit(context, arguments)),
                "process_run" => Task.FromResult(RunProcess(context, arguments)),
                "task_update" => Task.FromResult(_taskUpdates.Execute(taskId, arguments)),
                "dir" => Task.FromResult(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 Task.FromResult(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));
        }
    }

    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 = CanonicalPath.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 = CanonicalPath.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 = ExecutablePolicy.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 = CanonicalPath.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 ToolExecutionResult RunProcess(TaskContext context, JsonElement arguments)
    {
        Guid taskId = context.Journal.TaskId;
        string executable = arguments.GetProperty("executable").GetString()!;
        string cwd = arguments.GetProperty("cwd").GetString()!;
        _paths.Demand(executable, FileCapability.Execute);
        _paths.Demand(cwd, FileCapability.Read);
        string[] args = arguments.GetProperty("args").EnumerateArray().Select(static item => item.GetString()!).ToArray();
        string? managedEntryPoint = ExecutablePolicy.ManagedEntryPoint(executable, args);
        if (managedEntryPoint is not null) _paths.Demand(managedEntryPoint, FileCapability.Execute);
        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(_executables, audit, AppContainerProfile.EnsureFixed().Sid, settings.JobLimits)
            .Run(taskId, executable, args, cwd, TimeSpan.FromSeconds(seconds), stdin, context.SessionId);
        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}");
    }

    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.process_timeout_seconds ({configuredTimeoutSeconds}).");
        return seconds;
    }

    private async Task CleanupAsync(Guid taskId, TaskContext context)
    {
        List<Exception> failures = [];
        try { context.Network?.Dispose(); } 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();
        await Task.CompletedTask;
        if (failures.Count > 0) 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;
        using JsonDocument document = JsonDocument.Parse(patch);
        string oldText = document.RootElement.GetProperty("old_text").GetString()!;
        string newText = document.RootElement.GetProperty("new_text").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, Guid sessionId)
    {
        public TaskResourceJournal Journal { get; } = journal;
        public Guid SessionId { get; } = sessionId;
        public TaskAclManager Acl { get; } = acl;
        public List<AclSnapshot> Acls { get; } = [];
        public TaskNetworkPolicy? Network { get; set; }
    }
}