e2e-irc/src/main.rs

use base64::{engine::general_purpose, Engine as _};
use ircparser;
use pgp::crypto::sym::SymmetricKeyAlgorithm;
use pgp::ser::Serialize;
use pgp::{Deserializable, Message, SignedPublicKey, SignedSecretKey};
use rand::prelude::*;
use std::collections::HashMap;
use std::net::{Shutdown, TcpStream};
use std::sync::mpsc;
use std::thread;
use std::{env, fs};

mod encryption;
mod listener_server;
mod message_stream;
mod writer_client;

macro_rules! try_recv {
    ($rx:ident) => {
        match $rx.try_recv() {
            Ok(recieved) => recieved,
            Err(_e) => match _e {
                mpsc::TryRecvError::Empty => String::new(),
                mpsc::TryRecvError::Disconnected => {
                    panic!("Error!")
                }
            },
        }
    };
}

fn bytes_to_privmsg_base64(message: Vec<u8>, reciever: &str) -> String {
    let mut command = String::new();
    for line in general_purpose::STANDARD
        .encode(message)
        .chars()
        .collect::<Vec<char>>()
        .chunks(50)
        .map(|c| c.iter().collect::<String>())
    {
        command.push_str(&format!("PRIVMSG {reciever} {line}\n"));
    }

    command
}

fn encrypt(key: &SignedPublicKey, message: &str) -> Result<Vec<u8>, pgp::errors::Error> {
    let message = Message::new_literal("none", message);
    let mut rng = StdRng::from_entropy();

    let message = message.encrypt_to_keys(&mut rng, SymmetricKeyAlgorithm::AES128, &[key])?;

    Ok(message.to_bytes()?)
}

