online_security_project/client/Program.cs

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Globalization;
using System.IO;
using System.Linq;
using System.Net.Sockets;
using System.Security.Cryptography;
using System.Text;
using System.Text.Json;
using System.Threading.Tasks;
using client;
using lib;

namespace Client;

public class Program
{
    static string User = "";
    static readonly Random Rand = new((int)DateTime.Now.Ticks);
    static byte counter = (byte)Rand.Next();

    static async Task Main(string[] args)
    {
        User = args
            .SkipWhile(a => !new[] { "-p", "--phone" }.Contains(a)) // search for the option
            .Skip(1) // skip the `-u/--user` itself to get the value
            .FirstOrDefault() ?? "0000"; // get the value or deafult if it doesnt exist

        // On boot, check if a key is available
        RSA? serverKey = LoadRSAFromFile("server_key.pem");
        if (serverKey == null) { Console.WriteLine("Could not find server key, please run server before clients!"); return; }
        RSA pubKey = LoadRSAFromFile($"pubkey_{User}.pem") ?? RSA.Create(2048);
        RSA privKey = LoadRSAFromFile($"privkey_{User}.pem") ?? pubKey;

        SaveRSAKeys(User, pubKey, privKey);
        using TcpClient client = new("127.0.0.1", 12345);

        var stream = client.GetStream();

        // First contact init
        bool needsRegister = pubKey == privKey || args.Any(a => new[] { "-fr", "--force-register" }.Contains(a));
        Aes sk = Aes.Create(); // creates an AES-256 key
        // establish secure connection
        byte[] skEnc = serverKey.Encrypt([.. sk.Key, .. sk.IV], RSAEncryptionPadding.OaepSHA256);
        await stream.WriteAsync(skEnc);
        // wait for the server to confirm it recieved the keys
        await stream.ReadExactlyAsync(new byte[1]);

        if (needsRegister)
        {
            try
            {
                await RegisterClient(User, pubKey, privKey, sk, stream);
            }
            catch (Exception ex)
            {
                Console.WriteLine("Failed registration process");
                Console.WriteLine("Exception: " + ex.Message);
                Console.WriteLine("Stack: " + ex.StackTrace);
                return;
            }
        }
        else
        {
            // start login process
            byte[] msg = Request.CreateRequest(RequestType.Login, ref counter, Utils.NumberToBytes(User));
            msg = sk.EncryptCfb(msg, sk.IV, PaddingMode.PKCS7);
            await stream.WriteAsync(msg);
            byte[] toSign = new byte[16];
            await stream.ReadExactlyAsync(toSign, 0, 16);
            toSign = sk.DecryptCfb(toSign, sk.IV, PaddingMode.None);
            byte[] signed = privKey.SignData(toSign, HashAlgorithmName.SHA256, RSASignaturePadding.Pkcs1);
            msg = Request.CreateRequest(RequestType.ConfirmLogin, ref counter, BitConverter.GetBytes(signed.Length));
            msg = sk.EncryptCfb(msg, sk.IV, PaddingMode.None);
            WriteColor($"Sending sig: {Convert.ToBase64String(signed)}", ConsoleColor.Green);
            msg = [.. msg, .. signed];
            await stream.WriteAsync(msg);
            byte[] buffer = new byte[512];
            int len = await stream.ReadAsync(buffer);
            byte[] dec = sk.DecryptCfb(buffer[..len], sk.IV, PaddingMode.PKCS7);
            string r = Encoding.UTF8.GetString(dec);
            if (r == "OK")
            {
                Console.WriteLine("Login successful");
            }
            else
            {
                Console.WriteLine($"Failed login: {r}");
                client.Dispose();
                return;
            }
        }

        await HandleUserInput(client, stream, sk, privKey);

    }

