Documentation
// SPDX-FileCopyrightText: Heiko Schaefer <heiko@schaefer.name>
// SPDX-License-Identifier: MIT OR Apache-2.0

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

use minipgp6_lock::{argon2_small, lock};
use minipgp6_packet::{
    body::{
        secret_key::{SecretKey, SecretKeyMaterial},
        signature::{
            SignatureType,
            Subpacket,
            subpacket_data::{
                Features,
                IssuerFingerprint,
                KeyFlags,
                PreferredAeadCiphersuites,
                SignatureCreationTime,
            },
        },
        userid::UserId,
    },
    transferable::{SignedComponent, Tsk},
    types::{
        KeyRole,
        algorithm::{AeadAlgorithm, PublicKeyAlgorithm, SymmetricAlgorithm},
        time::Timestamp,
    },
    unlocked::UnlockedSoftwareKey,
};
use minipgp6_sign::component::{sign_direct_key, sign_key_binding, sign_user_id};
use rand::{CryptoRng, Rng, thread_rng};

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::regular(SignatureCreationTime(Timestamp::now()))?,
        Subpacket::regular(IssuerFingerprint(primary.public().fingerprint()))?,
        // CS (FIXME: don't hardcode literal)
        Subpacket::regular(KeyFlags(Box::new([0x03])))?,
        // SEIPD v2 only (FIXME: don't hardcode literal)
        Subpacket::regular(Features(Box::new([0x08])))?,
        // AES128, Ocb
        Subpacket::regular(PreferredAeadCiphersuites(Box::new([0x07, 0x02])))?,
    ];

    let config = minipgp6_sign::new_signature(
        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.clone(), primary.public());

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

    Ok(Tsk {
        primary: SignedComponent::<SecretKey> {
            component: primary,
            revocations: Vec::new(),
            signatures: 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, Timestamp::now(), KeyRole::Subkey)?;

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

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

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

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

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

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

            tsk.subkeys.push(SignedComponent::<SecretKey> {
                component: subkey,
                revocations: Vec::new(),
                signatures: 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::regular(SignatureCreationTime(Timestamp::now()))?,
        Subpacket::regular(IssuerFingerprint(
            tsk.primary.component.public().fingerprint(),
        ))?,
    ];

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

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

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

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

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

    tsk.user_ids.push(SignedComponent::<UserId> {
        component: uid,
        revocations: Vec::new(),
        signatures: 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.component, password)?;
        set_password(&mut tsk.subkeys[0].component, 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, Some(pw)) {
            key.secret_key_material = SecretKeyMaterial::Unprotected(decrypted.into_plain());

            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 unlocked = UnlockedSoftwareKey::new(plain.clone(), key.public());

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

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

    Ok(())
}