Relay/RelayServer/Services/ChatTest.cs

using System.Text.Json;
using RelayServer.Models;
using WebSocketSharp;
using WebSocketSharp.Server;

namespace RelayServer.Services;

public class ChatTest : WebSocketBehavior
{
    public static ClientKeyService? ClientKeyService { get; set; }
    public static string? ServerPublicKey { get; set; }
    public static string? ServerPrivateKey { get; set; }
    public static string? ChannelDbKey { get; set; }
    public static SurrealDb.Net.SurrealDbClient? Db { get; set; }

    protected override void OnMessage(MessageEventArgs e)
    {
        var msg = e.Data;
        Console.WriteLine(msg);

        if (msg.StartsWith("REGISTER_KEY|"))
        {
            HandleRegisterKey(msg);
            return;
        }

        if (msg == "GET_SERVER_KEY")
        {
            HandleGetServerKey();
            return;
        }

        if (msg == "GET_CHANNELS")
        {
            HandleGetChannels();
            return;
        }

        if (msg.StartsWith("GET_HISTORY|"))
        {
            HandleGetHistory(msg);
            return;
        }
        
        SocketRtcSignalMessage? rtcProbe = null;
        try
        {
            rtcProbe = JsonSerializer.Deserialize<SocketRtcSignalMessage>(msg);
        }
        catch
        {
            // ignored
        }

        if (rtcProbe?.Type == "encrypted_rtc_signal")
        {
            HandleEncryptedRtcSignal(msg);
            return;
        }

        HandleEncryptedClientMessage(msg);
    }
    
    private static string ExtractUsernameFromUserId(string senderUserId)
    {
        if (string.IsNullOrWhiteSpace(senderUserId))
            return "Unknown";

        var parts = senderUserId.Split(':', 2);
        return parts.Length == 2 ? parts[1] : senderUserId;
    }

    private void HandleRegisterKey(string msg)
    {
        var parts = msg.Split('|', 3);

        if (parts.Length < 3)
        {
            Console.WriteLine("Invalid REGISTER_KEY payload.");
            return;
        }

        var username = parts[1];
        var publicKey = parts[2];

        if (ClientKeyService is null)
        {
            Console.WriteLine("ClientKeyService is not initialized.");
            return;
        }

        Task.Run(async () =>
        {
            await ClientKeyService.RegisterOrUpdateKeyAsync(username, publicKey);
        }).GetAwaiter().GetResult();

        Send($"SERVER:REGISTERED_KEY:{username}");
    }

    private void HandleGetChannels()
    {
        if (Db is null)
        {
            Console.WriteLine("Db is not initialized.");
            return;
        }

        var channels = Task.Run(async () => await Db.Select<Channels>("channels"))
            .GetAwaiter()
            .GetResult()
            .OrderBy(c => c.CreatedAt)
            .Select(c => new SocketChannelInfo
            {
                ChannelId = GetRecordId(c.Id),
                Name = c.Name,
                CreatedAt = c.CreatedAt
            })
            .ToList();

        var payload = new SocketChannelList
        {
            Type = "channel_list",
            Channels = channels
        };

        Send(JsonSerializer.Serialize(payload));
    }
    
    private void HandleGetServerKey()
    {
        if (string.IsNullOrWhiteSpace(ServerPublicKey))
        {
            Console.WriteLine("Server public key is not initialized.");
            return;
        }

        var payload = new ServerPublicKeyMessage
        {
            Type = "server_public_key",
            PublicKey = ServerPublicKey
        };

        Send(JsonSerializer.Serialize(payload));
    }

    private void HandleEncryptedClientMessage(string msg)
    {
        SocketEncryptedMessage? clientPayload;

        try
        {
            clientPayload = JsonSerializer.Deserialize<SocketEncryptedMessage>(msg);
        }
        catch
        {
            Console.WriteLine("Failed to parse encrypted client payload.");
            return;
        }

        if (clientPayload is null || clientPayload.Type != "client_encrypted_chat")
            return;

        if (ClientKeyService is null ||
            Db is null ||
            string.IsNullOrWhiteSpace(ServerPrivateKey) ||
            string.IsNullOrWhiteSpace(ChannelDbKey))
        {
            Console.WriteLine("Server crypto/database dependencies are not initialized.");
            return;
        }

        string plainText;

        try
        {
            plainText = E2EeHelper.DecryptForRecipient(
                new EncryptedPayload
                {
                    CipherText = clientPayload.CipherText,
                    Nonce = clientPayload.Nonce,
                    Tag = clientPayload.Tag,
                    EncryptedKey = clientPayload.EncryptedKey
                },
                ServerPrivateKey
            );
        }
        catch (Exception ex)
        {
            Console.WriteLine($"Failed to decrypt client payload: {ex.Message}");
            return;
        }

        Console.WriteLine($"Server decrypted message from {clientPayload.SenderUsername}: {plainText}");
        try
        {
            var channelCrypto = new ChannelCryptoService();
            var dbEncrypted = channelCrypto.Encrypt(plainText, ChannelDbKey);

            var savedMessage = Task.Run(async () =>
                await Db.Create("channel_messages", new ChannelMessages
                {
                    ChannelId = clientPayload.ChannelId,
                    SenderUserId = $"users:{clientPayload.SenderUsername.ToLower()}",
                    CipherText = dbEncrypted.cipherText,
                    Nonce = dbEncrypted.nonce,
                    Tag = dbEncrypted.tag,
                    CreatedAt = DateTime.UtcNow
                })
            ).GetAwaiter().GetResult();

            Console.WriteLine($"Live message saved to DB: {JsonSerializer.Serialize(savedMessage)}");
        }
        catch (Exception ex)
        {
            Console.WriteLine($"Failed to save live message to DB: {ex.Message}");
            return;
        }

        var allKeys = Task.Run(async () => await ClientKeyService.GetAllAsync())
            .GetAwaiter()
            .GetResult();

        foreach (var client in allKeys)
        {
            var encrypted = E2EeHelper.EncryptForRecipient(plainText, client.PublicKey);

            Console.WriteLine($"Encrypting outbound message from {clientPayload.SenderUsername} for {client.Username}");

            var outbound = new SocketEncryptedMessage
            {
                Type = "encrypted_chat",
                SenderUsername = clientPayload.SenderUsername,
                RecipientUsername = client.Username,
                ChannelId = clientPayload.ChannelId,
                CipherText = encrypted.CipherText,
                Nonce = encrypted.Nonce,
                Tag = encrypted.Tag,
                EncryptedKey = encrypted.EncryptedKey
            };

            Sessions.Broadcast(JsonSerializer.Serialize(outbound));
        }
    }

