Documentation
//! Generate OpenPGP Transferable Secret Keys
//!
//! <https://www.rfc-editor.org/rfc/rfc9580#name-transferable-secret-keys>

use minipgp6_lock::{UnlockedSoftwareKey, argon2_small, lock};
use minipgp6_packet::{
    packet::{
        secret_key::{SecretKey, SecretKeyMaterial},
        signature::{SignatureType, subpacket::Subpacket},
        userid::UserId,
    },
    transferable::{SignedPrimary, SignedSubkey, SignedUserId, Tsk},
    types::{AeadAlgorithm, KeyRole, PublicKeyAlgorithm, SymmetricAlgorithm, Timestamp},
};
use minipgp6_sign::component::{sign_direct_key, sign_subkey_binding, sign_user_id};
use rand::{CryptoRng, Rng, thread_rng};
use secrecy::{ExposeSecret, SecretBox};

use crate::Error;

/// Specify what capability a subkey should have
#[derive(Debug, Clone)]
pub enum KeyUsage {
    /// Subkey for encryption
    Encryption,

    /// Subkey for signing data
    SignData,

    /// Subkey for authentication (usually SSH)
    Auth,
}

/// Create a new minimal Tsk, consisting of only a primary key and a direct key signature on it
pub fn generate_minimal<R: Rng + CryptoRng>(
    mut rng: R,
    primary_algo: PublicKeyAlgorithm,
) -> Result<Tsk, Error> {
    let primary = minipgp6_sign::generate_key(&mut rng, primary_algo, KeyRole::Primary)?;

    let hashed_area = vec![
        Subpacket::SignatureCreationTime(Timestamp::now()).into(),
        Subpacket::IssuerFingerprint(primary.public().fingerprint()).into(),
        Subpacket::KeyFlags(vec![0x03]).into(), // CS (FIXME: don't hardcode literal)
        Subpacket::Features(vec![0x08]).into(), // SEIPD v2 only (FIXME: don't hardcode literal)
        Subpacket::PreferredAeadCiphersuites(vec![0x07, 0x02]).into(), // AES128, Ocb
    ];

    let config = minipgp6_sign::config::new(
        rng,
        SignatureType::DirectKey,
        primary_algo,
        primary_algo.hash_algorithm()?,
        hashed_area,
        Vec::new(),
    );

    let SecretKeyMaterial::Unprotected(plain) = &primary.secret_key_material else {
        return Err(Error::LockingUnsupported);
    };

    let signer = UnlockedSoftwareKey::new(plain.expose_secret().clone(), primary.public());

    let dks = sign_direct_key(config, &signer, primary.public())?;

    Ok(Tsk {
        primary: SignedPrimary {
            key: primary,
            revocations: Vec::new(),
            direct: vec![dks],
        },
        subkeys: Vec::new(),
        user_ids: Vec::new(),
    })
}

/// Add a subkey to `tsk`, and create the required binding self-signature for it
pub fn add_subkey<R: Rng + CryptoRng>(
    tsk: &mut Tsk,
    mut rng: R,
    algo: PublicKeyAlgorithm,
    usage: KeyUsage,
) -> Result<(), Error> {
    match usage {
        KeyUsage::Encryption => {
            let subkey = minipgp6_encrypt::generate_key(&mut rng, algo, KeyRole::Subkey)?;

            let hashed_area = vec![
                Subpacket::SignatureCreationTime(Timestamp::now()).into(),
                Subpacket::IssuerFingerprint(tsk.primary.key.public().fingerprint()).into(),
                Subpacket::KeyFlags(vec![0x04 | 0x08]).into(), // E
            ];

            let primary_pk = tsk.primary.key.public().public_key_algorithm;

            let config = minipgp6_sign::config::new(
                rng,
                SignatureType::SubkeyBinding,
                primary_pk,
                primary_pk.hash_algorithm()?,
                hashed_area,
                Vec::new(),
            );

            let SecretKeyMaterial::Unprotected(plain) = &tsk.primary.key.secret_key_material else {
                return Err(Error::LockingUnsupported);
            };

            let signer =
                UnlockedSoftwareKey::new(plain.expose_secret().clone(), tsk.primary.key.public());

            let binding =
                sign_subkey_binding(config, &signer, tsk.primary.key.public(), subkey.public())?;

            tsk.subkeys.push(SignedSubkey {
                key: subkey,
                revocations: Vec::new(),
                bindings: vec![binding],
            });

            Ok(())
        }
        _ => Err(Error::Message(format!("unsupported usage {:?}", usage))),
    }
}