    async static Task RegisterClient(string user, RSA pub, RSA priv, Aes sk, NetworkStream stream)
    {
        Console.WriteLine("Attempting to register with public key:");
        Console.WriteLine(pub.ExportRSAPublicKeyPem());
        // Generate aes key and send it forward
        Console.WriteLine($"Session key: {string.Join(' ', sk.Key)}");
        Console.WriteLine($"Session IV:  {string.Join(' ', sk.IV)}");

        // Generate the Register msg
        Console.WriteLine("Sending rsa public key thing");
        byte[] pubBytes = pub.ExportRSAPublicKey();
        byte[] data = new byte[12];
        Array.Copy(Utils.NumberToBytes(user), data, 8);
        Array.Copy(BitConverter.GetBytes(pubBytes.Length), 0, data, 8, 4);
        byte[] msg = Request.CreateRequest(RequestType.Register, ref counter, data);
        // Encrypt msg and send it
        byte[] payload = [.. msg, .. pubBytes];
        byte[] enc = sk.EncryptCfb(payload, sk.IV, PaddingMode.PKCS7);
        Console.WriteLine($"payload length: {enc.Length}");
        await stream.WriteAsync(enc);

        // get the 6 digit code (from "secure channel", actually an OK message is expected here but the 6 digit code kinda replaces it)
        byte[] digits = new byte[6];
        int len = 0;
        while (len != 6)
        {
            len = await stream.ReadAsync(digits);
        }
        // print the 6 digit code
        Console.WriteLine($"[{DateTime.Now}] 6 digit code: {string.Join(' ', digits.Select(d => d.ToString()))}");
        // get the 6 digit code from the user
        while (true)
        {
            Console.Write("> ");
            string? code = Console.ReadLine()?.Trim();
            if (code == null || code.Take(6).Any(c => !char.IsDigit(c)))
            {
                Console.WriteLine("Invalid code!");
                continue;
            }
            byte[] codeBytes = code
                .Take(6) // take the first 6 characters
                .Select(d => byte.Parse(d.ToString())) // parse into bytes
                .ToArray();
            // Debug print the inserted value to see it works :)
            Console.WriteLine(string.Join(' ', codeBytes.Select(b => b.ToString())));

            // Sign the 6 digit code & Generate ConfirmRegister message
            byte[] signed = priv.SignData(codeBytes, HashAlgorithmName.SHA256, RSASignaturePadding.Pkcs1);
            msg = Request.CreateRequest(RequestType.ConfirmRegister, ref counter, [.. codeBytes, .. BitConverter.GetBytes(signed.Length)]);
            enc = sk.EncryptCfb(msg, sk.IV, PaddingMode.None); // no reason to encrpy the signature
            payload = [.. enc, .. signed]; // should be 128 (enc) + 256 (signed)
            await stream.WriteAsync(payload);
            // wait for OK/NACK response (anything other than OK is a NACK)
            int incoming = await stream.ReadAsync(enc);
            msg = sk.DecryptCfb(enc[..incoming], sk.IV, PaddingMode.PKCS7);
            string r = Encoding.UTF8.GetString(msg);
            if (r == "OK")
            {
                Console.WriteLine("Registration process complete");
                break;
            }
            else
            {
                Console.WriteLine("Failed reigstreation: " + r);
            }
        }
    }

    static RSA? LoadRSAFromFile(string file)
    {
        if (File.Exists(file))
        {
            string content = File.ReadAllText(file);
            RSA k = RSA.Create();
            k.ImportFromPem(content);
            return k;
        }
        return null;
    }

    static void SaveRSAKeys(string user, RSA pub, RSA priv)
    {
        File.WriteAllText($"pubkey_{user}.pem", pub.ExportRSAPublicKeyPem());
        File.WriteAllText($"privkey_{user}.pem", priv.ExportRSAPrivateKeyPem());
    }

