kkv 1.2.1

/**************************************************************************
 *   src/server.rs  --  This file is part of kkv.                         *
 *                                                                        *
 *   Copyright (C) 2025 Mateo Lafalce                                     *
 *                                                                        *
 *   kkv is free software: you can redistribute it and/or modify          *
 *   it under the terms of the GNU General Public License as published    *
 *   by the Free Software Foundation, either version 3 of the License,    *
 *   or (at your option) any later version.                               *
 *                                                                        *
 *   kkv is distributed in the hope that it will be useful,               *
 *   but WITHOUT ANY WARRANTY; without even the implied warranty          *
 *   of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.              *
 *   See the GNU General Public License for more details.                 *
 *                                                                        *
 *   You should have received a copy of the GNU General Public License    *
 *   along with this program.  If not, see http://www.gnu.org/licenses/.  *
 *                                                                        *
 **************************************************************************/

use crate::{
    axum::ValidatedForm,
    message::Message,
    types::{
        DataCollection, HexPubkey, IndexMsg, KeyExchange, Keypair, OnionPubkey, PoW, RequestPubkey,
        SentMessage,
    },
    utils::{
        derive_key, deserialize_keys, get_last_elements, get_publickey_bytes, new_nonce,
        sign_message,
    },
};
use argon2::{Argon2, PasswordHash, PasswordVerifier};
use axum::{
    extract::State,
    http::{header::AUTHORIZATION, HeaderMap, StatusCode},
    routing::{get, post},
    Form, Json, Router,
};
use blake3_pow::verify;
use chacha20poly1305::{aead::Aead, ChaCha20Poly1305, KeyInit};
use colored::*;
use ed25519_dalek::{Signature, SigningKey, Verifier, VerifyingKey};
use std::{
    collections::HashMap,
    io::Error,
    sync::{Arc, Mutex},
    time::{Duration, SystemTime, UNIX_EPOCH},
};
use tokio::net::TcpListener;
use x25519_dalek::{PublicKey, ReusableSecret};

/// Port where kkv works
pub const PORT: &str = "45069";

/// State of the app
#[derive(Clone)]
pub struct AppState {
    valid_pubkey_for_ms: Arc<Mutex<HashMap<OnionPubkey, Keypair>>>,
    messages: Arc<Mutex<Vec<Message>>>,
    max_per_day: usize,
    cost_pow: u32,
    expected_time: u64,
    ed25519_bytes: [u8; 64],
    hashed_psw: String,
}

impl AppState {
    fn new(max_per_day: usize, cost_pow: u32, expected_time: u64, hashed_psw: String) -> Self {
        let mut csprng = rand::rngs::OsRng;
        let signing_key: SigningKey = SigningKey::generate(&mut csprng);
        // it's secure since it created here (it's not a )
        let ed25519_bytes = signing_key.to_keypair_bytes();
        AppState {
            valid_pubkey_for_ms: Arc::new(Mutex::new(HashMap::new())),
            messages: Arc::new(Mutex::new(Vec::new())),
            max_per_day,
            cost_pow,
            expected_time,
            ed25519_bytes,
            hashed_psw,
        }
    }
}

pub async fn server(
    max_per_day: usize,
    cost_pow: u32,
    expected_time: u64,
    hashed_psw: String,
) -> Result<(), Error> {
    let state = AppState::new(max_per_day, cost_pow, expected_time, hashed_psw);
    let app: Router = Router::new()
        .route("/get_domain_pk", get(get_domain_pk))
        .route("/pow", post(pow))
        .route("/exchange_keys", post(exchange_keys))
        .route("/msg", post(msg))
        .route("/status", get(status))
        .route("/get_msg", post(get_msg))
        .route("/read_msg", post(read_msg))
        .route("/del_msg", post(del_msg))
        .with_state(state);

    let listener: TcpListener = match TcpListener::bind(format!("127.0.0.1:{PORT}")).await {
        Ok(listener) => listener,
        Err(e) => {
            println!("{}", e.to_string().red());
            return Ok(());
        }
    };

    match axum::serve(listener, app).await {
        Ok(_) => (),
        Err(e) => {
            println!("{}", e.to_string().red());
            return Ok(());
        }
    };

    Ok(())
}