/// Add a subkey to `tsk`, and create the required certifying self-signature for it
pub fn add_userid<R: Rng + CryptoRng>(tsk: &mut Tsk, rng: R, user_id: String) -> Result<(), Error> {
    let uid = UserId::new(user_id)?;

    let hashed_area = vec![
        Subpacket::SignatureCreationTime(Timestamp::now()).into(),
        Subpacket::IssuerFingerprint(tsk.primary.key.public().fingerprint()).into(),
    ];

    let primary_pk = tsk.primary.key.public().public_key_algorithm;

    let config = minipgp6_sign::config::new(
        rng,
        SignatureType::PositiveCertification,
        primary_pk,
        primary_pk.hash_algorithm()?,
        hashed_area,
        Vec::new(),
    );

    let SecretKeyMaterial::Unprotected(plain) = &tsk.primary.key.secret_key_material else {
        return Err(Error::LockingUnsupported);
    };

    let signer = UnlockedSoftwareKey::new(plain.expose_secret().clone(), tsk.primary.key.public());

    let certification = sign_user_id(config, &signer, tsk.primary.key.public(), &uid)?;

    tsk.user_ids.push(SignedUserId {
        user_id: uid,
        revocations: Vec::new(),
        certifications: vec![certification],
    });

    Ok(())
}

/// Generate a new Tsk
///
/// TODO: generalize
pub fn generate_tsk<RNG: Rng + CryptoRng>(
    mut rng: RNG,
    primary_type: PublicKeyAlgorithm,
    encryption_type: Option<PublicKeyAlgorithm>,
    user_ids: &[impl AsRef<str>],
    key_password: Option<&[u8]>,
) -> Result<Tsk, Error> {
    // Generate primary key
    let mut tsk = generate_minimal(&mut rng, primary_type)?;

    // Add encryption subkey (if not signing_only)
    if let Some(encryption_type) = encryption_type {
        add_subkey(&mut tsk, &mut rng, encryption_type, KeyUsage::Encryption)?;
    }

    // add user ids
    for user_id in user_ids {
        add_userid(&mut tsk, &mut rng, user_id.as_ref().to_string())?;
    }

    // Set password for primary and encryption subkey
    // FIXME: generalize
    if let Some(password) = key_password {
        set_password(&mut tsk.primary.key, password)?;
        set_password(&mut tsk.subkeys[0].key, password)?;
    }

    Ok(tsk)
}

/// Remove the password locking from `key`, if possible
pub fn remove_password(key: &mut SecretKey, pw: &[u8]) -> Result<(), Error> {
    remove_password_with_list(key, &[pw])
}

/// Remove the password locking from `key`, if possible, using a list of passwords
pub fn remove_password_with_list(key: &mut SecretKey, pws: &[&[u8]]) -> Result<(), Error> {
    if matches!(key.secret_key_material, SecretKeyMaterial::Unprotected(_)) {
        return Ok(()); // no need to unlock
    }

    for pw in pws {
        if let Ok(decrypted) = minipgp6_lock::unlock(key, pw) {
            key.secret_key_material =
                SecretKeyMaterial::Unprotected(SecretBox::new(decrypted.into()));

            return Ok(());
        }
    }

    Err(Error::CantUnlock)
}

/// Lock `key` with `password`
pub fn set_password(key: &mut SecretKey, password: &[u8]) -> Result<(), Error> {
    let SecretKeyMaterial::Unprotected(plain) = &key.secret_key_material else {
        return Err(Error::Message(
            "Secret key packet is already locked".to_string(),
        ));
    };

    let mut rng = thread_rng();

    let s2k = argon2_small(&mut rng);
    let locked = lock(
        &mut rng,
        plain.expose_secret().clone(),
        password,
        SymmetricAlgorithm::AES256,
        AeadAlgorithm::Ocb,
        s2k,
        key.public(),
    )
    .map_err(|e| Error::Message(format!("Failed to lock secret key material: {}", e)))?;

    key.secret_key_material = SecretKeyMaterial::Aead(locked);

    Ok(())
}