IT-Nexus/agent-cs/Services/RtcService.cs

using System.Diagnostics;
using System.Net;
using System.Net.Sockets;
using System.Net.WebSockets;
using System.Security.Cryptography;
using System.Text;
using Newtonsoft.Json;
using Newtonsoft.Json.Linq;

namespace ITNexusAgent.Services;

// WebSocket-Brücke: empfängt rdp_start vom Browser, spawnt Screen-Capture-Helper
// in der User-Session (via schtasks), leitet JPEG-Frames als rdp_frame weiter
public class RtcService
{
    private readonly string _serverUrl;
    private readonly string _agentKey;
    private readonly string _hostname;
    private int _indicatorPid = -1;
    private TcpListener? _indicatorListener;
    private CancellationTokenSource? _userDisconnectCts;

    public RtcService(string serverUrl, string agentKey, string hostname)
    {
        _serverUrl = serverUrl;
        _agentKey = agentKey;
        _hostname = hostname;
    }

    public async Task RunAsync(CancellationToken ct)
    {
        while (!ct.IsCancellationRequested)
        {
            try { await ConnectAsync(ct); }
            catch (Exception ex) { AgentWorker.Log($"RDP: {ex.Message}"); }

            if (!ct.IsCancellationRequested)
                await Task.Delay(TimeSpan.FromSeconds(15), ct).ContinueWith(_ => { });
        }
    }

    private async Task ConnectAsync(CancellationToken ct)
    {
        var wsUrl = _serverUrl
            .Replace("https://", "wss://")
            .Replace("http://", "ws://")
            + $"/ws?type=rdp-agent&key={Uri.EscapeDataString(_agentKey)}&hostname={Uri.EscapeDataString(_hostname)}";

        using var ws = new ClientWebSocket();
        await ws.ConnectAsync(new Uri(wsUrl), ct);
        AgentWorker.Log("RDP: Bereit");

        CancellationTokenSource? captureCts = null;

        var buf = new byte[4096];
        while (ws.State == WebSocketState.Open && !ct.IsCancellationRequested)
        {
            WebSocketReceiveResult res;
            try { res = await ws.ReceiveAsync(new ArraySegment<byte>(buf), ct); }
            catch { break; }

            if (res.MessageType == WebSocketMessageType.Close) break;

            var raw = Encoding.UTF8.GetString(buf, 0, res.Count);
            JObject? obj;
            try { obj = JObject.Parse(raw); } catch { continue; }

            var type = obj["type"]?.ToString();

            if (type == "rdp_start" && captureCts == null)
            {
                var screenIdx = obj["screen"]?.ToObject<int>() ?? 0;
                captureCts = CancellationTokenSource.CreateLinkedTokenSource(ct);
                _ = ConsentAndCaptureAsync(ws, screenIdx, captureCts.Token);
                AgentWorker.Log($"RDP: Consent angefordert (screen={screenIdx})");
            }
            else if (type == "rdp_switch_screen" && captureCts != null)
            {
                // Consent bereits erteilt — nur Screen wechseln, kein erneuter Dialog
                var newScreen = obj["screen"]?.ToObject<int>() ?? 0;
                captureCts.Cancel();
                captureCts = CancellationTokenSource.CreateLinkedTokenSource(ct);
                var switchToken = _userDisconnectCts != null
                    ? CancellationTokenSource.CreateLinkedTokenSource(captureCts.Token, _userDisconnectCts.Token).Token
                    : captureCts.Token;
                _ = CapturePipeLoopAsync(ws, newScreen, switchToken);
                AgentWorker.Log($"RDP: Screen gewechselt zu {newScreen}");
            }
            else if (type == "rdp_stop" && captureCts != null)
            {
                captureCts.Cancel();
                captureCts = null;
                StopIndicator();
                AgentWorker.Log("RDP: Screen-Capture gestoppt");
            }
        }

        captureCts?.Cancel();
        StopIndicator();
        AgentWorker.Log("RDP: Getrennt");
    }

    private void StopIndicator()
    {
        _userDisconnectCts?.Cancel();
        _userDisconnectCts = null;
        try { _indicatorListener?.Stop(); } catch { }
        _indicatorListener = null;
        if (_indicatorPid > 0)
        {
            try { Process.GetProcessById(_indicatorPid).Kill(); } catch { }
            _indicatorPid = -1;
        }
    }

    private static string GenerateSecret() => Convert.ToHexString(RandomNumberGenerator.GetBytes(16));

