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

//! Wrapper types and logic around OpenPGP private keys, aka TSKs ("transferable secret keys").

use std::io;

use chrono::{DateTime, SubsecRound, Utc};
use pgp::cleartext::CleartextSignedMessage;
use pgp::crypto::hash::HashAlgorithm;
use pgp::packet::{SignatureConfig, SignatureType, Subpacket, SubpacketData};
use pgp::ser::Serialize;
use pgp::types::{Fingerprint, KeyVersion, PublicKeyTrait};
use pgp::{
    ArmorOptions, KeyType, Message, PublicOrSecret, SecretKeyParamsBuilder, SignedPublicKey,
    SignedSecretKey, SubkeyParamsBuilder,
};
use rand::thread_rng;

use crate::certificate::{Certificate, Checked};
use crate::key::{ComponentKeySec, KeyFlagMatch, SignedComponentKey, SignedComponentKeySec};
use crate::policy::{
    PREFERRED_AEAD_ALGORITHMS, PREFERRED_COMPRESSION_ALGORITHMS, PREFERRED_HASH_ALGORITHMS,
    PREFERRED_SYMMETRIC_KEY_ALGORITHMS,
};
use crate::signature::SigStack;
use crate::{ComponentKeyPub, Error};

/// A "transferable secret key (TSK)," also known as an "OpenPGP private/secret key."
#[derive(Clone)]
pub struct Tsk {
    ssk: SignedSecretKey,
    checked: Checked, // keep a "checked" version of the ssk, to find valid component keys on
}

impl TryFrom<&[u8]> for Tsk {
    type Error = Error;

    fn try_from(input: &[u8]) -> Result<Self, Self::Error> {
        use pgp::Deserializable;

        let (ssk, _) = SignedSecretKey::from_reader_single(input)?;
        Ok(ssk.into())
    }
}

impl From<SignedSecretKey> for Tsk {
    fn from(ssk: SignedSecretKey) -> Self {
        let spk = SignedPublicKey::from(ssk.clone());
        let checked = (&Certificate::from(spk)).into();

        Self { ssk, checked }
    }
}

/// A signer that's based on a private component key that is marked as capable of producing data
/// signatures.
pub struct DataSigner {
    key: ComponentKeySec,
}

impl DataSigner {
    /// Takes a message as input, and produces a signed message
    pub fn sign_msg<F>(
        &self,
        msg: Message,
        key_pw: F,
        hash_algo: HashAlgorithm,
    ) -> Result<Message, Error>
    where
        F: FnOnce() -> String,
    {
        self.key.sign_msg(msg, key_pw, hash_algo)
    }

    /// Takes a message as input, and produces a cleartext signed message
    pub fn sign_csf<F>(&self, body: &str, key_pw: F) -> Result<CleartextSignedMessage, Error>
    where
        F: FnOnce() -> String,
    {
        match &self.key {
            // FIXME: DRY?
            ComponentKeySec::Primary(sk) => Ok(CleartextSignedMessage::sign(
                thread_rng(),
                body,
                &sk,
                key_pw,
            )?),
            ComponentKeySec::Subkey(ssk) => Ok(CleartextSignedMessage::sign(
                thread_rng(),
                body,
                &ssk,
                key_pw,
            )?),
        }
    }

    /// The fingerprint of the underlying component private key
    pub fn fingerprint(&self) -> Fingerprint {
        self.key.fingerprint()
    }
}

impl Tsk {
    /// Load a set of `Tsk`s from a source.
    ///
    /// The source data may be armored or binary.
    pub fn load<R: io::Read>(read: &mut R) -> Result<Vec<Tsk>, Error> {
        let (keys, _headers) = pgp::composed::signed_key::from_reader_many(read)?;

        let mut tsk = vec![];

        for res in keys {
            match res {
                Ok(pos) => match pos {
                    PublicOrSecret::Public(spk) => {
                        eprintln!("Bad data {:?}", spk);
                    }
                    PublicOrSecret::Secret(ssk) => {
                        tsk.push(ssk.into());
                    }
                },
                Err(_) => eprintln!("Bad data {:?}", res),
            }
        }

        if tsk.is_empty() {
            Err(Error::Message("No TSKs found".to_string()))
        } else {
            Ok(tsk)
        }
    }