    static async Task HandleUserInput(TcpClient client, NetworkStream stream, Aes sk, RSA privKey)
    {
        string? currentChat = null;
        Dictionary<string, RSA> publicKeys = [];
        while (client.Connected)
        {
            Console.Write("> ");
            string? input = Console.ReadLine();
            if (input == null)
            {
                await Task.Delay(100);
                continue;
            }
            else if (input.StartsWith('/'))
            {
                string[] words = input.Split(' ');
                // Commands :D, i like commands
                switch (words[0].ToLower())
                {
                    case "/quit":
                    case "/exit":
                    case "/q!":
                        return;
                    case "/chat":
                    case "/msg":
                        string? old = currentChat;
                        currentChat = words.Length > 1 ? words[1] : null;
                        if (currentChat != null && currentChat.Length > 16)
                        {
                            Console.WriteLine("Invalid number: too long");
                            currentChat = old;
                            continue;
                        }
                        if (currentChat != null && !publicKeys.ContainsKey(currentChat))
                        {
                            // attempt to get the currnet chat's key, this is also possible to do just before sending a message
                            // but i decided to do it here since if better fits the current structure
                            RSA? key = await GetPublicKey(stream, sk, currentChat);
                            if (key != null)
                            {
                                publicKeys[currentChat] = key;
                            }
                            else
                            {
                                currentChat = old;
                                Console.WriteLine($"Reverting to previous chat: {currentChat ?? "none"}");
                            }
                        }
                        break;
                    case "/read":
                    case "/get":
                    case "/fetch":
                    case "/pull":
                        await GetMessages(stream, sk, privKey, publicKeys);
                        break;
                }
            }
            else
            {
                if (currentChat == null)
                {
                    Console.WriteLine("No chat is active, please select chat using '/msg [number]' or '/chat [number]'");
                }
                else if (publicKeys.TryGetValue(currentChat!, out RSA? key))
                {
                    // Pick the id as a random uint, good enough for the sake of testing
                    Message m = new((uint)Rand.Next(), User, false, input);
                    m.CalculateSignature(privKey);
                    byte[] userMsg = [.. key.Encrypt(m.Bytes(), RSAEncryptionPadding.OaepSHA256), .. m.Signature];
                    byte[] req = Request.CreateRequest(
                        RequestType.SendMessage,
                        ref counter,
                        [.. Utils.NumberToBytes(currentChat!), .. BitConverter.GetBytes(userMsg.Length)]);
                    req = sk.EncryptCfb(req, sk.IV);
                    stream.Write([.. req, .. userMsg]);
                    Console.WriteLine($"[{DateTime.Now}] Sent to server: {input}");
                }
                else
                {
                    Console.WriteLine($"active chat exists, but no key was found...");
                }
            }
        }
    }

    static async Task<RSA?> GetPublicKey(NetworkStream stream, Aes sk, string chat)
    {
        byte[] req = Request.CreateRequest(RequestType.GetUserKey, ref counter, Utils.NumberToBytes(chat));
        req = sk.EncryptCfb(req, sk.IV, PaddingMode.None); // no need for padding this is exactly 128 bytes
        await stream.WriteAsync(req);
        byte[] response = new byte[1024];
        int len = await stream.ReadAsync(response);
        byte[] key = sk.DecryptCfb(response[..len], sk.IV, PaddingMode.PKCS7);
        if (key[0] == 1)
        {
            Console.WriteLine($"failed getting key for {chat}: {Encoding.UTF8.GetString(key[1..])}");
            return null;
        }
        else
        {
            RSA bobsKey = RSA.Create();
            bobsKey.ImportRSAPublicKey(key.AsSpan()[1..], out int _);
            Console.WriteLine($"Got key:\n {bobsKey.ExportRSAPublicKeyPem()}\n");
            return bobsKey;
        }
    }