    // Liest exakt secret.Length Bytes und vergleicht zeitkonstant — verhindert, dass ein beliebiger
    // lokaler Prozess sich als der gespawnte Helper ausgibt und Consent/Disconnect/Frames vortäuscht.
    private static async Task<bool> ValidateSecretAsync(NetworkStream stream, string secret, CancellationToken ct)
    {
        var expected = Encoding.ASCII.GetBytes(secret);
        var buf = new byte[expected.Length];
        var read = 0;
        while (read < buf.Length)
        {
            int n;
            try { n = await stream.ReadAsync(buf.AsMemory(read, buf.Length - read), ct); }
            catch { return false; }
            if (n == 0) return false;
            read += n;
        }
        return CryptographicOperations.FixedTimeEquals(buf, expected);
    }

    private async Task ConsentAndCaptureAsync(ClientWebSocket ws, int screenIdx, CancellationToken ct)
    {
        var exePath = System.Diagnostics.Process.GetCurrentProcess().MainModule!.FileName;
        var secret = GenerateSecret();

        var consentListener = new TcpListener(IPAddress.Loopback, 0);
        consentListener.Start();
        var consentPort = ((IPEndPoint)consentListener.LocalEndpoint).Port;

        if (!SpawnConsentHelper(exePath, consentPort.ToString(), secret))
        {
            consentListener.Stop();
            AgentWorker.Log("RDP: Consent-Helper fehlgeschlagen, verweigere Zugriff");
            var failed = Encoding.UTF8.GetBytes("{\"type\":\"rdp_denied\",\"reason\":\"consent_spawn_failed\"}");
            if (ws.State == WebSocketState.Open)
                await ws.SendAsync(new ArraySegment<byte>(failed), WebSocketMessageType.Text, true, CancellationToken.None);
            return;
        }

        // Browser informieren dass auf User-Zustimmung gewartet wird
        var pending = Encoding.UTF8.GetBytes("{\"type\":\"rdp_consent_pending\"}");
        if (ws.State == WebSocketState.Open)
            await ws.SendAsync(new ArraySegment<byte>(pending), WebSocketMessageType.Text, true, CancellationToken.None);

        bool accepted = false;
        try
        {
            using var timeoutCts = new CancellationTokenSource(TimeSpan.FromSeconds(35));
            using var linked = CancellationTokenSource.CreateLinkedTokenSource(ct, timeoutCts.Token);
            while (!linked.IsCancellationRequested)
            {
                using var tcp = await consentListener.AcceptTcpClientAsync(linked.Token);
                var stream = tcp.GetStream();
                if (!await ValidateSecretAsync(stream, secret, linked.Token)) continue; // fremder Connect-Versuch ohne gültiges Secret — ignorieren, weiter warten
                var b = stream.ReadByte();
                accepted = b == 1;
                break;
            }
        }
        catch { accepted = false; }
        finally { consentListener.Stop(); }

        if (!accepted)
        {
            AgentWorker.Log("RDP: User hat Zugriff abgelehnt");
            var denied = Encoding.UTF8.GetBytes("{\"type\":\"rdp_denied\"}");
            if (ws.State == WebSocketState.Open)
                await ws.SendAsync(new ArraySegment<byte>(denied), WebSocketMessageType.Text, true, CancellationToken.None);
            return;
        }

        AgentWorker.Log("RDP: User hat Zugriff erlaubt, starte Capture");
        SpawnIndicator(exePath);

        var combinedToken = _userDisconnectCts != null
            ? CancellationTokenSource.CreateLinkedTokenSource(ct, _userDisconnectCts.Token).Token
            : ct;
        await CapturePipeLoopAsync(ws, screenIdx, combinedToken);
    }