fn decrypt(key: &SignedSecretKey, message: Vec<u8>, password: &str) -> Result<&str, pgp::errors::Error> {
    let message = Message::from_bytes(message.as_slice())?;

    let (message, _) = message.decrypt(|| password.to_string(), &[key])?;

    Ok(String::from_utf8_lossy(
        message.get_content()?.unwrap_or(String::new()),
    ))
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let args = env::args().collect::<Vec<String>>();

    let server = args[1].clone();
    let default_password = String::new();
    let passwd = args.get(2).unwrap_or(&default_password);

    let stream = TcpStream::connect(format!("{server}:6697"))?;

    let public_key = fs::read("public.gpg")?;
    let secret_key = SignedSecretKey::from_bytes(fs::read("secret.gpg")?.as_slice())?;

    let reader_stream = match stream.try_clone() {
        Ok(stream) => stream,
        Err(_error) => {
            let _ = stream.shutdown(Shutdown::Both);
            panic!("Failed to create the reader stream")
        }
    };

    let writer_stream = match stream.try_clone() {
        Ok(stream) => stream,
        Err(_error) => {
            let _ = stream.shutdown(Shutdown::Both);
            let _ = reader_stream.shutdown(Shutdown::Both);
            panic!("Failed to create the writer stream")
        }
    };

    let (listener_channel_send_tx, listener_channel_rx) = mpsc::channel();
    let (listener_channel_tx, listener_channel_recv_rx) = mpsc::channel();

    let (writer_channel_tx, writer_channel_send_rx) = mpsc::channel();
    let (writer_channel_recv_tx, writer_channel_rx) = mpsc::channel();

    thread::spawn(|| {
        listener_server::listen_to_client(listener_channel_send_tx, listener_channel_recv_rx)
    });
    thread::spawn(|| {
        writer_client::write_to_server(
            writer_stream,
            server,
            writer_channel_send_rx,
            writer_channel_recv_tx,
        )
    });

    let mut keys: HashMap<String, SignedPublicKey> = HashMap::new();
    let mut userstring: Option<String> = None;

    loop {
        let recieved = try_recv!(listener_channel_rx);

        if recieved != "" {
            let command = &ircparser::parse(&recieved).expect("Got an invalid IRC instruction")[0];

            if command.command == "PRIVMSG" && !command.params[0].starts_with("#") {
                let reciever = &command.params[0];
                if !keys.contains_key(reciever) {
                    dbg!(recieved);

                    writer_channel_tx.send(format!("PRIVMSG {reciever} START_KEY\n"))?;
                    writer_channel_tx
                        .send(bytes_to_privmsg_base64(public_key.clone(), reciever))?;
                    writer_channel_tx.send(format!("PRIVMSG {reciever} END_KEY\n"))?;

                    let mut foreign_key: Vec<String> = Vec::new();
                    while !foreign_key.contains(&"END_KEY".to_string()) {
                        let recieved_raw = writer_channel_rx.recv()?;

                        let recieved = &ircparser::parse(&recieved_raw).unwrap()[0];
                        let begin_source_reciever = format!(":{reciever}!");
                        if recieved.command != "PRIVMSG"
                            || !recieved
                                .source
                                .clone()
                                .unwrap_or("".to_string())
                                .starts_with(&begin_source_reciever)
                        {
                            writer_channel_tx.send(recieved_raw)?;
                            continue;
                        }

                        foreign_key.push(recieved.params[1].clone());
                    }
                    foreign_key.pop();
                    let foreign_key = foreign_key.concat();

                    let foreign_key = general_purpose::STANDARD.decode(foreign_key)?;

                    let key = SignedPublicKey::from_bytes(foreign_key.as_slice())?;

                    keys.insert(reciever.to_string(), key);
                }

                let foreign_key = keys.get(reciever).unwrap();

                writer_channel_tx.send("START_MESSAGE".to_string())?;
                writer_channel_tx.send(bytes_to_privmsg_base64(
                    encrypt(&foreign_key, &command.params[1])?,
                    reciever,
                ))?;
                writer_channel_tx.send("END_MESSAGE".to_string())?;
            } else {
                dbg!(&recieved);
                writer_channel_tx.send(recieved)?;
            }
        }

        let recieved = try_recv!(writer_channel_rx);
        if recieved != "" {
            let recieved_parsed = &ircparser::parse(&recieved).unwrap()[0];

            if recieved_parsed.command != "PRIVMSG"
                || recieved_parsed
                    .params
                    .get(0)
                    .unwrap_or(&"".to_string())
                    .starts_with("#")
            {
                continue;
            }

            let reciever = recieved_parsed.source.clone();
            let reciever = {
                let full_reciever = reciever.unwrap().chars().collect::<Vec<char>>();
                let start_pos = full_reciever.iter().position(|&x| x == ':').unwrap() + 1;
                let end_pos = full_reciever.iter().position(|&x| x == '!').unwrap();
                full_reciever[start_pos..end_pos].iter().collect::<String>()
            };

            if recieved_parsed.params[1] == "START_MESSAGE" {
                let mut encrypted_message: Vec<String> = Vec::new();
                while !encrypted_message.contains(&"END_MESSAGE".to_string()) {
                    let recieved_raw = writer_channel_rx.recv()?;
                    let recieved = &ircparser::parse(&recieved_raw).unwrap()[0];

                    let begin_source_reciever = format!(":{reciever}!");
                    if recieved.command != "PRIVMSG"
                        || !recieved
                            .source
                            .clone()
                            .unwrap_or("".to_string())
                            .starts_with(&begin_source_reciever)
                    {
                        writer_channel_tx.send(recieved_raw)?;
                        continue;
                    }

                    encrypted_message.push(recieved.params[1].clone());
                }
                encrypted_message.pop();

                let encrypted_message = encrypted_message.join("");
                let encrypted_message = general_purpose::STANDARD.decode(encrypted_message)?;
                let encrypted_message = Message::from_bytes(encrypted_message.as_slice())?;

                let passwd = passwd.clone();
                let decrypted_message =
                    encrypted_message.decrypt(|| passwd.to_string(), &[&secret_key]);
            }
            listener_channel_tx.send(recieved)?;
        }
    }
}