online_security_project/server/Program.cs

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

namespace Server;

public class Program
{
    const int MSG_LEN = 16; // msg len is 128 bits = 16 bytes

    static readonly Data Data = new();
    static readonly Random Rnd = new((int)DateTime.Now.Ticks);

    static async Task Main()
    {
        // Generally this key would be static but since its not production yet we can generate it every time to make sure
        // the users has the key and could load it from file
        RSA key = RSA.Create(1024);
        File.WriteAllText("server_key.pem", key.ExportRSAPublicKeyPem());

        int port = 12345;
        TcpListener server = new(IPAddress.Parse("0.0.0.0"), port);
        int connectionCounter = 0;
        try
        {
            server.Start();
            byte[] buffer = new byte[256];

            while (true)
            {
                // Currently, every time it gets a block, it will simply send it back but ToUpper
                TcpClient client = await server.AcceptTcpClientAsync();
                _ = Task.Run(async () => await HandleClient(client, connectionCounter));
                connectionCounter += 1;
            }
        }
        catch (Exception ex)
        {
            Console.WriteLine($"Server error: {ex.Message}");
            Console.WriteLine("Trace: " + ex.StackTrace);
        }
        finally
        {
            server.Stop();
        }
    }

    static async Task HandleClient(TcpClient client, int id)
    {
        NetworkStream stream = client.GetStream();
        byte[] buffer = new byte[1024];
        byte counter = 0;
        // Get AES session key
        int len = await stream.ReadAsync(buffer);
        Aes sk = Aes.Create();
        sk.Key = buffer[..256]; // just to make sure no one sends a too big to be true key
        sk.IV = buffer[256..len];
        Write(id, "key + iv: " + len.ToString());

        // Get first message (should be either login or )
        len = await stream.ReadAsync(buffer);
        byte[] msg = sk.DecryptCfb(buffer[..MSG_LEN], sk.IV, PaddingMode.PKCS7);
        if (msg[0] != 0)
        {
            Write(id, "Invalid session id!");
            client.Dispose();
            return;
        }
        counter = IncrementCounter(msg[2]); // allow counter to start at a random position
        if (msg[1] == (byte)RequestType.Register)
        {
            // Do register stuff
            // get phone number
            string phone = Utils.BytesToNumber(msg[3..11]);
            Write(id, $"Client wants to register as {phone}");
            int keyLen = BitConverter.ToInt32(msg, 11);
            RSA pub = RSA.Create();
            pub.ImportRSAPublicKey(buffer.AsSpan()[MSG_LEN..], out int bytesRead);
            Write(id, $"Imported key len: {bytesRead} while client claims it is {keyLen}");
            // generate the 6 digit code and send it
            byte[] code = [
                (byte)Rnd.Next(10),
                (byte)Rnd.Next(10),
                (byte)Rnd.Next(10),
                (byte)Rnd.Next(10),
                (byte)Rnd.Next(10),
                (byte)Rnd.Next(10),
                ];
            await Send6DigitCodeInSecureChannel(stream, code);
            // wait for the code to be back with a key
            int tries = 5; // allow 5 tries before closing the connection and forcing a restart
            while (tries > 0)
            {
                tries -= 1;
                len = await stream.ReadAsync(buffer);
                Write(id, $"Got 6 digit code with sig, len: {len}");
                msg = sk.DecryptCfb(buffer[..MSG_LEN], sk.IV, PaddingMode.PKCS7);
                byte[] sig = buffer[MSG_LEN..len];
                if (msg[0] != 0 || msg[1] != (byte)RequestType.ConfirmRegister || msg[2] != counter)
                {
                    // invalid or unexpected req, someone might be sending dups
                    continue;
                }
                counter = IncrementCounter(counter);
                byte[] gottenCode = msg[3..9];
                int expectedSigLen = BitConverter.ToInt32(msg, 9);
                if (expectedSigLen != len - MSG_LEN)
                {
                    Write(id, $"expected sig len doesnt match read len: {expectedSigLen} / {len - MSG_LEN}");
                }
                // check if the codes are equal
                if (code.Zip(gottenCode).Any(a => a.First != a.Second))
                {
                    // codes are not equal, send a nack
                    msg = sk.EncryptCfb(Encoding.UTF8.GetBytes("BAD CODE"), sk.IV, PaddingMode.PKCS7);
                    await stream.WriteAsync(msg);
                }
                else
                {
                    // codes are equal - verify sig
                    bool sigValid = pub.VerifyData(gottenCode, sig, HashAlgorithmName.SHA256, RSASignaturePadding.Pkcs1);
                    if (sigValid)
                    {
                        msg = sk.EncryptCfb(Encoding.UTF8.GetBytes("OK"), sk.IV, PaddingMode.PKCS7);
                        await stream.WriteAsync(msg);
                        Data.Keys[phone] = pub; // save the key
                        break;
                    }
                    else
                    {
                        msg = sk.EncryptCfb(Encoding.UTF8.GetBytes("SIG INVALID"), sk.IV, PaddingMode.PKCS7);
                        await stream.WriteAsync(msg);
                    }
                }
            }

        }
        else if (msg[1] == (byte)RequestType.Login)
        {
            // verify login
        }
        else
        {
            // invalid connection, quit
            Write(id, "Client didnt register or login as first message");
            client.Dispose();
            return;
        }
        // Client registered/logged in, do main messages loop
        try
        {
            while (client.Connected)
            {
                // while the client is connected, simply read messages from the client and handle accordingly,
                // either by getting new messages for other ppl, or sending back keys/pending messages
                len = await stream.ReadAsync(buffer);
                msg = sk.DecryptCfb(buffer[..MSG_LEN], sk.IV, PaddingMode.None);
                // verify that the counter message is correct
                if (msg[0] != 0 || msg[2] != counter)
                {
                    msg = sk.EncryptCfb(Encoding.UTF8.GetBytes("DUPLICATE"), sk.IV, PaddingMode.PKCS7);
                    await stream.WriteAsync(msg);
                    continue;
                }
                counter = IncrementCounter(counter);
                switch ((RequestType)msg[1])
                {
                    case RequestType.GetMessages:
                        break;
                    case RequestType.GetUserKey:
                        string phone = Utils.BytesToNumber(msg[3..11]);
                        RSA? key = Data.GetKey(phone);
                        if (key != null)
                        {
                            msg = sk.EncryptCfb(key.ExportRSAPublicKey(), sk.IV, PaddingMode.PKCS7);
                            await stream.WriteAsync(msg);
                        }
                        else
                        {
                            msg = sk.EncryptCfb(Encoding.UTF8.GetBytes("USER DOES NOT EXIST"), sk.IV, PaddingMode.PKCS7);
                            await stream.WriteAsync(msg);
                        }
                        break;
                    case RequestType.SendMessage:
                        break;
                    default:
                        msg = sk.EncryptCfb(Encoding.UTF8.GetBytes("INVALID REQUEST"), sk.IV, PaddingMode.PKCS7);
                        await stream.WriteAsync(msg);
                        break;
                }

            }

        }
        catch (Exception ex)
        {
            Write(id, $"Client failed with error {ex.Message}");
            Write(id, $"Stack: {ex.StackTrace}");
        }

        client.Dispose();
    }

    static byte IncrementCounter(byte counter)
    {
        return counter == byte.MaxValue ? (byte)0 : (byte)(counter + 1);
    }

    static async Task Send6DigitCodeInSecureChannel(NetworkStream stream, byte[] code)
    {
        await stream.WriteAsync(code);
    }

    /// Helper log message so it would print both the time and the id
    static void Write(int id, string Message)
    {
        Console.WriteLine($"[{DateTime.Now:HH:mm:ss}] {id} - {Message}");
    }
}