gott/Endpoints/RemoteEndpoint.cs

using System;
using System.Buffers;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Net.WebSockets;
using System.Text;
using System.Text.Json;
using System.Threading;
using System.Threading.Tasks;
using Gott.Services;
using Microsoft.AspNetCore.Builder;
using Microsoft.AspNetCore.Http;
using Microsoft.AspNetCore.Routing;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Hosting;

namespace Gott.Endpoints;

public static class RemoteEndpoint
{
    public static void MapRemoteEndpoints(this IEndpointRouteBuilder app)
    {
        // WebSocket-Endpoint für Remote-Clients (kein Auth)
        app.Map("/api/remote", HandleWebSocket);

        // REST-APIs für Frontend (Auth required)
        app.MapGet("/api/remote/connections", GetConnections);
        app.MapPost("/api/remote/assign", AssignConnection);
        app.MapPost("/api/remote/unassign", UnassignConnection);

        // PowerShell-Script ausliefern
        app.MapGet("/remote.ps1", (Delegate)ServeScript);
        app.MapGet("/remote", (Delegate)ServeScript);
    }

    private static async Task HandleWebSocket(HttpContext ctx, RemoteConnectionManager mgr)
    {
        if (!ctx.WebSockets.IsWebSocketRequest)
        {
            ctx.Response.StatusCode = 400;
            await ctx.Response.WriteAsync("WebSocket expected");
            return;
        }

        // App-Shutdown-Token verwenden damit WebSocket-Loops sauber beendet werden
        var appLifetime = ctx.RequestServices.GetRequiredService<IHostApplicationLifetime>();
        var ct = appLifetime.ApplicationStopping;

        using var ws = await ctx.WebSockets.AcceptWebSocketAsync();
        var conn = mgr.Register(ws);

        using var heartbeatCts = CancellationTokenSource.CreateLinkedTokenSource(ct);

        try
        {
            var receiveTask = ReceiveLoop(ws, conn, mgr, heartbeatCts.Token);
            var heartbeatTask = HeartbeatLoop(ws, conn, heartbeatCts);

            // Warten bis einer von beiden endet (Receive-Fehler oder Heartbeat-Timeout)
            await Task.WhenAny(receiveTask, heartbeatTask);
            heartbeatCts.Cancel();

            // Beide Tasks sauber abwarten
            try { await receiveTask; } catch { }
            try { await heartbeatTask; } catch { }
        }
        catch (WebSocketException)
        {
            // Client disconnected
        }
        catch (OperationCanceledException)
        {
            // Shutdown oder Ctrl+C
        }
        finally
        {
            mgr.Unregister(conn.Id);
            Console.WriteLine($"[Remote] Client disconnected: {conn.Id} ({conn.Hostname})");

            if (ws.State == WebSocketState.Open || ws.State == WebSocketState.CloseReceived)
            {
                try
                {
                    using var closeCts = new CancellationTokenSource(TimeSpan.FromSeconds(3));
                    await ws.CloseAsync(WebSocketCloseStatus.NormalClosure, "bye", closeCts.Token);
                }
                catch { }
            }
        }
    }

    private static async Task ReceiveLoop(WebSocket ws, RemoteConnection conn, RemoteConnectionManager mgr, CancellationToken ct)
    {
        var buffer = new byte[64 * 1024]; // 64 KB Buffer
        var messageBuffer = new MemoryStream();

        while (ws.State == WebSocketState.Open && !ct.IsCancellationRequested)
        {
            WebSocketReceiveResult result;
            messageBuffer.SetLength(0);

            do
            {
                result = await ws.ReceiveAsync(buffer, ct);
                if (result.MessageType == WebSocketMessageType.Close) return;
                messageBuffer.Write(buffer, 0, result.Count);
            } while (!result.EndOfMessage);

            if (result.MessageType != WebSocketMessageType.Text) continue;

            var json = Encoding.UTF8.GetString(messageBuffer.GetBuffer(), 0, (int)messageBuffer.Length);

            JsonElement msg;
            try { msg = JsonSerializer.Deserialize<JsonElement>(json); }
            catch { continue; }

            conn.LastActivityAt = DateTime.UtcNow;

            var type = msg.TryGetProperty("type", out var tp) ? tp.GetString() : "";

            switch (type)
            {
                case "hello":
                    conn.Hostname = msg.TryGetProperty("hostname", out var h) ? h.GetString() ?? "" : "";
                    conn.Username = msg.TryGetProperty("username", out var u) ? u.GetString() ?? "" : "";
                    conn.OsVersion = msg.TryGetProperty("os", out var o) ? o.GetString() ?? "" : "";
                    Console.WriteLine($"[Remote] Client connected: {conn.Username}@{conn.Hostname} ({conn.OsVersion})");
                    break;

                case "result":
                    var cmdResult = new RemoteCommandResult
                    {
                        ExitCode = msg.TryGetProperty("exitCode", out var ec) ? ec.GetInt32() : -1,
                        Output = msg.TryGetProperty("output", out var ot) ? ot.GetString() ?? "" : "",
                        Error = msg.TryGetProperty("error", out var er) ? er.GetString() ?? "" : ""
                    };
                    mgr.HandleResult(conn.Id, cmdResult);
                    break;

                case "ping":
                    var pong = Encoding.UTF8.GetBytes("{\"type\":\"pong\"}");
                    await ws.SendAsync(pong, WebSocketMessageType.Text, true, ct);
                    break;

                case "pong":
                    // Antwort auf Server-Ping → LastActivityAt ist schon gesetzt
                    break;
            }
        }
    }

