online_security_project/client/Program.cs

using System;
using System.Net.Sockets;
using System.Text;
using System.Threading.Tasks;
using System.Security.Cryptography;
using System.Linq;
using System.IO;
using lib;

namespace Client;

public class Program
{
    static async Task Main(string[] args)
    {
        string 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("key_server.pem")!;
        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
        if (needsRegister)
        {
            while(true) {
                try {
                    await RegisterClient(user, pubKey, privKey, serverKey, sk, stream);
                }
                catch (Exception ex) {
                    Console.WriteLine("Failed registration process");
                    Console.WriteLine("Exception: " + ex.Message);
                    Console.WriteLine("Stack: " + ex.StackTrace);
                }
            }
        }
        else
        {

        }


        var inputTask = Task.Run(async () => await HandleUserInput(client, stream));
        var serverInput = Task.Run(async () => await HandleServerInput(client, stream));

        _ = Task.WaitAny(inputTask, serverInput);

    }

    async static Task RegisterClient(string user, RSA pub, RSA priv, RSA server, Aes sk, NetworkStream stream)
    {
        byte counter = 0;
        // first send the `Register` message
        byte[] pubBytes = pub.ExportRSAPublicKey();
        byte[] data = new byte[16];
        Array.Copy(Utils.NumberToBytes(user), data, 8);
        Array.Copy(BitConverter.GetBytes(pubBytes.Length), 0, data, 8, 4);
        Array.Copy(BitConverter.GetBytes(sk.Key.Length), 0, data, 12, 4);
        byte[] msg = Request.CreateRequest(RequestType.Register, ref counter, data);
        byte[] finalPayload = [.. msg, .. pubBytes, .. sk.Key];
        // signing a hash here is pretty useless tbh
        // if oscar is changing this message somehow it will just fail the registration process and will restart eventually
        // so kinda pointless, esp when its signed with the server's key
        byte[] enc = server.Encrypt(finalPayload, RSAEncryptionPadding.OaepSHA256);
        await stream.WriteAsync(enc);
        // get the 6 digit code
        byte[] digits = new byte[6];
        await stream.ReadExactlyAsync(digits, 0, 6);
        // 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) {
            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();
            msg = Request.CreateRequest(RequestType.ConfirmRegister, ref counter, codeBytes);
            byte[] signed = priv.SignData(msg, HashAlgorithmName.SHA256, RSASignaturePadding.Pkcs1);
            enc = sk.EncryptCfb(msg, Array.Empty<byte>(), PaddingMode.None); // no reason to encrpy the signature
            finalPayload = [.. enc, .. signed]; // should be 128 (enc) + 256 (signed)
            await stream.WriteAsync(finalPayload);
            int incoming = await stream.ReadAsync(enc);
            msg = sk.DecryptCfb(enc[..incoming], Array.Empty<byte>(), PaddingMode.Zeros);
            string r = Encoding.UTF8.GetString(msg);
            if(r == "OK") {
                Console.WriteLine("Registration process complete");
                break;
            }
        }
    }

    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)
    {
        while (client.Connected)
        {
            string? input = Console.ReadLine();
            if (input == null)
            {
                await Task.Delay(100);
                continue;
            }
            else if (input.StartsWith('/'))
            {
                // Commands :D, i like commands
                switch (input.ToLower())
                {
                    case "/quit":
                    case "/exit":
                    case "/q!":
                        return;
                }
            }
            else
            {
                stream.Write(Encoding.ASCII.GetBytes(input));
                Console.WriteLine($"[{DateTime.Now}]Sent to server: {input}");
            }
        }
    }

    static async Task HandleServerInput(TcpClient client, NetworkStream stream)
    {
        byte[] buffer = new byte[1024];
        while (client.Connected)
        {
            int readLen = await stream.ReadAsync(buffer);
            if (readLen != 0)
            {
                string fromServer = Encoding.ASCII.GetString(buffer[..readLen]);
                Console.WriteLine($"[{DateTime.Now}]  From server: {fromServer}");

            }
        }
    }
}