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

//! Handling of OpenPGP messages.

use std::io;
use std::io::Read;

use pgp::crypto::aead::AeadAlgorithm;
use pgp::crypto::hash::HashAlgorithm;
use pgp::crypto::sym::SymmetricKeyAlgorithm;
use pgp::packet::{
    LiteralData, Packet, PublicKeyEncryptedSessionKey, SymEncryptedProtectedData,
    SymKeyEncryptedSessionKey,
};
use pgp::ser::Serialize;
use pgp::types::StringToKey;
use pgp::{decrypt_session_key_with_password, ArmorOptions, PlainSessionKey};
use pgp::{Deserializable, Edata, Esk};
use pgp::{Message, Signature};
use zeroize::Zeroizing;

use crate::certificate::Certificate;
use crate::certificate::Checked;
use crate::key::{ComponentKeyPub, SignedComponentKey};
use crate::policy::MAX_RECURSION;
use crate::tsk::Tsk;
use crate::Error;

/// Any mechanism that can decrypt a PKESK into a PlainSessionKey.
/// This abstraction allows use of either software-backed or hardware-backed private keys.
pub trait PkeskDecryptor {
    fn decrypt(&self, pkesk: &PublicKeyEncryptedSessionKey) -> Option<PlainSessionKey>;
}

/// A decryptor backed by a "softkey" [Tsk] (as opposed to hardware-backed private key material).
///
/// A SoftkeyPkeskDecryptor can be used to decrypt PKESKs and obtain decrypted session keys.
pub struct SoftkeyPkeskDecryptor<'a> {
    tsk: Tsk,
    key_passwords: Vec<&'a [u8]>,
}

impl<'a> SoftkeyPkeskDecryptor<'a> {
    pub fn new(tsk: Tsk, key_passwords: Vec<&'a [u8]>) -> Self {
        Self { tsk, key_passwords }
    }
}

impl PkeskDecryptor for SoftkeyPkeskDecryptor<'_> {
    fn decrypt(&self, pkesk: &PublicKeyEncryptedSessionKey) -> Option<PlainSessionKey> {
        for key in self.tsk.decryption_capable_component_keys() {
            let key_pw = match self.key_passwords.len() {
                0 => "".to_string(),
                1 => String::from_utf8_lossy(self.key_passwords[0]).to_string(),
                _ => {
                    // FIXME
                    return None;
                }
            };

            // Match key id/fp or check for wildcard.
            if pkesk.match_identity(&key) {
                match key {
                    SignedComponentKey::Sec(s) => {
                        match s.decrypt_session_key(pkesk, || key_pw.clone()) {
                            Ok(sk0) => {
                                return Some(sk0);
                            }
                            Err(e) => eprintln!("err {:?}", e),
                        }
                    }
                    SignedComponentKey::Pub(_) => {
                        // we can't decrypt with this key
                    }
                }
            };
        }

        None
    }
}

/// Result of decrypting and/or verifying the signatures of a message with [unpack].
///
/// Depending on the format of the processed message, this result contains:
/// - The cleartext of the message.
/// - The session key, if the message was encrypted.
/// - A list of signatures that were found to be valid.
pub struct MessageResult {
    pub validated: Vec<(Certificate, ComponentKeyPub, Signature)>,
    pub session_key: Option<(u8, Vec<u8>)>,
    pub cleartext: LiteralData,
}

