#![forbid(unsafe_code)]

mod support;

use std::{collections::BTreeMap, io::Cursor, path::Path};

use ironstorage::{
    crypto::{
        CryptoError, DetachedSignatureBytes, KeyInfo, KeyStore, SecretProvider, SecretProviderError,
    },
    read::{ShowResult, VaultReader},
    repository::{EncryptedEntry, Repository, SecretBytes},
};
use pgp::{
    composed::{
        Edata, EncryptionCaps, Esk, KeyType, Message, SecretKeyParamsBuilder, SubkeyParamsBuilder,
    },
    crypto::{
        aead::AeadAlgorithm, ecc_curve::ECCCurve, hash::HashAlgorithm, sym::SymmetricKeyAlgorithm,
    },
    packet::{ProtectedDataConfig, SymEncryptedProtectedDataConfig},
    ser::Serialize as _,
    types::{Password, Timestamp},
};
use rand_chacha::{ChaCha20Rng, rand_core::SeedableRng};
use sha2::{Digest as _, Sha256};
use smallvec::smallvec;
use support::compatibility::{FixtureSet, TestResult};

#[derive(Default)]
struct FixtureSecrets {
    values: BTreeMap<String, Vec<u8>>,
    requests: Vec<String>,
    accepted: Vec<String>,
    rejected: Vec<String>,
}

impl FixtureSecrets {
    fn all(fixture: &FixtureSet) -> Self {
        Self {
            values: fixture
                .generated
                .keys
                .iter()
                .map(|key| {
                    (
                        key.primary_fingerprint.clone(),
                        key.passphrase.as_bytes().to_vec(),
                    )
                })
                .collect(),
            requests: Vec::new(),
            accepted: Vec::new(),
            rejected: Vec::new(),
        }
    }

    fn one(fingerprint: &str, passphrase: impl AsRef<[u8]>) -> Self {
        Self {
            values: BTreeMap::from([(fingerprint.to_owned(), passphrase.as_ref().to_vec())]),
            requests: Vec::new(),
            accepted: Vec::new(),
            rejected: Vec::new(),
        }
    }
}

impl SecretProvider for FixtureSecrets {
    fn secret_for(&mut self, key: &KeyInfo) -> Result<SecretBytes, SecretProviderError> {
        let fingerprint = key.fingerprint().as_str().to_owned();
        self.requests.push(fingerprint.clone());
        self.values
            .get(&fingerprint)
            .cloned()
            .map(SecretBytes::new)
            .ok_or(SecretProviderError::Unavailable)
    }

    fn secret_accepted(&mut self, key: &KeyInfo) -> Result<(), SecretProviderError> {
        self.accepted.push(key.fingerprint().as_str().to_owned());
        Ok(())
    }

    fn secret_rejected(&mut self, key: &KeyInfo) {
        self.rejected.push(key.fingerprint().as_str().to_owned());
    }
}

struct CancelledProvider;

impl SecretProvider for CancelledProvider {
    fn secret_for(&mut self, _key: &KeyInfo) -> Result<SecretBytes, SecretProviderError> {
        Err(SecretProviderError::Cancelled)
    }
}

#[test]
fn imports_armored_and_binary_public_and_protected_secret_keys() -> TestResult {
    let fixture = FixtureSet::load()?;
    let alice = fixture.key("alice")?;
    let bob = fixture.key("bob")?;
    let mut keys = KeyStore::new();

    keys.import(&fixture.read(&alice.public_armor)?)?;
    keys.import(&fixture.read(&alice.secret_binary)?)?;
    keys.import(&fixture.read(&bob.public_binary)?)?;
    keys.import(&fixture.read(&bob.secret_armor)?)?;

    assert_eq!(keys.len(), 2);
    let infos = keys.infos().collect::<Vec<_>>();
    assert!(infos.iter().all(KeyInfo::has_secret));
    assert!(infos.iter().all(KeyInfo::can_encrypt));
    assert!(infos.iter().all(KeyInfo::can_sign));
    assert!(format!("{keys:?}").contains("key_count: 2"));
    assert!(!format!("{keys:?}").contains(&alice.passphrase));

    let loaded = KeyStore::load(fixture.path("keys"))?;
    assert_eq!(loaded.len(), 2, "duplicate armor/binary exports are merged");
    Ok(())
}

