//! Encrypted identity backup. Onion (locator) is not included.

use argon2::{Algorithm, Argon2, Params, Version};
use chacha20poly1305::aead::{Aead, AeadCore, KeyInit, OsRng};
use chacha20poly1305::{ChaCha20Poly1305, Key, Nonce};
use rand::RngCore;

pub const MAGIC: &[u8] = b"owbak1";
const SALT_LEN: usize = 16;
const NONCE_LEN: usize = 12;
const KEY_LEN: usize = 32;
const PLAIN_LEN: usize = KEY_LEN * 4;

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct BackupKeys {
    pub identity_sk: [u8; 32],
    pub identity_pk: [u8; 32],
    pub prekey_sk: [u8; 32],
    pub prekey_pk: [u8; 32],
}

impl BackupKeys {
    pub fn to_bytes(&self) -> [u8; PLAIN_LEN] {
        let mut out = [0u8; PLAIN_LEN];
        out[0..32].copy_from_slice(&self.identity_sk);
        out[32..64].copy_from_slice(&self.identity_pk);
        out[64..96].copy_from_slice(&self.prekey_sk);
        out[96..128].copy_from_slice(&self.prekey_pk);
        out
    }

    pub fn from_bytes(b: &[u8]) -> Result<Self, String> {
        if b.len() != PLAIN_LEN {
            return Err("backup plaintext length".into());
        }
        Ok(Self {
            identity_sk: b[0..32].try_into().unwrap(),
            identity_pk: b[32..64].try_into().unwrap(),
            prekey_sk: b[64..96].try_into().unwrap(),
            prekey_pk: b[96..128].try_into().unwrap(),
        })
    }
}

pub(crate) fn kdf(passphrase: &str, salt: &[u8]) -> Result<[u8; 32], String> {
    if passphrase.is_empty() {
        return Err("empty passphrase".into());
    }
    if salt.len() != SALT_LEN {
        return Err("salt length".into());
    }
    let params = Params::new(19_456, 2, 1, Some(32)).map_err(|e| e.to_string())?;
    let argon = Argon2::new(Algorithm::Argon2id, Version::V0x13, params);
    let mut key = [0u8; 32];
    argon
        .hash_password_into(passphrase.as_bytes(), salt, &mut key)
        .map_err(|e| e.to_string())?;
    Ok(key)
}

pub fn seal(passphrase: &str, keys: &BackupKeys) -> Result<Vec<u8>, String> {
    let mut salt = [0u8; SALT_LEN];
    rand::rngs::OsRng.fill_bytes(&mut salt);
    let key = kdf(passphrase, &salt)?;
    let cipher = ChaCha20Poly1305::new(Key::from_slice(&key));
    let nonce = ChaCha20Poly1305::generate_nonce(&mut OsRng);
    if nonce.len() != NONCE_LEN {
        return Err("nonce length".into());
    }
    let ct = cipher
        .encrypt(&nonce, keys.to_bytes().as_ref())
        .map_err(|_| "encrypt failed".to_string())?;
    let mut out = Vec::with_capacity(MAGIC.len() + SALT_LEN + NONCE_LEN + ct.len());
    out.extend_from_slice(MAGIC);
    out.extend_from_slice(&salt);
    out.extend_from_slice(&nonce);
    out.extend_from_slice(&ct);
    Ok(out)
}

pub fn open(passphrase: &str, blob: &[u8]) -> Result<BackupKeys, String> {
    let min = MAGIC.len() + SALT_LEN + NONCE_LEN + 16;
    if blob.len() < min || !blob.starts_with(MAGIC) {
        return Err("not an onionwire backup".into());
    }
    let salt = &blob[MAGIC.len()..MAGIC.len() + SALT_LEN];
    let nonce_off = MAGIC.len() + SALT_LEN;
    let nonce = Nonce::from_slice(&blob[nonce_off..nonce_off + NONCE_LEN]);
    let ct = &blob[nonce_off + NONCE_LEN..];
    let key = kdf(passphrase, salt)?;
    let cipher = ChaCha20Poly1305::new(Key::from_slice(&key));
    let pt = cipher
        .decrypt(nonce, ct)
        .map_err(|_| "wrong passphrase or corrupt backup".to_string())?;
    BackupKeys::from_bytes(&pt)
}

pub(crate) fn aead_encrypt(key: &[u8; 32], plaintext: &[u8]) -> Result<Vec<u8>, String> {
    let cipher = ChaCha20Poly1305::new(Key::from_slice(key));
    let nonce = ChaCha20Poly1305::generate_nonce(&mut OsRng);
    let ct = cipher
        .encrypt(&nonce, plaintext)
        .map_err(|_| "encrypt failed".to_string())?;
    let mut out = Vec::with_capacity(NONCE_LEN + ct.len());
    out.extend_from_slice(&nonce);
    out.extend_from_slice(&ct);
    Ok(out)
}

pub(crate) fn aead_decrypt(key: &[u8; 32], blob: &[u8]) -> Result<Vec<u8>, String> {
    if blob.len() < NONCE_LEN + 16 {
        return Err("ciphertext length".into());
    }
    let nonce = Nonce::from_slice(&blob[..NONCE_LEN]);
    let cipher = ChaCha20Poly1305::new(Key::from_slice(key));
    cipher
        .decrypt(nonce, &blob[NONCE_LEN..])
        .map_err(|_| "wrong passphrase or corrupt".to_string())
}