    private static async Task HeartbeatLoop(WebSocket ws, RemoteConnection conn, CancellationTokenSource heartbeatCts)
    {
        const int pingIntervalSeconds = 30;
        const int deadTimeoutSeconds = 90;
        var pingBytes = Encoding.UTF8.GetBytes("{\"type\":\"ping\"}");

        try
        {
            while (!heartbeatCts.IsCancellationRequested && ws.State == WebSocketState.Open)
            {
                await Task.Delay(TimeSpan.FromSeconds(pingIntervalSeconds), heartbeatCts.Token);

                // Prüfe ob der Client noch lebt
                var silenceSec = (DateTime.UtcNow - conn.LastActivityAt).TotalSeconds;
                if (silenceSec > deadTimeoutSeconds)
                {
                    Console.WriteLine($"[Remote] Heartbeat timeout ({silenceSec:F0}s): {conn.Id} ({conn.Hostname})");
                    heartbeatCts.Cancel();
                    return;
                }

                // Ping senden
                try
                {
                    await ws.SendAsync(pingBytes, WebSocketMessageType.Text, true, heartbeatCts.Token);
                }
                catch
                {
                    // Send fehlgeschlagen → Verbindung tot
                    heartbeatCts.Cancel();
                    return;
                }
            }
        }
        catch (OperationCanceledException)
        {
            // Normal – ReceiveLoop hat beendet
        }
    }

    private static IResult GetConnections(RemoteConnectionManager mgr)
    {
        var connections = mgr.GetAll().Select(c => new
        {
            id = c.Id,
            hostname = c.Hostname,
            username = c.Username,
            os = c.OsVersion,
            connectedAt = c.ConnectedAt.ToString("o"),
            assignedSessionId = c.AssignedSessionId
        });
        return Results.Ok(connections);
    }

    private static IResult AssignConnection(HttpContext ctx, RemoteConnectionManager mgr)
    {
        using var reader = new StreamReader(ctx.Request.Body);
        var body = reader.ReadToEndAsync().Result;
        var json = JsonSerializer.Deserialize<JsonElement>(body);

        var connectionId = json.TryGetProperty("connectionId", out var cid) ? cid.GetString() : "";
        var sessionId = json.TryGetProperty("sessionId", out var sid) ? sid.GetString() : "";

        if (string.IsNullOrEmpty(connectionId) || string.IsNullOrEmpty(sessionId))
            return Results.BadRequest(new { error = "connectionId und sessionId erforderlich" });

        if (!mgr.Assign(connectionId, sessionId))
            return Results.NotFound(new { error = "Verbindung nicht gefunden" });

        return Results.Ok(new { ok = true });
    }

    private static IResult UnassignConnection(HttpContext ctx, RemoteConnectionManager mgr)
    {
        using var reader = new StreamReader(ctx.Request.Body);
        var body = reader.ReadToEndAsync().Result;
        var json = JsonSerializer.Deserialize<JsonElement>(body);

        var connectionId = json.TryGetProperty("connectionId", out var cid) ? cid.GetString() : "";

        if (string.IsNullOrEmpty(connectionId))
            return Results.BadRequest(new { error = "connectionId erforderlich" });

        if (!mgr.Unassign(connectionId))
            return Results.NotFound(new { error = "Verbindung nicht gefunden" });

        return Results.Ok(new { ok = true });
    }

    private static async Task<IResult> ServeScript(HttpContext ctx)
    {
        var scriptPath = Path.Combine(AppContext.BaseDirectory, "wwwroot", "remote.ps1");
        if (!File.Exists(scriptPath))
            return Results.NotFound("Script nicht gefunden");

        // Host aus X-Forwarded-Host (Reverse-Proxy) oder Request.Host
        var host = ctx.Request.Headers["X-Forwarded-Host"].FirstOrDefault()
                ?? ctx.Request.Host.ToString();

        // Schema aus X-Forwarded-Proto oder Request.Scheme
        var httpScheme = ctx.Request.Headers["X-Forwarded-Proto"].FirstOrDefault()
                      ?? ctx.Request.Scheme;
        var wsScheme = httpScheme == "https" ? "wss" : "ws";

        var wsUrl = $"{wsScheme}://{host}/api/remote";

        var script = await File.ReadAllTextAsync(scriptPath);
        script = script.Replace("{{WS_URL}}", wsUrl);

        return Results.Content(script, "text/plain; charset=utf-8");
    }
}