#[test]
fn resolves_pass_identities_and_recipient_file_rules() -> TestResult {
    let fixture = FixtureSet::load()?;
    let keys = KeyStore::load(fixture.path("keys"))?;
    let alice = fixture.key("alice")?;
    let info = keys
        .infos()
        .find(|info| info.fingerprint().as_str() == alice.primary_fingerprint)
        .expect("Alice info");

    for identity in [
        alice.primary_fingerprint.as_str(),
        &alice.primary_fingerprint[24..],
        &alice.primary_fingerprint[32..],
        info.key_id(),
        alice.user_id.as_str(),
        alice.encryption_subkey_fingerprint.as_str(),
        &alice.encryption_subkey_fingerprint[24..],
    ] {
        assert_eq!(
            keys.resolve(identity)?.fingerprint().as_str(),
            alice.primary_fingerprint
        );
    }
    let recipients = keys.resolve_recipients(
        format!(
            "  {}  # primary recipient\n\n{} # duplicate\n",
            alice.primary_fingerprint, alice.user_id
        )
        .as_bytes(),
    )?;
    assert_eq!(recipients.len(), 1);

    assert!(matches!(
        keys.resolve("missing@example.invalid"),
        Err(CryptoError::MissingIdentity { .. })
    ));
    assert!(matches!(
        keys.resolve_recipients(b" # comments only\n"),
        Err(CryptoError::MissingRecipients)
    ));
    assert!(matches!(
        keys.resolve_recipients(&[0xff]),
        Err(CryptoError::InvalidRecipientFile)
    ));
    Ok(())
}

#[test]
fn reports_exact_user_id_ambiguity() -> TestResult {
    let fixture = FixtureSet::load()?;
    let alice = fixture.key("alice")?;
    let mut keys = KeyStore::new();
    keys.import(&fixture.read(&alice.public_binary)?)?;
    keys.import(&duplicate_user_id_key(&alice.user_id)?)?;

    assert!(matches!(
        keys.resolve(&alice.user_id),
        Err(CryptoError::AmbiguousIdentity { .. })
    ));
    Ok(())
}

#[test]
fn decrypts_every_gnupg_audited_pass_fixture() -> TestResult {
    let fixture = FixtureSet::load()?;
    let keys = KeyStore::load(fixture.path("keys"))?;
    let mut provider = FixtureSecrets::all(&fixture);

    for entry in &fixture.generated.entries {
        let ciphertext = fixture.read(Path::new("stores").join(&entry.store).join(&entry.path))?;
        let expected = fixture.read(
            Path::new("expected")
                .join(&entry.store)
                .join(&entry.plaintext),
        )?;
        let plaintext = keys.decrypt(&EncryptedEntry::new(ciphertext), &mut provider)?;
        assert_eq!(
            plaintext.expose(),
            expected,
            "fixture {}/{}",
            entry.store,
            entry.path
        );
    }
    assert!(!provider.requests.is_empty());
    Ok(())
}

