#![forbid(unsafe_code)]

mod support;

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

use ironstorage::{
    crypto::{
        CryptoError, DetachedSignatureBytes, KeyInfo, KeyStore, SecretProvider, SecretProviderError,
    },
    repository::{EncryptedEntry, SecretBytes},
};
use pgp::{
    composed::{EncryptionCaps, KeyType, Message, SecretKeyParamsBuilder, SubkeyParamsBuilder},
    crypto::{ecc_curve::ECCCurve, hash::HashAlgorithm, sym::SymmetricKeyAlgorithm},
    ser::Serialize as _,
    types::{Password, Timestamp},
};
use rand_chacha::{ChaCha20Rng, rand_core::SeedableRng};
use smallvec::smallvec;
use support::compatibility::{FixtureSet, TestResult};

#[derive(Default)]
struct FixtureSecrets {
    values: BTreeMap<String, Vec<u8>>,
    requests: 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(),
        }
    }

    fn one(fingerprint: &str, passphrase: impl AsRef<[u8]>) -> Self {
        Self {
            values: BTreeMap::from([(fingerprint.to_owned(), passphrase.as_ref().to_vec())]),
            requests: 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)
    }
}

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 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()))?;
    assert!(parsed.is_encrypted());
    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()?)
}