/// Process an existing message: Decrypt message, and/or verify signatures.
///
/// This function handles messages that are encrypted, signed or both.
///
/// More specifically: OpenPGP messages may consist of layers of encryption, signing and
/// compression. This function processes arbitrary combinations of these layers, up to a maximum
/// layering depth.
///
/// The return value encodes (to some degree) the properties the message:
///
/// - It returns the cleartext of the message, and
/// - A list of the signatures on the message that were found to be valid (if any).
///
/// **Decryption**
///
/// Decryption can be attempted using two distinct mechanisms:
///
/// 1. Using a private OpenPGP component key provided via `decryptor`.
///    The OpenPGP component key may optionally be protected with a passphrase. Unlocking such
///    protected keys will be attempted with the passphrases provided in `key_passwords` (if any).
///
/// 2. A symmetric key, represented by a passphrase, in `skesk_passwords` (OpenPGP messages can be
///    encrypted to a recipient who is not using an OpenPGP key, with this method).
///
/// **Signature verification**
///
/// If the message has been signed, signatures will be verified against the certificates in `verifier`.
/// Any correct signatures will be reported in the `MessageResult`.
///
/// **Compression layers**
///
/// If the message has compression layers, they will be unpacked, silently.
///
/// FIXME: This function should take `msg: &Message`
pub fn unpack(
    msg: Message,
    pkesk_decryptors: &[Box<dyn PkeskDecryptor + '_>],
    skesk_passwords: Vec<&[u8]>,
    verifier: &[Certificate],
) -> Result<MessageResult, Error> {
    unpack_int(&msg, pkesk_decryptors, skesk_passwords, None, verifier, 0)
}

/// Try to decrypt msg with all known session key versions
pub fn unpack_by_session_key(
    msg: &Message,
    key: &[u8],
    sym_alg: SymmetricKeyAlgorithm,
    verifier: &[Certificate],
) -> Result<MessageResult, Error> {
    // Try to decrypt with a v3/4 session key
    let sk_34 = PlainSessionKey::V3_4 {
        sym_alg,
        key: key.into(),
    };
    if let Ok(mr) = unpack_int(msg, &[], vec![], Some(sk_34), verifier, 0) {
        return Ok(mr);
    }

    // Try to decrypt with a v6 session key
    let sk_6 = PlainSessionKey::V6 { key: key.into() };
    if let Ok(mr) = unpack_int(msg, &[], vec![], Some(sk_6), verifier, 0) {
        return Ok(mr);
    }

    // All attempts to decrypt failed
    Err(Error::Message("Failed to decrypt by session key".into()))
}

// internal variant of the unpack() function with depth handling and limitation
fn unpack_int(
    msg: &Message,
    pkesk_decryptors: &[Box<dyn PkeskDecryptor + '_>],
    skesk_passwords: Vec<&[u8]>,
    session_key: Option<PlainSessionKey>,
    verifier: &[Certificate],
    depth: usize,
) -> Result<MessageResult, Error> {
    if depth > MAX_RECURSION {
        return Err(Error::Message(
            "Excessive message nesting depth".to_string(),
        ));
    }

    match &msg {
        Message::Encrypted { ref esk, ref edata } => {
            let mut sk = session_key.clone(); // Session key, if any is found

            if sk.is_none() {
                'esks: for e in esk {
                    match e {
                        Esk::PublicKeyEncryptedSessionKey(pkesk) => {
                            for dec in pkesk_decryptors {
                                if let Some(esk) = dec.decrypt(pkesk) {
                                    sk = Some(esk);
                                    break 'esks;
                                }
                            }
                        }
                        Esk::SymKeyEncryptedSessionKey(skesk) => {
                            if !skesk_passwords.is_empty() {
                                let sym_pw = match skesk_passwords.len() {
                                    0 => "".to_string(),
                                    1 => String::from_utf8_lossy(skesk_passwords[0]).to_string(),
                                    _ => {
                                        // FIXME
                                        return Err(Error::Message(
                                            "More than one SKESK password currently unsupported"
                                                .to_string(),
                                        ));
                                    }
                                };

                                if let Ok(sk0) = decrypt_session_key_with_password(skesk, || sym_pw)
                                {
                                    sk = Some(sk0);
                                    break 'esks;
                                }
                            }
                        }
                    }
                }
            }

            let Some(sk) = sk else {
                // FIXME: return more specific error type
                return Err(Error::Message("Failed to get session key".to_string()));
            };

            let (sym_alg, key) = match &sk {
                PlainSessionKey::V3_4 { sym_alg, key } => (*sym_alg, key.to_vec()),

                PlainSessionKey::V6 { key } => match edata {
                    Edata::SymEncryptedProtectedData(seipd) => match &seipd.data() {
                        pgp::packet::Data::V2 { sym_alg, .. } => (*sym_alg, key.to_vec()),
                        pgp::packet::Data::V1 { .. } => {
                            return Err(Error::Rpgp(pgp::errors::Error::Message(format!(
                                "V6 session key with unexpected symmetric data version {}",
                                seipd.version()
                            ))));
                        }
                    },
                    Edata::SymEncryptedData(_) => {
                        return Err(Error::Rpgp(pgp::errors::Error::Message(
                            "SED packet is not supported with V6 session keys".to_string(),
                        )));
                    }
                },

                psk => {
                    return Err(Error::Rpgp(pgp::errors::Error::Message(format!(
                        "Unsupported plain session key version {:?}",
                        psk
                    ))));
                }
            };

            let msg = edata.decrypt(sk.clone())?;

            let mut mr = unpack_int(
                &msg,
                pkesk_decryptors,
                skesk_passwords.clone(),
                session_key,
                verifier,
                depth + 1,
            )?;

            // Return the session key in MessageResult
            mr.session_key = Some(((sym_alg).into(), key));

            Ok(mr)
        }
        Message::Literal(lit) => Ok(MessageResult {
            validated: vec![],
            session_key: None,
            cleartext: lit.clone(), // FIXME: avoid clone
        }),
        Message::Compressed(cd) => {
            let payload = cd.decompress()?;
            let msg = Message::from_bytes(payload)?;

            unpack_int(
                &msg,
                pkesk_decryptors,
                skesk_passwords,
                session_key,
                verifier,
                depth + 1,
            )
        }
        Message::Signed {
            one_pass_signature: _,
            message: inner,
            signature,
        } => {
            let mut vals: Vec<(Certificate, ComponentKeyPub, Signature)> = vec![];

            // FIXME: handle in streaming mode?

            // only consider "signature" if it passes our policy check, skip validation if not
            if crate::signature::signature_acceptable(signature) {
                // get the data we need to validate this signature against (either the bare literal
                // data [without pgp framing], or the raw bytes of the message).
                fn grab_data(m: &Message, depth: usize) -> Result<Vec<u8>, Error> {
                    if depth > MAX_RECURSION {
                        return Err(Error::Message(
                            "Excessive message nesting depth".to_string(),
                        ));
                    }

                    match m {
                        Message::Literal(lit) => Ok(lit.data().to_vec()),
                        Message::Compressed(cd) => {
                            let payload = cd.decompress()?;
                            let msg = Message::from_bytes(payload)?;

                            grab_data(&msg, depth + 1)
                        }
                        Message::Signed { message, .. } => {
                            let Some(message) = message else {
                                return Err(Error::Message(
                                    "No inner message found in signed message".to_string(),
                                ));
                            };
                            // FIXME: return raw message data, if notarizing signature?
                            grab_data(message.as_ref(), depth + 1)
                        }
                        Message::Encrypted { .. } => Ok(m.to_bytes()?),
                    }
                }

                let Some(inner) = inner else {
                    return Err(Error::Message("Inner message is None".to_string()));
                };
                let data = grab_data(inner.as_ref(), 0)?;

                for vcert in verifier {
                    if let Some(reference) = signature.created() {
                        // Only consider signers that are valid at signature creation time.
                        let cv = Checked::from(vcert);
                        let verifiers = cv.valid_signing_capable_component_keys_at(reference);

                        for v in verifiers {
                            if v.verify(signature, &data).is_ok() {
                                vals.push((
                                    vcert.clone(),
                                    v.as_componentkey().clone(),
                                    signature.clone(),
                                ));
                            }
                        }
                    }
                }
            }

            let Some(inner) = inner else {
                return Err(Error::Message("Inner message is None".to_string()));
            };
            let mut mr = unpack_int(
                inner.as_ref(),
                pkesk_decryptors,
                skesk_passwords,
                session_key,
                verifier,
                depth + 1,
            )?;

            mr.validated.append(&mut vals);

            // FIXME: deduplicate validations?

            Ok(mr)
        }
    }
}

/// Configuration for encryption, either as SeipdV1 or SeipdV2
#[derive(Debug, Copy, Clone, PartialEq)]
pub enum EncryptionMechanism {
    SeipdV1(SymmetricKeyAlgorithm),
    SeipdV2(AeadAlgorithm, SymmetricKeyAlgorithm),
}

/// Encrypt (and optionally sign) a message.
///
/// NOTE: `source` is expected to contain raw data, not an OpenPGP Message
///
/// FIXME: we need passwords to unlock signers!
///
/// FIXME: pass recipient primary (to set as intended recipient)
#[allow(clippy::too_many_arguments)]
pub fn encrypt(
    mechanism: EncryptionMechanism,
    recipients: Vec<ComponentKeyPub>,
    skesk_passwords: Vec<&[u8]>,
    signers: Vec<Tsk>,
    hash_algo: Option<HashAlgorithm>,
    source: &mut (dyn Read + Send + Sync),
    mut sink: &mut (dyn io::Write + Send + Sync),
    armor: bool,
) -> Result<Zeroizing<Vec<u8>>, Error> {
    let mut rng = rand::thread_rng();

    // 1. Define a new session key
    let session_key = match mechanism {
        EncryptionMechanism::SeipdV1(sym) | EncryptionMechanism::SeipdV2(_, sym) => {
            sym.new_session_key(&mut rng)
        }
    };

    let mut esk = vec![];

    // 2. Encrypt (pub) the session key, to each PublicKey.
    for recipient in recipients {
        let pkesk = match mechanism {
            EncryptionMechanism::SeipdV1(sym) => {
                recipient.pkesk_from_session_key_v3(&mut rng, &session_key, sym)?
            }
            EncryptionMechanism::SeipdV2(_, _) => {
                recipient.pkesk_from_session_key_v6(&mut rng, &session_key)?
            }
        };

        esk.push(Esk::PublicKeyEncryptedSessionKey(pkesk));
    }

    // 3. Encrypt the session key to each symmetric password
    for pw in skesk_passwords {
        let pass = String::from_utf8_lossy(pw).to_string();

        let skesk = match mechanism {
            EncryptionMechanism::SeipdV1(sym) => SymKeyEncryptedSessionKey::encrypt_v4(
                || pass,
                &session_key,
                StringToKey::new_default(&mut rng),
                sym,
            )?,
            EncryptionMechanism::SeipdV2(aead, sym) => {
                // "If much less memory is available, a uniformly safe option is Argon2id with
                // t=3 iterations, p=4 lanes, m=2^(16) (64 MiB of RAM),
                // 128-bit salt, and 256-bit tag size. This is the SECOND RECOMMENDED option."

                // FIXME: move these settings up to rPGP, as part of a convenience constructor?
                let s2k = StringToKey::new_argon2(&mut rng, 3, 4, 16);

                SymKeyEncryptedSessionKey::encrypt_v6(
                    &mut rng,
                    || pass,
                    &session_key,
                    s2k,
                    sym,
                    aead,
                )?
            }
        };

        esk.push(Esk::SymKeyEncryptedSessionKey(skesk));
    }

    // 4. Wrap the plaintext into a literal
    let mut data = vec![];
    source.read_to_end(&mut data)?;

    let lit = LiteralData::from_bytes((&[]).into(), &data);

    // 5. Maybe sign the message.
    let hash_algo = hash_algo.unwrap_or_default();

    let msg = if !signers.is_empty() {
        // FIXME: set Intended Recipient Fingerprint? (at least for v6)
        // "The OpenPGP Key fingerprint of the intended recipient primary key"
        // "When generating this subpacket in a version 6 signature, it SHOULD be marked as critical."

        let mut packets = vec![];
        packets.push(Packet::from(lit.clone()));

        for signer in signers {
            let data_signers: Vec<_> = signer.signing_capable_component_keys().collect();

            // FIXME: use all signing capable keys?
            if let Some(ds) = data_signers.first() {
                // FIXME: key decryption password!

                if let Message::Signed {
                    one_pass_signature,
                    signature,
                    ..
                } = ds.sign_msg(Message::Literal(lit.clone()), String::default, hash_algo)?
                {
                    if let Some(mut ops) = one_pass_signature {
                        if packets.len() > 1 {
                            // only the innermost signature should be marked "last",
                            // so we mark all others as non-last.
                            ops.last = 0;
                        }
                        packets.insert(0, Packet::from(ops));
                    }

                    packets.push(Packet::from(signature));
                }
            } else {
                return Err(Error::Message(
                    "No signing capable component key found for signer".to_string(),
                ));
            }
        }

        if let Some(Ok(msg)) = Message::from_packets(packets.into_iter().map(Ok).peekable()).next()
        {
            msg
        } else {
            // This shouldn't happen, if we construct a reasonable "packets"
            return Err(Error::Message(
                "Failed to construct Message from packets".to_string(),
            ));
        }
    } else {
        // we're not signing
        Message::Literal(lit)
    };

    // 6. Symmetrically Encrypt the message to the session key.
    let edata = match mechanism {
        EncryptionMechanism::SeipdV1(sym) => {
            Edata::SymEncryptedProtectedData(SymEncryptedProtectedData::encrypt_seipdv1(
                &mut rng,
                sym,
                &session_key,
                &msg.to_bytes()?,
            )?)
        }
        EncryptionMechanism::SeipdV2(aead, sym) => {
            Edata::SymEncryptedProtectedData(SymEncryptedProtectedData::encrypt_seipdv2(
                &mut rng,
                sym,
                aead,
                crate::policy::AEAD_CHUNK_SIZE,
                &session_key,
                &msg.to_bytes()?,
            )?)
        }
    };

    // 7. Put together encrypted message
    let msg = Message::Encrypted { esk, edata };

    // 8. Write encrypted message to sink
    match armor {
        true => msg.to_armored_writer(&mut sink, ArmorOptions::default())?,
        false => msg.to_writer(&mut sink)?,
    }

    // 9. Return session key
    Ok(session_key)
}
