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);

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, 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);
        _taskUpdates.BeginTask(taskId);
        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;
                context.Roots.Add(path);
                foreach (string target in ExistingAclTargets(path)) context.Acls.Add(acl.Grant(taskId, target, capabilities));
            }
            IReadOnlyList<NetworkEndpoint> allowances = settings.NetworkMode == "allow"
                ? await NetworkAllowanceSnapshot.ResolveAsync(settings.NetworkAllowances, cancellationToken).ConfigureAwait(false)
                : [];
            context.Network = TaskNetworkPolicy.Apply(journal, profile.Sid, allowances, audit);
            if (!_contexts.TryAdd(taskId, context)) throw new InvalidOperationException("Task context already exists.");
        }
        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);
        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(taskId, arguments)),
            "task_update" => Task.FromResult(_taskUpdates.Execute(taskId, arguments)),
            _ => throw new ToolValidationException($"Unknown internal tool '{toolName}'.")
        };
    }

    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); }
        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.");
            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
        {
            _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();
        }
        audit.Write(new(DateTimeOffset.UtcNow, context.Journal.TaskId, "file." + operation, "completed", new Dictionary<string, string> { ["path"] = Path.GetFullPath(path) }));
        return Success($"{operation} completed for {Path.GetFullPath(path)}", externalEffect: true);
    }

    private ToolExecutionResult RunProcess(Guid taskId, JsonElement arguments)
    {
        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 = arguments.TryGetProperty("timeout_seconds", out JsonElement timeout) ? timeout.GetInt32() : 300;
        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);
        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 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); }
        var existingSnapshots = context.Acls.Select(static snapshot => snapshot.Path).ToHashSet(StringComparer.OrdinalIgnoreCase);
        foreach (string entry in context.Roots.SelectMany(ExistingAclTargets).Distinct(StringComparer.OrdinalIgnoreCase)
            .Where(entry => !existingSnapshots.Contains(entry))
            .OrderBy(static entry => entry.Count(character => character == Path.DirectorySeparatorChar)))
            try { context.Acl.NormalizeNewEntry(taskId, entry); } 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 IEnumerable<string> ExistingAclTargets(string path)
    {
        if (File.Exists(path))
        {
            yield return path;
            yield break;
        }
        if (!Directory.Exists(path)) yield break;
        var options = new EnumerationOptions
        {
            RecurseSubdirectories = true,
            IgnoreInaccessible = false,
            ReturnSpecialDirectories = false,
            AttributesToSkip = FileAttributes.ReparsePoint
        };
        foreach (string entry in Directory.EnumerateFileSystemEntries(path, "*", options)
            .OrderByDescending(static entry => entry.Count(character => character == Path.DirectorySeparatorChar)))
            yield return entry;
        yield return path;
    }
    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)
    {
        public TaskResourceJournal Journal { get; } = journal;
        public TaskAclManager Acl { get; } = acl;
        public List<AclSnapshot> Acls { get; } = [];
        public HashSet<string> Roots { get; } = new(StringComparer.OrdinalIgnoreCase);
        public TaskNetworkPolicy? Network { get; set; }
    }
}