#[test]
fn decrypts_gnupg_packet_20_aes256_ocb_through_the_vault_reader() -> TestResult {
    let fixture = FixtureSet::load()?;
    let aead = &fixture.gnupg_aead;
    let alice = fixture.key("alice")?;
    assert_eq!(aead.format, 1);
    assert_eq!(aead.producer, "GnuPG 2.4.8");
    assert!(aead.command_profile.contains("--force-ocb"));
    assert_eq!(aead.packet_tag, 20);
    assert_eq!(aead.cipher, "AES256");
    assert_eq!(aead.aead_mode, "OCB");
    assert_eq!(
        aead.recipient_primary_fingerprint,
        alice.primary_fingerprint
    );
    assert_eq!(
        aead.recipient_encryption_subkey_fingerprint,
        alice.encryption_subkey_fingerprint
    );

    let ciphertext = fixture.read(&aead.entry)?;
    assert_eq!(
        hex::encode(Sha256::digest(&ciphertext)),
        aead.ciphertext_sha256
    );
    let message = Message::from_bytes(Cursor::new(&ciphertext))?;
    let Message::Encrypted { esk, edata, .. } = &message else {
        panic!("GnuPG AEAD fixture must be encrypted");
    };
    let [Esk::PublicKeyEncryptedSessionKey(pkesk)] = esk.as_slice() else {
        panic!("GnuPG AEAD fixture must have one public-key recipient");
    };
    let secret = fixture.parse_secret_key(alice)?;
    assert!(pkesk.match_identity(&secret.secret_subkeys[0].key));
    assert!(!pkesk.match_identity(&secret.primary_key));
    let Edata::GnupgAeadData { reader } = edata else {
        panic!("GnuPG fixture must use packet type 20");
    };
    let ProtectedDataConfig::GnupgAead(config) = reader.config() else {
        panic!("packet type 20 must carry GnuPG AEAD configuration");
    };
    assert_eq!(config.sym_alg, SymmetricKeyAlgorithm::AES256);
    assert_eq!(config.aead, AeadAlgorithm::Ocb);

    let repository = Repository::open(fixture.path(format!("gnupg-aead/{}", aead.store)))?;
    let keys = KeyStore::load(fixture.path("keys"))?;
    let reader = VaultReader::new(&repository, &keys);
    let mut provider = FixtureSecrets::one(&alice.primary_fingerprint, &alice.passphrase);
    let ShowResult::Entry(plaintext) = reader.show(Some("paypal"), &mut provider)? else {
        panic!("paypal must resolve as an entry");
    };
    assert_eq!(plaintext.expose(), fixture.read(&aead.plaintext)?);
    assert_eq!(
        provider.requests,
        std::slice::from_ref(&alice.primary_fingerprint)
    );
    assert_eq!(
        provider.accepted,
        std::slice::from_ref(&alice.primary_fingerprint)
    );
    assert!(provider.rejected.is_empty());
    Ok(())
}

#[test]
fn rejects_invalid_gnupg_packet_20_inputs_without_accepting_secrets() -> TestResult {
    let fixture = FixtureSet::load()?;
    let aead = &fixture.gnupg_aead;
    let alice = fixture.key("alice")?;
    let ciphertext = fixture.read(&aead.entry)?;
    let keys = KeyStore::load(fixture.path(&alice.secret_armor))?;

    let mut wrong = FixtureSecrets::one(&alice.primary_fingerprint, b"wrong-aead-passphrase");
    assert!(matches!(
        keys.decrypt(&EncryptedEntry::new(ciphertext.clone()), &mut wrong),
        Err(CryptoError::DecryptionFailed)
    ));
    assert!(wrong.accepted.is_empty());
    assert_eq!(
        wrong.rejected,
        std::slice::from_ref(&alice.primary_fingerprint)
    );

    let mut tampered = ciphertext.clone();
    *tampered.last_mut().expect("non-empty AEAD fixture") ^= 0x01;
    let mut correct = FixtureSecrets::one(&alice.primary_fingerprint, &alice.passphrase);
    assert!(matches!(
        keys.decrypt(&EncryptedEntry::new(tampered), &mut correct),
        Err(CryptoError::DecryptionFailed)
    ));
    assert!(correct.accepted.is_empty());
    assert_eq!(
        correct.rejected,
        std::slice::from_ref(&alice.primary_fingerprint)
    );

    let public_only = KeyStore::load(fixture.path(&alice.public_armor))?;
    let mut unavailable = FixtureSecrets::default();
    assert!(matches!(
        public_only.decrypt(&EncryptedEntry::new(ciphertext.clone()), &mut unavailable),
        Err(CryptoError::MissingSecretKey)
    ));
    assert!(unavailable.requests.is_empty());

    let mut unsupported = ciphertext;
    let config = unsupported
        .windows(3)
        .position(|window| window == [1, 9, 2])
        .expect("packet-20 version, AES-256, and OCB configuration");
    unsupported[config + 2] = 1;
    let mut unused = FixtureSecrets::one(&alice.primary_fingerprint, &alice.passphrase);
    assert!(matches!(
        keys.decrypt(&EncryptedEntry::new(unsupported), &mut unused),
        Err(CryptoError::CorruptMessage)
    ));
    assert!(unused.requests.is_empty());
    assert!(unused.accepted.is_empty());
    assert!(unused.rejected.is_empty());
    Ok(())
}

