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; }
    public static string? DefaultChannelId { 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.StartsWith("GET_HISTORY|"))
        {
            HandleGetHistory(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 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) ||
            string.IsNullOrWhiteSpace(DefaultChannelId))
        {
            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 = DefaultChannelId,
                    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,
                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('|', 2);

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

        var username = parts[1];

        if (ClientKeyService is null ||
            Db is null ||
            string.IsNullOrWhiteSpace(ChannelDbKey) ||
            string.IsNullOrWhiteSpace(DefaultChannelId))
        {
            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 == DefaultChannelId)
            .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,
                CipherText = encrypted.CipherText,
                Nonce = encrypted.Nonce,
                Tag = encrypted.Tag,
                EncryptedKey = encrypted.EncryptedKey
            };

            Send(JsonSerializer.Serialize(outbound));
        }
    }
}