    private void HandleGetHistory(string msg)
    {
        var parts = msg.Split('|', 3);

        if (parts.Length < 3)
        {
            Console.WriteLine("Invalid GET_HISTORY payload.");
            return;
        }

        var username = parts[1];
        var channelId = parts[2];

        if (ClientKeyService is null ||
            Db is null ||
            string.IsNullOrWhiteSpace(ChannelDbKey))
        {
            Console.WriteLine("History dependencies are not initialized.");
            return;
        }

        var targetClient = Task.Run(async () => await ClientKeyService.GetByUsernameAsync(username))
            .GetAwaiter()
            .GetResult();

        if (targetClient is null)
        {
            Console.WriteLine($"No public key found for history request user {username}");
            return;
        }

        var allMessages = Task.Run(async () => await Db.Select<ChannelMessages>("channel_messages"))
            .GetAwaiter()
            .GetResult();

        var channelMessages = allMessages
            .Where(m => m.ChannelId == channelId)
            .OrderBy(m => m.CreatedAt)
            .ToList();

        Console.WriteLine($"Sending {channelMessages.Count} history messages to {username}");

        var channelCrypto = new ChannelCryptoService();

        foreach (var dbMessage in channelMessages)
        {
            string plainText;

            try
            {
                plainText = channelCrypto.Decrypt(
                    dbMessage.CipherText,
                    dbMessage.Nonce,
                    dbMessage.Tag,
                    ChannelDbKey
                );
            }
            catch (Exception ex)
            {
                Console.WriteLine($"Failed to decrypt DB history row {dbMessage.Id}: {ex.Message}");
                continue;
            }

            var encrypted = E2EeHelper.EncryptForRecipient(plainText, targetClient.PublicKey);

            var outbound = new SocketEncryptedMessage
            {
                Type = "encrypted_chat",
                SenderUsername = ExtractUsernameFromUserId(dbMessage.SenderUserId),
                RecipientUsername = username,
                ChannelId = channelId,
                CipherText = encrypted.CipherText,
                Nonce = encrypted.Nonce,
                Tag = encrypted.Tag,
                EncryptedKey = encrypted.EncryptedKey
            };

            Send(JsonSerializer.Serialize(outbound));
        }
    }
    
    private static string GetRecordId(object? id)
    {
        if (id is null)
            return string.Empty;

        var json = JsonSerializer.Serialize(id);

        using var doc = JsonDocument.Parse(json);

        var root = doc.RootElement;

        var recordId = root.GetProperty("Id").GetString() ?? string.Empty;
        var table = root.GetProperty("Table").GetString() ?? string.Empty;

        return $"{table}:{recordId}";
    }
    
    private void HandleEncryptedRtcSignal(string msg)
    {
        SocketRtcSignalMessage? clientPayload;

        try
        {
            clientPayload = JsonSerializer.Deserialize<SocketRtcSignalMessage>(msg);
        }
        catch
        {
            Console.WriteLine("Failed to parse encrypted RTC signal payload.");
            return;
        }

        if (clientPayload is null || clientPayload.Type != "encrypted_rtc_signal")
            return;

        if (ClientKeyService is null || string.IsNullOrWhiteSpace(ServerPrivateKey))
        {
            Console.WriteLine("Server RTC crypto dependencies are not initialized.");
            return;
        }

        string plainJson;

        try
        {
            plainJson = E2EeHelper.DecryptForRecipient(
                new EncryptedPayload
                {
                    CipherText = clientPayload.CipherText,
                    Nonce = clientPayload.Nonce,
                    Tag = clientPayload.Tag,
                    EncryptedKey = clientPayload.EncryptedKey
                },
                ServerPrivateKey
            );
        }
        catch (Exception ex)
        {
            Console.WriteLine($"Failed to decrypt RTC signal payload: {ex.Message}");
            return;
        }

        var targetClient = Task.Run(async () => await ClientKeyService.GetByUsernameAsync(clientPayload.RecipientUsername))
            .GetAwaiter()
            .GetResult();

        if (targetClient is null)
        {
            Console.WriteLine($"No target RTC client key found for {clientPayload.RecipientUsername}");
            return;
        }

        var encrypted = E2EeHelper.EncryptForRecipient(plainJson, targetClient.PublicKey);

        var outbound = new SocketRtcSignalMessage
        {
            Type = "encrypted_rtc_signal",
            SenderUsername = clientPayload.SenderUsername,
            RecipientUsername = clientPayload.RecipientUsername,
            CipherText = encrypted.CipherText,
            Nonce = encrypted.Nonce,
            Tag = encrypted.Tag,
            EncryptedKey = encrypted.EncryptedKey
        };

        Sessions.Broadcast(JsonSerializer.Serialize(outbound));
    }
}