    static async Task GetMessages(NetworkStream stream, Aes sk, RSA privKey, Dictionary<string, RSA> publicKeys)
    {
        byte[] req = Request.CreateRequest(RequestType.GetMessages, ref counter, []);
        req = sk.EncryptCfb(req, sk.IV, PaddingMode.None); // no need for padding this is exactly 128 bytes
        await stream.WriteAsync(req);
        byte[] buffer = new byte[4096];
        int len = await stream.ReadAsync(buffer);
        byte[] lengths = buffer[..16];
        lengths = sk.DecryptCfb(lengths, sk.IV, PaddingMode.None);
        Console.ForegroundColor = ConsoleColor.Cyan;
        Console.WriteLine($"Got messages: {string.Join(", ", lengths)}");
        Console.ResetColor();
        byte[] msg = new byte[1024]; // msg buffer

        int start = 16; // skip the first 16 bytes since its the lengths message
        for (int i = 0; i < lengths.Length - 1; i += 2)
        {
            int l = (lengths[i] << 8) | (lengths[i + 1]);
            Console.WriteLine($"got message of length: {l}");
            if (l == 0) { break; } // a 0 means we are done actually, as empty messages shouldn't be allowed
            // get the msg
            int end = start + l;
            if (end > len)
            {
                // TODO: properly handle it as a buffered stream and properly buffer it
            }

            WriteColor($"got ecnryped message: {Convert.ToBase64String(buffer[start..end])}", ConsoleColor.Green);
            // decrypt the message
            int msgLen = privKey.KeySize / 8;
            if (privKey.TryDecrypt(buffer.AsSpan()[start..(msgLen + start)], msg, RSAEncryptionPadding.OaepSHA256, out int written))
            {
                byte[] dec = msg[..written];
                Message m = new(dec)
                {
                    Signature = buffer[(start + msgLen)..end]
                };

                bool sigValid = false;
                if (publicKeys.TryGetValue(m.Sender, out RSA? pk) || (pk = await GetPublicKey(stream, sk, m.Sender)) != null)
                {
                    publicKeys[m.Sender] = pk; // update the key value stored in case it was requested for in GetPublicKey
                    sigValid = m.IsSignatureValid(pk);
                }
                else
                {
                    Console.WriteLine("Failed getting sender's public key");
                }
                WriteColor($"decrypted message: {Convert.ToBase64String(dec)}", ConsoleColor.Green);
                Console.WriteLine($"Got message from {m.Sender} (id: {m.Id}) (valid: {sigValid}): " + (m.IsAck ? "ACK" : m.Content));
                // if the signature is valid we also want to send back a message ack for the same id :)
                if (sigValid)
                {
                    Message mAck = new(m.Id, User, true, "");
                    mAck.CalculateSignature(privKey);
                    byte[] userMsg = [.. publicKeys[m.Sender].Encrypt(mAck.Bytes(), RSAEncryptionPadding.OaepSHA256), .. mAck.Signature];
                    byte[] reqAck = Request.CreateRequest(
                        RequestType.SendMessage,
                        ref counter,
                        [.. Utils.NumberToBytes(m.Sender), .. BitConverter.GetBytes(userMsg.Length)]);
                    reqAck = sk.EncryptCfb(reqAck, sk.IV);
                    stream.Write([.. reqAck, .. userMsg]);
                }
            }
            else
            {
                // what if we cant decrypt the message? well, i doubt there is anything we can do about it
                // even if we know who sent the message we dont know which message it is unless the server also knows
                // and i dont want the server to be aware of that, i think it makes more sense for the server to act as a relay
                // and a buffer than an actual participant, so if the message is failing to decrypt that will go unnoticed.
                // supposedly the sender will notice the lack of response and send it again
                Console.WriteLine("Incoming message failed to decrypt, unknown sender");
            }
            start = end;
        }

        bool hasMoreMessages = lengths[15] != 0;
        Console.WriteLine($"hasMoreMessages: {hasMoreMessages}");
    }

    static void WriteColor(string Line, ConsoleColor Color, ConsoleColor Background = ConsoleColor.Black)
    {
        Console.ForegroundColor = Color;
        Console.BackgroundColor = Background;
        Console.WriteLine(Line);
        Console.ResetColor();
    }
}