pub async fn status(method: HeaderMap, State(state): State<AppState>) -> StatusCode {
    let client_psw = &method[AUTHORIZATION].as_bytes();
    let parsed_hash = PasswordHash::new(&state.hashed_psw).unwrap();

    if Argon2::default()
        .verify_password(client_psw, &parsed_hash)
        .is_err()
    {
        return StatusCode::UNAUTHORIZED;
    }

    let mut valids = match state.valid_pubkey_for_ms.lock() {
        Ok(guard) => guard,
        Err(_) => return StatusCode::UNAUTHORIZED,
    };
    update_senders_hashmap(&mut valids);
    StatusCode::OK
}

pub async fn del_msg(
    method: HeaderMap,
    State(state): State<AppState>,
    Form(input): Form<IndexMsg>,
) -> StatusCode {
    let client_psw = &method[AUTHORIZATION].as_bytes();
    let parsed_hash = PasswordHash::new(&state.hashed_psw).unwrap();

    if Argon2::default()
        .verify_password(client_psw, &parsed_hash)
        .is_err()
    {
        return StatusCode::UNAUTHORIZED;
    }

    let mut valids = match state.valid_pubkey_for_ms.lock() {
        Ok(guard) => guard,
        Err(_) => return StatusCode::UNAUTHORIZED,
    };
    update_senders_hashmap(&mut valids);

    let index = input.get_index();

    let actual_size = match state.messages.lock() {
        Ok(guard) => guard.len(),
        Err(_) => return StatusCode::UNAUTHORIZED,
    };

    if index >= actual_size {
        return StatusCode::UNAUTHORIZED;
    }

    match state.messages.lock() {
        Ok(mut guard) => {
            guard.remove(index);
            StatusCode::OK
        }
        Err(_) => StatusCode::UNAUTHORIZED,
    }
}

pub async fn read_msg(
    method: HeaderMap,
    State(state): State<AppState>,
    Form(input): Form<IndexMsg>,
) -> StatusCode {
    let client_psw = &method[AUTHORIZATION].as_bytes();
    let parsed_hash = PasswordHash::new(&state.hashed_psw).unwrap();

    if Argon2::default()
        .verify_password(client_psw, &parsed_hash)
        .is_err()
    {
        return StatusCode::UNAUTHORIZED;
    }

    let mut valids = match state.valid_pubkey_for_ms.lock() {
        Ok(guard) => guard,
        Err(_) => return StatusCode::UNAUTHORIZED,
    };
    update_senders_hashmap(&mut valids);

    let index = input.get_index();

    let actual_size = match state.messages.lock() {
        Ok(guard) => guard.len(),
        Err(_) => return StatusCode::UNAUTHORIZED,
    };

    if index >= actual_size {
        return StatusCode::UNAUTHORIZED;
    }

    match state.messages.lock() {
        Ok(mut guard) => {
            guard[index].read_message();
            StatusCode::OK
        }
        Err(_) => StatusCode::UNAUTHORIZED,
    }
}

pub async fn get_msg(
    method: HeaderMap,
    State(state): State<AppState>,
    ValidatedForm(input): ValidatedForm<IndexMsg>,
) -> Json<DataCollection> {
    let client_psw = &method[AUTHORIZATION].as_bytes();
    let parsed_hash = PasswordHash::new(&state.hashed_psw).unwrap();

    if Argon2::default()
        .verify_password(client_psw, &parsed_hash)
        .is_err()
    {
        return Json(DataCollection::default());
    }

    let mut valids = match state.valid_pubkey_for_ms.lock() {
        Ok(guard) => guard,
        Err(_) => return Json(DataCollection::default()),
    };
    update_senders_hashmap(&mut valids);

    let index = input.get_index();

    match state.messages.lock() {
        Ok(frame_of_msg) => {
            let last = get_last_elements(frame_of_msg.to_vec(), index);
            let data = DataCollection::new(frame_of_msg.len(), last);
            Json(data)
        }
        Err(_) => Json(DataCollection::default()),
    }
}