    /// Save this Tsk to a writer
    pub fn save(&self, armored: bool, sink: &mut dyn io::Write) -> Result<(), Error> {
        Self::save_all([self], armored, sink)
    }

    /// Save a set of Certificates to a writer
    pub fn save_all<'a>(
        tsks: impl IntoIterator<Item = &'a Self>,
        armored: bool,
        mut sink: &mut dyn io::Write,
    ) -> Result<(), Error> {
        if armored {
            let spks: Vec<_> = tsks.into_iter().map(|tsk| &tsk.ssk).collect();

            let opts = ArmorOptions::default();

            pgp::armor::write(
                &spks,
                pgp::armor::BlockType::PrivateKey,
                &mut sink,
                opts.headers,
                opts.include_checksum,
            )?;
        } else {
            for tsk in tsks {
                tsk.ssk.to_writer(&mut sink)?;
            }
        }

        Ok(())
    }

    /// FIXME: remove this from the API again?
    pub fn key(&self) -> &SignedSecretKey {
        &self.ssk
    }

    pub fn generate_v4(
        key_type_primary: KeyType,
        key_type_encrypt: impl Into<Option<KeyType>>,
        primary_user_id: Option<String>,
        secondary_user_ids: Vec<String>,
        key_password: Option<&str>,
    ) -> Result<Self, Error> {
        let Some(primary_user_id) = primary_user_id else {
            return Err(Error::Rpgp(pgp::errors::Error::Message(
                "Generating UID-less keys not supported".to_string(),
            )));
        };

        let mut rng = thread_rng();

        let subkeys = if let Some(key_type_encrypt) = key_type_encrypt.into() {
            vec![SubkeyParamsBuilder::default()
                .key_type(key_type_encrypt)
                .can_encrypt(true)
                .build()?]
        } else {
            vec![]
        };

        let mut key_params = SecretKeyParamsBuilder::default();
        key_params
            .key_type(key_type_primary)
            .can_certify(true)
            .can_sign(true)
            .primary_user_id(primary_user_id)
            .user_ids(secondary_user_ids)
            .preferred_symmetric_algorithms(PREFERRED_SYMMETRIC_KEY_ALGORITHMS.into())
            .preferred_hash_algorithms(PREFERRED_HASH_ALGORITHMS.into())
            .preferred_compression_algorithms(PREFERRED_COMPRESSION_ALGORITHMS.into())
            .subkeys(subkeys);

        let secret_key_params = key_params.build()?;
        let secret_key = secret_key_params.generate(&mut rng)?;

        let passwd_fn = String::new;

        let mut signed_secret_key = secret_key.sign(&mut thread_rng(), passwd_fn)?;
        if let Some(key_password) = key_password {
            signed_secret_key
                .primary_key
                .set_password(&mut rng, || key_password.to_string())?;

            for sk in &mut signed_secret_key.secret_subkeys {
                sk.key.set_password(&mut rng, || key_password.to_string())?;
            }
        };

        Ok(signed_secret_key.into())
    }

    pub fn generate_v6(
        key_type_primary: KeyType,
        key_type_encrypt: KeyType,
        primary_user_id: Option<String>,
        secondary_user_ids: Vec<String>,
        key_password: Option<&str>,
    ) -> Result<Self, Error> {
        let mut rng = thread_rng();

        let (primary_user_id, uidless) = match primary_user_id {
            Some(uid) => (uid, false),
            None => (String::default(), true),
        };

        let mut key_params = SecretKeyParamsBuilder::default();
        key_params
            .version(KeyVersion::V6)
            .key_type(key_type_primary)
            .can_certify(true)
            .can_sign(true)
            .primary_user_id(primary_user_id)
            .user_ids(secondary_user_ids)
            .preferred_symmetric_algorithms(PREFERRED_SYMMETRIC_KEY_ALGORITHMS.into())
            .preferred_hash_algorithms(PREFERRED_HASH_ALGORITHMS.into())
            .preferred_compression_algorithms(PREFERRED_COMPRESSION_ALGORITHMS.into())
            .preferred_aead_algorithms(PREFERRED_AEAD_ALGORITHMS.into())
            .subkeys(vec![SubkeyParamsBuilder::default()
                .version(KeyVersion::V6)
                .key_type(key_type_encrypt)
                .can_encrypt(true)
                .build()?]);

        let secret_key_params = key_params.build()?;
        let secret_key = secret_key_params.generate(&mut rng)?;

        let passwd_fn = String::new;
        let mut signed_secret_key = secret_key.sign(&mut thread_rng(), passwd_fn)?;

        // drop userid if it's just a placeholder
        if uidless {
            signed_secret_key.details.users = vec![];
        }

        // Add a DKS for v6 keys
        let hash_alg = HashAlgorithm::SHA2_512; // FIXME?

        let mut config = SignatureConfig::v6(
            &mut thread_rng(),
            SignatureType::Key,
            signed_secret_key.algorithm(),
            hash_alg,
        )?;
        config.hashed_subpackets = vec![
            Subpacket::regular(SubpacketData::SignatureCreationTime(
                chrono::Utc::now().trunc_subsecs(0),
            )),
            Subpacket::regular(SubpacketData::KeyFlags([3].into())), // FIXME, set CS
            Subpacket::regular(SubpacketData::Features([0x01 | 0x08].into())), // FIXME, set SEIPDv2 + SEIPDv1
            Subpacket::regular(SubpacketData::PreferredSymmetricAlgorithms(
                PREFERRED_SYMMETRIC_KEY_ALGORITHMS.into(),
            )),
            Subpacket::regular(SubpacketData::PreferredHashAlgorithms(
                PREFERRED_HASH_ALGORITHMS.into(),
            )),
            Subpacket::regular(SubpacketData::PreferredCompressionAlgorithms(
                PREFERRED_COMPRESSION_ALGORITHMS.into(),
            )),
            Subpacket::regular(SubpacketData::IssuerFingerprint(
                signed_secret_key.fingerprint(),
            )),
        ];

        let passwd_fn = String::new;
        let dks = config.sign_key(&signed_secret_key, passwd_fn, &signed_secret_key)?;

        signed_secret_key.details.direct_signatures.push(dks);

        if let Some(key_password) = key_password {
            signed_secret_key
                .primary_key
                .set_password(&mut rng, || key_password.to_string())?;

            for sk in &mut signed_secret_key.secret_subkeys {
                sk.key.set_password(&mut rng, || key_password.to_string())?;
            }
        };

        Ok(signed_secret_key.into())
    }

    fn component_keys(&self) -> Vec<SignedComponentKey> {
        let x = SignedComponentKeySec::Primary(self.ssk.clone());
        let mut v = vec![SignedComponentKey::Sec(x)];

        self.ssk.secret_subkeys.iter().for_each(|sssk| {
            let dks = SigStack::from_iter(self.ssk.details.direct_signatures.iter()).active();

            let x = SignedComponentKeySec::Subkey((sssk.clone(), dks.cloned()));
            v.push(SignedComponentKey::Sec(x));
        });

        v
    }

    fn get_matching_secret_key(&self, ckp: &ComponentKeyPub) -> Option<ComponentKeySec> {
        match ckp {
            ComponentKeyPub::Primary(pri) => {
                let sk = &self.ssk.primary_key;

                if pri.fingerprint() == sk.fingerprint() {
                    // FIXME: compare more efficiently?
                    Some(ComponentKeySec::Primary(sk.clone()))
                } else {
                    None
                }
            }

            ComponentKeyPub::Subkey(sub) => {
                for sssk in &self.ssk.secret_subkeys {
                    // FIXME: compare more efficiently?
                    if sub.fingerprint() == sssk.fingerprint() {
                        return Some(ComponentKeySec::Subkey(sssk.key.clone()));
                    }
                }

                // Found no match
                None
            }
        }
    }

    /// Get list of all valid signing capable component keys
    /// (this fn is specific to *signing*, not validation: it uses "now" as its reference time and is thus stricter)
    pub fn signing_capable_component_keys(&self) -> impl Iterator<Item = DataSigner> + '_ {
        let now = Utc::now();
        let sv = self.checked.valid_signing_capable_component_keys_at(&now);

        let mut ds = vec![];

        for ckp in sv.iter().map(|sv| sv.as_componentkey()) {
            // Find associated secret key packet
            let key = self.get_matching_secret_key(ckp);

            if let Some(key) = key {
                ds.push(DataSigner { key });
            } else {
                log::warn!(
                    "signing_capable_component_keys: failed to find secret key packet for {:?}",
                    ckp.fingerprint()
                );
            }
        }

        ds.into_iter()
    }

    /// Get list of all decryption capable component keys
    /// (this fn is specific to *decryption*, not encryption: it is very lenient in allowing use of keys)
    pub fn decryption_capable_component_keys(
        &self,
    ) -> impl Iterator<Item = SignedComponentKey> + '_ {
        let now: DateTime<Utc> = chrono::offset::Utc::now();

        // FIXME: filter out unknown notations

        self.component_keys()
            .into_iter()
            .filter(move |sck| sck.is_encryption_capable(&now))
    }

    fn sec_components(&self) -> Vec<SignedComponentKeySec> {
        let x = SignedComponentKeySec::Primary(self.ssk.clone());
        let mut v = vec![x];

        self.ssk.secret_subkeys.iter().for_each(|sssk| {
            let dks = SigStack::from_iter(self.ssk.details.direct_signatures.iter()).active();

            let x = SignedComponentKeySec::Subkey((sssk.clone(), dks.cloned()));
            v.push(x);
        });

        v
    }

    /// Get all component secret keys that have a "signing capable" key flag.
    ///
    /// NOTE: this function doesn't cryptographically verify self-signatures, it's not safe to use
    /// on potentially attacker-controlled Tsks.
    pub fn signing_keys_sec(&self) -> Vec<ComponentKeySec> {
        // FIXME: filter out revoked subkeys

        self.sec_components()
            .iter()
            .filter(|x| x.has_key_flag(KeyFlagMatch::Sign))
            .map(Into::into)
            .collect()
    }

    /// Get all component secret keys that have an "encryption capable" key flag.
    ///
    /// NOTE: this function doesn't cryptographically verify self-signatures, it's not safe to use
    /// on potentially attacker-controlled Tsks.
    pub fn decryption_keys_sec(&self) -> Vec<ComponentKeySec> {
        // This should actually not filter out revoked subkeys - we may still want to decrypt with them!

        self.sec_components()
            .iter()
            .filter(|x| x.has_key_flag(KeyFlagMatch::Enc))
            .map(Into::into)
            .collect()
    }

    /// Get all component secret keys that have an "authentication capable" key flag.
    ///
    /// NOTE: this function doesn't cryptographically verify self-signatures, it's not safe to use
    /// on potentially attacker-controlled Tsks.
    pub fn auth_keys_sec(&self) -> Vec<ComponentKeySec> {
        // FIXME: filter out revoked subkeys

        self.sec_components()
            .iter()
            .filter(|x| x.has_key_flag(KeyFlagMatch::Auth))
            .map(Into::into)
            .collect()
    }
}