#[test]
fn encrypts_for_one_or_many_recipients_without_compression() -> TestResult {
    let fixture = FixtureSet::load()?;
    let keys = KeyStore::load(fixture.path("keys"))?;
    let alice = fixture.key("alice")?;
    let bob = fixture.key("bob")?;
    let recipients = keys.resolve_recipients(
        format!(
            "{}\n{}\n",
            alice.primary_fingerprint, bob.primary_fingerprint
        )
        .as_bytes(),
    )?;
    let expected = b"generated-by-ironstorage\nlogin: fixture\n";
    let ciphertext = keys.encrypt(SecretBytes::new(expected.to_vec()), &recipients)?;

    let parsed = Message::from_bytes(Cursor::new(ciphertext.as_bytes()))?;
    let Message::Encrypted { edata, .. } = &parsed else {
        panic!("IronStorage output must be encrypted");
    };
    let Edata::SymEncryptedProtectedData { reader } = edata else {
        panic!("IronStorage output must remain packet type 18");
    };
    assert_eq!(
        reader.config(),
        &ProtectedDataConfig::Seipd(SymEncryptedProtectedDataConfig::V1)
    );
    for key in [alice, bob] {
        let mut provider = FixtureSecrets::one(&key.primary_fingerprint, &key.passphrase);
        let plaintext = keys.decrypt(&ciphertext, &mut provider)?;
        assert_eq!(plaintext.expose(), expected);
        assert_eq!(provider.requests.last(), Some(&key.primary_fingerprint));
    }
    let independent = fixture.parse_secret_key(alice)?;
    assert_eq!(
        decrypt_independently(ciphertext.as_bytes(), &independent, &alice.passphrase)?,
        expected
    );
    assert!(matches!(
        keys.encrypt(SecretBytes::new(Vec::new()), &[]),
        Err(CryptoError::MissingRecipients)
    ));
    Ok(())
}

#[test]
fn protected_key_failures_are_typed_and_redacted() -> TestResult {
    let fixture = FixtureSet::load()?;
    let alice = fixture.key("alice")?;
    let entry = fixture
        .generated
        .entries
        .iter()
        .find(|entry| entry.recipients == [alice.primary_fingerprint.clone()])
        .expect("Alice fixture");
    let ciphertext = EncryptedEntry::new(
        fixture.read(Path::new("stores").join(&entry.store).join(&entry.path))?,
    );
    let keys = KeyStore::load(fixture.path(&alice.secret_armor))?;
    let mut wrong = FixtureSecrets::one(&alice.primary_fingerprint, b"very-secret-wrong-value");
    let error = keys.decrypt(&ciphertext, &mut wrong).unwrap_err();
    assert_eq!(error, CryptoError::DecryptionFailed);
    assert!(!format!("{error:?}").contains("very-secret-wrong-value"));
    assert!(!error.to_string().contains("very-secret-wrong-value"));

    let public_only = KeyStore::load(fixture.path(&alice.public_armor))?;
    let mut unavailable = FixtureSecrets::default();
    assert!(matches!(
        public_only.decrypt(&ciphertext, &mut unavailable),
        Err(CryptoError::MissingSecretKey)
    ));
    assert!(unavailable.requests.is_empty());

    let mut cancelled = FixtureSecrets::default();
    let error = keys.decrypt(&ciphertext, &mut cancelled).unwrap_err();
    assert!(matches!(
        error,
        CryptoError::SecretProvider {
            reason: SecretProviderError::Unavailable,
            ..
        }
    ));
    assert!(matches!(
        keys.decrypt(&ciphertext, &mut CancelledProvider),
        Err(CryptoError::SecretProvider {
            reason: SecretProviderError::Cancelled,
            ..
        })
    ));
    Ok(())
}