pub async fn msg(
    State(state): State<AppState>,
    ValidatedForm(input): ValidatedForm<SentMessage>,
) -> StatusCode {
    let mut valid_pubkey_for_ms = match state.valid_pubkey_for_ms.lock() {
        Ok(guard) => guard,
        Err(_) => return StatusCode::UNAUTHORIZED,
    };
    update_senders_hashmap(&mut valid_pubkey_for_ms);

    let pubkey = match hex::decode(input.get_pubkey()) {
        Ok(bytes) => bytes,
        Err(_) => return StatusCode::UNAUTHORIZED,
    };

    let chiphertext = match hex::decode(input.get_chiphertext()) {
        Ok(bytes) => bytes,
        Err(_) => return StatusCode::UNAUTHORIZED,
    };

    let signature = match hex::decode(input.get_signature()) {
        Ok(bytes) => bytes,
        Err(_) => return StatusCode::UNAUTHORIZED,
    };

    let alice_pk: [u8; 32] = match pubkey.try_into() {
        Ok(alice_pk) => alice_pk,
        Err(_) => return StatusCode::UNAUTHORIZED,
    };

    if valid_pubkey_for_ms.get(&alice_pk).is_none() {
        return StatusCode::UNAUTHORIZED;
    }

    let client_verifying_key = match VerifyingKey::from_bytes(&alice_pk) {
        Ok(client_verifying_key) => client_verifying_key,
        Err(_) => return StatusCode::UNAUTHORIZED,
    };

    let chiphertext_sig_bytes: [u8; 64] = match signature.try_into() {
        Ok(chiphertext_sig) => chiphertext_sig,
        Err(_) => return StatusCode::UNAUTHORIZED,
    };
    let chiphertext_sig = Signature::from_bytes(&chiphertext_sig_bytes);

    if client_verifying_key
        .verify(&chiphertext, &chiphertext_sig)
        .is_err()
    {
        return StatusCode::UNAUTHORIZED;
    }

    let secretkey = match valid_pubkey_for_ms.get(&alice_pk) {
        Some(secretkey) => secretkey.get_bob_diffie_hellman_sk(),
        None => return StatusCode::UNAUTHORIZED,
    };

    let alice_pubkey = match valid_pubkey_for_ms.get(&alice_pk) {
        Some(alice_pubkey) => alice_pubkey.get_alice_diffie_hellman_pk(),
        None => return StatusCode::UNAUTHORIZED,
    };

    let shared_secret = secretkey.diffie_hellman(&alice_pubkey);
    let fixed_12bytes: &[u8; 12] = shared_secret.as_bytes()[0..12]
        .try_into()
        .unwrap_or(&[0; 12]);
    let shared_nonce = chacha20poly1305::Nonce::from_slice(fixed_12bytes);

    let key = derive_key(shared_secret.as_bytes());
    let cipher = ChaCha20Poly1305::new(&key);

    let decrypted_msg = match cipher.decrypt(shared_nonce, chiphertext.as_ref()) {
        Ok(result) => result,
        Err(_) => return StatusCode::UNAUTHORIZED,
    };

    // if we doesn't have the signature, then revoke the method
    if decrypted_msg.len() < 64 {
        return StatusCode::UNAUTHORIZED;
    }

    let (signature_bytes, message_bytes) = decrypted_msg.split_at(64);
    let signature_fixed: [u8; 64] = match signature_bytes.try_into() {
        Ok(signature_fixed) => signature_fixed,
        Err(_) => return StatusCode::UNAUTHORIZED,
    };
    let signature = Signature::from_bytes(&signature_fixed);

    let filtered_bytes: Vec<u8> = message_bytes
        .iter()
        .copied()
        .filter(|&byte| byte != 0)
        .collect();
    let text = match String::from_utf8(filtered_bytes.clone()) {
        Ok(bytes) => bytes,
        Err(_) => return StatusCode::UNAUTHORIZED,
    };

    if client_verifying_key
        .verify(&filtered_bytes, &signature)
        .is_err()
    {
        return StatusCode::UNAUTHORIZED;
    }

    let new_message = match Message::new(alice_pk, signature_fixed, text) {
        Ok(result) => result,
        Err(_) => return StatusCode::UNAUTHORIZED,
    };
    let mut messages = match state.messages.lock() {
        Ok(messages) => messages,
        Err(_) => return StatusCode::UNAUTHORIZED,
    };
    if messages.len() >= state.max_per_day {
        return StatusCode::INSUFFICIENT_STORAGE;
    }

    messages.push(new_message.clone());
    valid_pubkey_for_ms.remove(&alice_pk);
    StatusCode::OK
}