    // Zeigt dauerhaftes "Bildschirm wird übertragen"-Overlay (TeamViewer-Style) solange Capture läuft.
    // User kann per Klick selbst trennen — Signal landet hier als CancellationTokenSource.Cancel().
    private void SpawnIndicator(string exePath)
    {
        try
        {
            _indicatorListener = new TcpListener(IPAddress.Loopback, 0);
            _indicatorListener.Start();
            var port = ((IPEndPoint)_indicatorListener.LocalEndpoint).Port;
            var secret = GenerateSecret();

            _userDisconnectCts = new CancellationTokenSource();
            var listener = _indicatorListener;
            var disconnectCts = _userDisconnectCts;
            _ = Task.Run(async () =>
            {
                try
                {
                    while (true)
                    {
                        using var tcp = await listener.AcceptTcpClientAsync();
                        var stream = tcp.GetStream();
                        if (!await ValidateSecretAsync(stream, secret, CancellationToken.None)) continue; // fremder Connect-Versuch ohne gültiges Secret
                        stream.ReadByte();
                        disconnectCts.Cancel();
                        AgentWorker.Log("RDP: User hat über Overlay getrennt");
                        break;
                    }
                }
                catch { }
            });

            _indicatorPid = SessionSpawner.SpawnInUserSession(exePath, $"--rdp-indicator {port} {secret}");
        }
        catch (Exception ex)
        {
            AgentWorker.Log($"RDP: Indicator-Start fehlgeschlagen: {ex.Message}");
        }
    }

    private static bool SpawnConsentHelper(string exePath, string portStr, string secret)
    {
        var pid = SessionSpawner.SpawnInUserSession(exePath, $"--rdp-consent {portStr} {secret}");
        return pid > 0;
    }

    private async Task CapturePipeLoopAsync(ClientWebSocket ws, int screenIdx, CancellationToken ct)
    {
        var exePath = System.Diagnostics.Process.GetCurrentProcess().MainModule!.FileName;
        var secret = GenerateSecret();

        // TCP Loopback: kein ACL-Problem zwischen SYSTEM-Service und User-Prozess
        var listener = new TcpListener(IPAddress.Loopback, 0);
        listener.Start();
        var port = ((IPEndPoint)listener.LocalEndpoint).Port;

        if (!SpawnCaptureHelper(exePath, port.ToString(), screenIdx, secret))
        {
            listener.Stop();
            AgentWorker.Log("RDP: Helper-Start fehlgeschlagen");
            return;
        }

        TcpClient? tcp = null;
        try
        {
            using var connectCts = new CancellationTokenSource(TimeSpan.FromSeconds(15));
            using var linked = CancellationTokenSource.CreateLinkedTokenSource(ct, connectCts.Token);
            while (true)
            {
                var candidate = await listener.AcceptTcpClientAsync(linked.Token);
                if (await ValidateSecretAsync(candidate.GetStream(), secret, linked.Token))
                {
                    tcp = candidate;
                    break;
                }
                candidate.Dispose(); // fremder Connect-Versuch ohne gültiges Secret — verwerfen, weiter warten
            }
        }
        catch
        {
            listener.Stop();
            AgentWorker.Log("RDP: Helper hat sich nicht verbunden");
            return;
        }
        finally { listener.Stop(); }

        AgentWorker.Log($"RDP: Helper verbunden (Port {port}), sende Frames...");

        using (tcp)
        {
            var stream = tcp.GetStream();
            var lenBuf = new byte[4];

            while (!ct.IsCancellationRequested && tcp.Connected && ws.State == WebSocketState.Open)
            {
                try
                {
                    var read = 0;
                    while (read < 4)
                    {
                        var n = await stream.ReadAsync(lenBuf.AsMemory(read, 4 - read), ct);
                        if (n == 0) goto done;
                        read += n;
                    }
                    var jpegLen = BitConverter.ToInt32(lenBuf);
                    if (jpegLen <= 0 || jpegLen > 5_000_000) continue;

                    var jpeg = new byte[jpegLen];
                    read = 0;
                    while (read < jpegLen)
                    {
                        var n = await stream.ReadAsync(jpeg.AsMemory(read, jpegLen - read), ct);
                        if (n == 0) goto done;
                        read += n;
                    }

                    var payload = Encoding.UTF8.GetBytes(JsonConvert.SerializeObject(new
                    {
                        type = "rdp_frame",
                        data = Convert.ToBase64String(jpeg)
                    }));

                    if (ws.State == WebSocketState.Open)
                        await ws.SendAsync(new ArraySegment<byte>(payload), WebSocketMessageType.Text, true, CancellationToken.None);
                }
                catch (OperationCanceledException) { break; }
                catch (Exception ex) { AgentWorker.Log($"RDP: TCP-Fehler: {ex.Message}"); break; }
            }
        }

        done:
        AgentWorker.Log("RDP: Frame-Loop beendet");
    }

    private static bool SpawnCaptureHelper(string exePath, string portStr, int screenIdx, string secret)
    {
        var pid = SessionSpawner.SpawnInUserSession(exePath, $"--rdp-capture {portStr} {screenIdx} {secret}");
        return pid > 0;
    }
}