#[test]
fn signs_and_verifies_gpg_id_with_an_explicit_trust_set() -> TestResult {
    let fixture = FixtureSet::load()?;
    let keys = KeyStore::load(fixture.path("keys"))?;
    let alice = fixture.key("alice")?;
    let bob = fixture.key("bob")?;
    let alice_handle = keys.resolve(&alice.primary_fingerprint)?;
    let bob_handle = keys.resolve(&bob.primary_fingerprint)?;
    let recipient_file = fixture.read("stores/basic/.gpg-id")?;
    let checked_in = DetachedSignatureBytes::new(fixture.read("stores/basic/.gpg-id.sig")?);

    assert_eq!(
        keys.verify(
            &recipient_file,
            &checked_in,
            std::slice::from_ref(&alice_handle)
        )?,
        alice_handle
    );
    assert_eq!(
        keys.verify(&recipient_file, &checked_in, &[bob_handle]),
        Err(CryptoError::InvalidSignature)
    );
    assert_eq!(
        keys.verify(
            b"tampered\n",
            &checked_in,
            std::slice::from_ref(&alice_handle)
        ),
        Err(CryptoError::InvalidSignature)
    );
    assert_eq!(
        keys.verify(
            &recipient_file,
            &DetachedSignatureBytes::new(vec![0xff]),
            std::slice::from_ref(&alice_handle),
        ),
        Err(CryptoError::CorruptSignature)
    );

    let mut provider = FixtureSecrets::one(&alice.primary_fingerprint, &alice.passphrase);
    let generated = keys.sign(&recipient_file, &alice_handle, &mut provider)?;
    assert_eq!(
        keys.verify(
            &recipient_file,
            &generated,
            std::slice::from_ref(&alice_handle)
        )?,
        alice_handle
    );
    let mut wrong = FixtureSecrets::one(&alice.primary_fingerprint, b"wrong");
    assert_eq!(
        keys.sign(&recipient_file, &alice_handle, &mut wrong),
        Err(CryptoError::SigningFailed)
    );
    Ok(())
}

#[test]
fn rejects_corrupt_keys_messages_and_signatures() -> TestResult {
    let fixture = FixtureSet::load()?;
    let keys = KeyStore::load(fixture.path("keys"))?;
    let mut provider = FixtureSecrets::all(&fixture);
    let mut empty = KeyStore::new();

    assert_eq!(
        empty.import(b"not an OpenPGP key"),
        Err(CryptoError::CorruptKeyMaterial)
    );
    assert!(empty.is_empty());
    assert!(matches!(
        keys.decrypt(&EncryptedEntry::new(vec![0xff, 0x00]), &mut provider),
        Err(CryptoError::CorruptMessage)
    ));
    Ok(())
}

fn duplicate_user_id_key(user_id: &str) -> TestResult<Vec<u8>> {
    let created_at = Timestamp::from_secs(1_704_067_201);
    let subkey = SubkeyParamsBuilder::default()
        .key_type(KeyType::ECDH(ECCCurve::Curve25519Legacy))
        .can_encrypt(EncryptionCaps::All)
        .created_at(created_at)
        .build()?;
    let mut rng = ChaCha20Rng::from_seed([0x55; 32]);
    let secret = SecretKeyParamsBuilder::default()
        .key_type(KeyType::Ed25519Legacy)
        .can_certify(true)
        .can_sign(true)
        .created_at(created_at)
        .primary_user_id(user_id.to_owned())
        .preferred_symmetric_algorithms(smallvec![SymmetricKeyAlgorithm::AES256])
        .preferred_hash_algorithms(smallvec![HashAlgorithm::Sha256])
        .preferred_compression_algorithms(smallvec![])
        .subkey(subkey)
        .build()?
        .generate(&mut rng)?;
    let mut bytes = Vec::new();
    secret.to_public_key().to_writer(&mut bytes)?;
    Ok(bytes)
}

fn decrypt_independently(
    ciphertext: &[u8],
    secret: &pgp::composed::SignedSecretKey,
    passphrase: &str,
) -> TestResult<Vec<u8>> {
    let message = Message::from_bytes(Cursor::new(ciphertext))?;
    let mut decrypted = message.decrypt(&Password::from(passphrase), secret)?;
    Ok(decrypted.as_data_vec()?)
}