pub async fn exchange_keys(
    State(state): State<AppState>,
    ValidatedForm(input): ValidatedForm<KeyExchange>,
) -> Json<KeyExchange> {
    let mut valids = match state.valid_pubkey_for_ms.lock() {
        Ok(guard) => guard,
        Err(_) => return Json(KeyExchange::default()),
    };
    update_senders_hashmap(&mut valids);

    let pubkey = match hex::decode(input.get_ed25519_pubkey()) {
        Ok(pubkey) => pubkey,
        Err(_) => return Json(KeyExchange::default()),
    };
    let nonce = match hex::decode(input.get_nonce()) {
        Ok(nonce) => nonce,
        Err(_) => return Json(KeyExchange::default()),
    };
    let pow = match hex::decode(input.get_pow()) {
        Ok(pow) => pow,
        Err(_) => return Json(KeyExchange::default()),
    };
    let pow_sig = match hex::decode(input.get_pow_sig()) {
        Ok(pow_sig) => pow_sig,
        Err(_) => return Json(KeyExchange::default()),
    };
    let diffie_hellman = match hex::decode(input.get_diffie_hellman()) {
        Ok(diffie_hellman) => diffie_hellman,
        Err(_) => return Json(KeyExchange::default()),
    };
    let diffie_hellman_sig = match hex::decode(input.get_diffie_hellman_sig()) {
        Ok(diffie_hellman_sig) => diffie_hellman_sig,
        Err(_) => return Json(KeyExchange::default()),
    };

    let alice_pk: [u8; 32] = match pubkey.try_into() {
        Ok(alice_pk) => alice_pk,
        Err(_) => return Json(KeyExchange::default()),
    };
    let client_verifying_key = match VerifyingKey::from_bytes(&alice_pk) {
        Ok(verifying_key) => verifying_key,
        Err(_) => return Json(KeyExchange::default()),
    };

    let nonce: [u8; 32] = match nonce.try_into() {
        Ok(nonce) => nonce,
        Err(_) => return Json(KeyExchange::default()),
    };

    let bob_pk = get_publickey_bytes(state.ed25519_bytes);
    let server_verifying_key = match VerifyingKey::from_bytes(&bob_pk) {
        Ok(verifying_key) => verifying_key,
        Err(_) => return Json(KeyExchange::default()),
    };

    let diffie_hellman_sig_bytes: [u8; 64] = match diffie_hellman_sig.try_into() {
        Ok(diffie_hellman_sig_bytes) => diffie_hellman_sig_bytes,
        Err(_) => return Json(KeyExchange::default()),
    };
    let diffie_hellman_sig = Signature::from_bytes(&diffie_hellman_sig_bytes);

    let pow_sig_bytes: [u8; 64] = match pow_sig.try_into() {
        Ok(pow_sig_bytes) => pow_sig_bytes,
        Err(_) => return Json(KeyExchange::default()),
    };
    let pow_sig = Signature::from_bytes(&pow_sig_bytes);

    let diffie_hellman_fixed: [u8; 32] = match diffie_hellman.clone().try_into() {
        Ok(diffie_hellman_fixed) => diffie_hellman_fixed,
        Err(_) => return Json(KeyExchange::default()),
    };
    let alice_diffie_hellman_pk = PublicKey::from(diffie_hellman_fixed);

    let mut dh_to_verify = Vec::new();
    dh_to_verify.extend_from_slice(&diffie_hellman);
    dh_to_verify.extend_from_slice(&nonce);

    let pow_fixed: [u8; 76] = match pow.clone().try_into() {
        Ok(pow) => pow,
        Err(_) => return Json(KeyExchange::default()),
    };

    let (pk_bytes, timestamp, bytes, cost) = deserialize_keys(&pow_fixed);

    let now = SystemTime::now();
    let duration_since_epoch = now
        .duration_since(UNIX_EPOCH)
        .unwrap_or(Duration::new(0, 0));
    let current_timestamp = duration_since_epoch.as_secs();

    if client_verifying_key
        .verify(&dh_to_verify, &diffie_hellman_sig)
        .is_err()
        || server_verifying_key.verify(&pow, &pow_sig).is_err()
        || pk_bytes != alice_pk
        || current_timestamp > timestamp
        || !verify(&bytes, nonce, cost)
    {
        return Json(KeyExchange::default());
    }

    let diffie_hellman_sk = ReusableSecret::random();
    let diffie_hellman = PublicKey::from(&diffie_hellman_sk).to_bytes();
    let hex_diffie_hellman = hex::encode(diffie_hellman);

    let diffie_hellman_sig = sign_message(&diffie_hellman, state.ed25519_bytes);
    let hex_diffie_hellman_sig = hex::encode(diffie_hellman_sig);

    let shared_secret = diffie_hellman_sk.diffie_hellman(&alice_diffie_hellman_pk);

    let fixed_12bytes: &[u8; 12] = shared_secret.as_bytes()[0..12]
        .try_into()
        .unwrap_or(&[0; 12]);
    let shared_nonce = chacha20poly1305::Nonce::from_slice(fixed_12bytes);

    let hex_nonce = hex::encode(nonce);

    let now = SystemTime::now();
    let duration_since_epoch = now
        .duration_since(UNIX_EPOCH)
        .unwrap_or(Duration::new(0, 0));
    let unix_timestamp_seconds = duration_since_epoch.as_secs();
    let delta = Duration::new(state.expected_time, 0);
    let end_timestamp = unix_timestamp_seconds + delta.as_secs();

    let keypair = Keypair::new(
        diffie_hellman_sk,
        alice_diffie_hellman_pk,
        *shared_nonce,
        hex_nonce,
        end_timestamp,
    );

    // this scenario doesnt look posible, but it's better
    // match every edge case :)
    if valids.len() == 50 {
        return Json(KeyExchange::default());
    }

    valids.insert(alice_pk, keypair);

    let hex_pubkey = hex::encode([0; 32]);
    let hex_null_nonce = hex::encode([0; 32]);
    let hex_null_pow = hex::encode([0; 32]);
    let hex_null_pow_sig = hex::encode([0; 80]);

    Json(KeyExchange::new(
        hex_pubkey,
        hex_null_nonce,
        hex_null_pow,
        hex_null_pow_sig,
        hex_diffie_hellman,
        hex_diffie_hellman_sig,
    ))
}

pub async fn get_domain_pk(State(state): State<AppState>) -> HexPubkey {
    let mut valid_pubkey_for_ms = match state.valid_pubkey_for_ms.lock() {
        Ok(guard) => guard,
        Err(_) => return hex::encode([0; 32]),
    };
    update_senders_hashmap(&mut valid_pubkey_for_ms);
    let bob_pubkey_bytes: [u8; 32] = get_publickey_bytes(state.ed25519_bytes);
    hex::encode(bob_pubkey_bytes)
}

pub async fn pow(
    State(state): State<AppState>,
    ValidatedForm(pubkey): ValidatedForm<RequestPubkey>,
) -> Json<PoW> {
    let mut valid_pubkey_for_ms = match state.valid_pubkey_for_ms.lock() {
        Ok(guard) => guard,
        Err(_) => return Json(PoW::default()),
    };
    update_senders_hashmap(&mut valid_pubkey_for_ms);

    let alice_public_bytes = match hex::decode(pubkey.get_bytes()) {
        Ok(bytes) => bytes,
        Err(_) => return Json(PoW::default()),
    };
    let alice_pk: [u8; 32] = match alice_public_bytes.try_into() {
        Ok(alice_pk) => alice_pk,
        Err(_) => return Json(PoW::default()),
    };

    let bytes_to_search: [u8; 32] = new_nonce();

    let now = SystemTime::now();
    let duration_since_epoch = now
        .duration_since(UNIX_EPOCH)
        .unwrap_or(Duration::new(0, 0));
    let unix_timestamp_seconds = duration_since_epoch.as_secs();
    let delta = Duration::new(state.expected_time, 0);
    let new_timestamp = unix_timestamp_seconds + delta.as_secs();

    let mut bytes = [].to_vec();
    bytes.extend(bytes_to_search);
    bytes.extend(state.cost_pow.to_le_bytes());
    bytes.extend(alice_pk);
    bytes.extend(new_timestamp.to_le_bytes());

    let signature = sign_message(&bytes, state.ed25519_bytes);

    let pow = PoW::new(hex::encode(bytes), hex::encode(signature));
    Json(pow)
}

pub fn update_senders_hashmap(valid_pubkeys: &mut HashMap<OnionPubkey, Keypair>) {
    let now = SystemTime::now();
    let duration_since_epoch = now
        .duration_since(UNIX_EPOCH)
        .unwrap_or(Duration::new(0, 0));
    let current_timestamp = duration_since_epoch.as_secs();

    let keys_to_remove: Vec<OnionPubkey> = valid_pubkeys
        .iter()
        .filter(|(_, keypair)| current_timestamp > keypair.get_unix_timestamp_seconds())
        .map(|(sender, _)| *sender)
        .collect();
    for sender in keys_to_remove {
        valid_pubkeys.remove(&sender);
    }
}