#![forbid(unsafe_code)]

mod support;

use std::collections::BTreeMap;

use data_encoding::BASE32_NOPAD;
use ironstorage::{
    command::{OtpAppendRequest, OtpInputSource, OtpInsertRequest},
    crypto::{KeyInfo, KeyStore, SecretProvider, SecretProviderError},
    otp::{OtpAlgorithm, OtpError, OtpInput, OtpKind, OtpService, OtpUri},
    recipient::RecipientPolicyManager,
    repository::{EntryPath, Repository, SecretBytes},
    write::{EntryCommit, EntryCommitError, EntryCommitter, OverwriteDecision},
};
use support::compatibility::{FixtureSet, TestResult};

struct FixtureSecrets(BTreeMap<String, Vec<u8>>);

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

impl SecretProvider for FixtureSecrets {
    fn secret_for(&mut self, key: &KeyInfo) -> Result<SecretBytes, SecretProviderError> {
        self.0
            .get(key.fingerprint().as_str())
            .cloned()
            .map(SecretBytes::new)
            .ok_or(SecretProviderError::Unavailable)
    }
}

#[derive(Default)]
struct Committer {
    changes: Vec<EntryCommit>,
    fail: bool,
}

impl EntryCommitter for Committer {
    fn commit(&mut self, change: &EntryCommit) -> Result<(), EntryCommitError> {
        self.changes.push(change.clone());
        if self.fail {
            Err(EntryCommitError::new("simulated Git failure"))
        } else {
            Ok(())
        }
    }
}

#[test]
fn key_uri_parsing_defaults_derivation_and_validation_are_typed() -> TestResult {
    let text = "otpauth://totp/Example:alice%40example.test?secret=jbswy3dpehpk3pxp&issuer=Example";
    let uri = OtpUri::parse_str(text)?;
    assert_eq!(uri.encoded().expose(), text.as_bytes());
    assert_eq!(uri.kind(), OtpKind::Totp);
    assert_eq!(uri.issuer(), Some("Example"));
    assert_eq!(uri.account(), "alice@example.test");
    assert_eq!(uri.algorithm(), OtpAlgorithm::Sha1);
    assert_eq!(uri.digits(), 6);
    assert_eq!(uri.period(), Some(30));
    assert_eq!(uri.counter(), None);
    assert_eq!(
        uri.derived_entry()?.to_string(),
        "Example/alice@example.test"
    );
    assert!(!format!("{uri:?}").contains("jbswy3dpehpk3pxp"));

    let derived = OtpUri::from_input(
        &OtpInputSource::Secret {
            issuer: Some("ACME Co".to_owned()),
            account: Some("alice@example.test".to_owned()),
        },
        OtpInput::line(b"JBSWY3DPEHPK3PXP".to_vec())?,
    )?;
    assert_eq!(
        derived.encoded().expose(),
        b"otpauth://totp/ACME+Co:alice%40example.test?secret=JBSWY3DPEHPK3PXP&issuer=ACME+Co"
    );
    assert_eq!(
        derived.derived_entry()?.to_string(),
        "ACME Co/alice@example.test"
    );

    let custom = OtpUri::parse_str(
        "otpauth://totp/custom?secret=JBSWY3DPEHPK3PXP&algorithm=SHA256&digits=8&period=60",
    )?;
    assert_eq!(custom.period(), Some(60));
    assert_eq!(custom.digits(), 8);
    assert_eq!(
        custom.code_at(119)?.expose(),
        custom.code_for_counter(1)?.expose()
    );

    for (invalid, expected) in [
        ("https://example.test/not-otp", OtpError::InvalidScheme),
        (
            "otpauth://totp/account?issuer=Example",
            OtpError::MissingSecret,
        ),
        (
            "otpauth://hotp/account?secret=JBSWY3DPEHPK3PXP",
            OtpError::MissingCounter,
        ),
        (
            "otpauth://totp/account?secret=not-base32!",
            OtpError::InvalidSecret,
        ),
        (
            "otpauth://totp/account?secret=MY========",
            OtpError::InvalidSecret,
        ),
        (
            "otpauth://totp/A:account?secret=JBSWY3DPEHPK3PXP&issuer=B",
            OtpError::IssuerMismatch,
        ),
    ] {
        assert_eq!(
            std::mem::discriminant(&OtpUri::parse_str(invalid).expect_err("invalid URI")),
            std::mem::discriminant(&expected)
        );
    }
    assert!(matches!(
        OtpUri::parse_str("otpauth://totp/account?secret=JBSWY3DPEHPK3PXP&secret=JBSWY3DPEHPK3PXP"),
        Err(OtpError::DuplicateParameter(_))
    ));
    assert!(matches!(
        OtpInput::hidden(b"one".to_vec(), b"two".to_vec()),
        Err(OtpError::ConfirmationMismatch)
    ));
    assert!(matches!(
        OtpInput::line(b"two\nlines".to_vec()),
        Err(OtpError::InvalidInput)
    ));
    Ok(())
}

#[test]
fn rfc4226_hotp_vectors_pass() -> TestResult {
    let secret = BASE32_NOPAD.encode(b"12345678901234567890");
    let uri = OtpUri::parse_str(&format!("otpauth://hotp/RFC4226?secret={secret}&counter=0"))?;
    for (counter, expected) in [
        "755224", "287082", "359152", "969429", "338314", "254676", "287922", "162583", "399871",
        "520489",
    ]
    .into_iter()
    .enumerate()
    {
        assert_eq!(
            uri.code_for_counter(counter as u64)?.expose(),
            expected.as_bytes()
        );
    }
    Ok(())
}

#[test]
fn rfc6238_totp_vectors_cover_all_algorithms() -> TestResult {
    let vectors = [
        (59, "94287082", "46119246", "90693936"),
        (1_111_111_109, "07081804", "68084774", "25091201"),
        (1_111_111_111, "14050471", "67062674", "99943326"),
        (1_234_567_890, "89005924", "91819424", "93441116"),
        (2_000_000_000, "69279037", "90698825", "38618901"),
        (20_000_000_000, "65353130", "77737706", "47863826"),
    ];
    let tokens = [
        token(b"12345678901234567890", "SHA1")?,
        token(b"12345678901234567890123456789012", "SHA256")?,
        token(
            b"1234567890123456789012345678901234567890123456789012345678901234",
            "SHA512",
        )?,
    ];
    for (timestamp, sha1, sha256, sha512) in vectors {
        for (token, expected) in tokens.iter().zip([sha1, sha256, sha512]) {
            assert_eq!(token.code_at(timestamp)?.expose(), expected.as_bytes());
        }
    }
    Ok(())
}

#[test]
fn pass_otp_fixtures_round_trip_code_uri_insert_and_append() -> TestResult {
    let fixture = FixtureSet::load()?;
    let store = fixture.materialize_store("basic")?;
    let repository = Repository::open(store.path())?;
    let keys = KeyStore::load(fixture.path("keys"))?;
    let service = OtpService::new(&repository, &keys);
    let mut provider = FixtureSecrets::all(&fixture);
    let mut committer = Committer::default();

    let fixture_uri = fixture
        .read("expected/basic/otp/totp.txt")?
        .split(|byte| *byte == b'\n')
        .find(|line| line.starts_with(b"otpauth://"))
        .expect("fixture URI")
        .to_vec();
    OtpService::validate(std::str::from_utf8(&fixture_uri)?)?;
    assert_eq!(
        service.uri("otp/totp", &mut provider)?.encoded().expose(),
        fixture_uri
    );
    let code = service.code("otp/totp", 59, None, &mut provider, &mut committer)?;
    assert_eq!(code.code().expose().len(), 6);
    assert!(code.code().expose().iter().all(u8::is_ascii_digit));
    assert!(committer.changes.is_empty());

    let replacement_uri = "otpauth://totp/Replaced:alice?secret=JBSWY3DPEHPK3PXP&issuer=Replaced";
    let replace_request = OtpAppendRequest {
        entry: "otp/totp".to_owned(),
        force: true,
        echo: true,
        source: OtpInputSource::Uri,
    };
    let replace_session = service.begin_append(&replace_request, &mut provider)?;
    assert!(!replace_session.requires_replace_confirmation());
    service.finish_append(
        replace_session,
        OtpInput::line(replacement_uri.as_bytes().to_vec())?,
        OverwriteDecision::Decline,
        None,
        &mut committer,
    )?;
    assert_eq!(
        decrypt(&repository, &keys, "otp/totp", &mut provider)?.expose(),
        format!("fixture-password\n{replacement_uri}\n").as_bytes()
    );
    assert_eq!(
        committer.changes.last().expect("replace commit").message(),
        "Replace OTP secret for otp/totp."
    );

    let hotp_path = EntryPath::parse("otp/hotp")?;
    let hotp_before = repository.read_entry(&hotp_path)?;
    let hotp = service.code("otp/hotp", 0, None, &mut provider, &mut committer)?;
    assert_eq!(hotp.counter(), Some(1));
    assert_eq!(hotp.code().expose().len(), 8);
    assert_eq!(
        committer.changes.last().expect("HOTP commit").message(),
        "Increment HOTP counter for otp/hotp."
    );
    let hotp_plaintext = keys.decrypt(&repository.read_entry(&hotp_path)?, &mut provider)?;
    assert!(
        hotp_plaintext
            .expose()
            .windows(9)
            .any(|part| part == b"counter=1")
    );
    assert_ne!(repository.read_entry(&hotp_path)?, hotp_before);

    let inserted_uri = "otpauth://totp/New:alice?secret=JBSWY3DPEHPK3PXP&issuer=New";
    let insert = OtpInsertRequest {
        entry: Some("otp/new".to_owned()),
        force: false,
        echo: false,
        source: OtpInputSource::Uri,
    };
    let plan =
        service.prepare_insert(&insert, OtpInput::line(inserted_uri.as_bytes().to_vec())?)?;
    assert!(!plan.requires_path_confirmation());
    assert!(!plan.requires_overwrite_confirmation());
    service.finish_insert(
        plan,
        OverwriteDecision::Allow,
        OverwriteDecision::Allow,
        None,
        &mut committer,
    )?;
    assert_eq!(
        decrypt(&repository, &keys, "otp/new", &mut provider)?.expose(),
        format!("{inserted_uri}\n").as_bytes()
    );

    let derived = OtpInsertRequest {
        entry: None,
        force: false,
        echo: false,
        source: OtpInputSource::Secret {
            issuer: Some("Issuer".to_owned()),
            account: Some("account".to_owned()),
        },
    };
    let plan = service.prepare_insert(
        &derived,
        OtpInput::hidden(b"JBSWY3DPEHPK3PXP".to_vec(), b"JBSWY3DPEHPK3PXP".to_vec())?,
    )?;
    assert_eq!(plan.path().to_string(), "Issuer/account");
    assert!(plan.requires_path_confirmation());
    service.finish_insert(
        plan,
        OverwriteDecision::Allow,
        OverwriteDecision::Allow,
        None,
        &mut committer,
    )?;
    assert_eq!(
        decrypt(&repository, &keys, "Issuer/account", &mut provider)?.expose(),
        b"otpauth://totp/Issuer:account?secret=JBSWY3DPEHPK3PXP&issuer=Issuer\n"
    );

    let original = decrypt(&repository, &keys, "email/personal", &mut provider)?;
    let append = OtpAppendRequest {
        entry: "email/personal".to_owned(),
        force: false,
        echo: false,
        source: OtpInputSource::Uri,
    };
    let session = service.begin_append(&append, &mut provider)?;
    assert!(!session.requires_replace_confirmation());
    service.finish_append(
        session,
        OtpInput::line(inserted_uri.as_bytes().to_vec())?,
        OverwriteDecision::Allow,
        None,
        &mut committer,
    )?;
    let appended = decrypt(&repository, &keys, "email/personal", &mut provider)?;
    assert!(appended.expose().starts_with(original.expose()));
    assert!(
        appended
            .expose()
            .ends_with(format!("{inserted_uri}\n").as_bytes())
    );
    assert_eq!(
        committer.changes.last().expect("append commit").message(),
        "Append OTP secret for email/personal."
    );
    Ok(())
}

#[test]
fn malformed_ambiguous_declined_and_commit_failures_never_mutate() -> TestResult {
    let fixture = FixtureSet::load()?;
    let store = fixture.materialize_store("basic")?;
    let repository = Repository::open(store.path())?;
    let keys = KeyStore::load(fixture.path("keys"))?;
    let service = OtpService::new(&repository, &keys);
    let mut provider = FixtureSecrets::all(&fixture);
    let mut committer = Committer::default();
    let path = EntryPath::parse("otp/totp")?;
    let original = repository.read_entry(&path)?;

    let request = OtpAppendRequest {
        entry: "otp/totp".to_owned(),
        force: false,
        echo: false,
        source: OtpInputSource::Uri,
    };
    let session = service.begin_append(&request, &mut provider)?;
    assert!(session.requires_replace_confirmation());
    assert!(matches!(
        service.finish_append(
            session,
            OtpInput::line(b"otpauth://totp/new?secret=JBSWY3DPEHPK3PXP".to_vec())?,
            OverwriteDecision::Decline,
            None,
            &mut committer,
        ),
        Err(OtpError::Cancelled)
    ));
    assert_eq!(repository.read_entry(&path)?, original);
    assert!(committer.changes.is_empty());

    let forced = OtpAppendRequest {
        force: true,
        ..request.clone()
    };
    let first = service.begin_append(&forced, &mut provider)?;
    let stale = service.begin_append(&forced, &mut provider)?;
    service.finish_append(
        first,
        OtpInput::line(b"otpauth://totp/first?secret=JBSWY3DPEHPK3PXP".to_vec())?,
        OverwriteDecision::Allow,
        None,
        &mut committer,
    )?;
    let committed = repository.read_entry(&path)?;
    assert!(matches!(
        service.finish_append(
            stale,
            OtpInput::line(b"otpauth://totp/stale?secret=JBSWY3DPEHPK3PXP".to_vec())?,
            OverwriteDecision::Allow,
            None,
            &mut committer,
        ),
        Err(OtpError::ConcurrentModification { .. })
    ));
    assert_eq!(repository.read_entry(&path)?, committed);
    assert_eq!(committer.changes.len(), 1);

    let non_otp = EntryPath::parse("email/personal")?;
    let non_otp_before = repository.read_entry(&non_otp)?;
    assert!(matches!(
        service.uri("email/personal", &mut provider),
        Err(OtpError::MissingUri { .. })
    ));
    assert_eq!(repository.read_entry(&non_otp)?, non_otp_before);
    assert_eq!(committer.changes.len(), 1);

    let hotp = EntryPath::parse("otp/hotp")?;
    let hotp_before = repository.read_entry(&hotp)?;
    committer.fail = true;
    assert!(matches!(
        service.code("otp/hotp", 0, None, &mut provider, &mut committer),
        Err(OtpError::Commit(_))
    ));
    assert!(
        decrypt(&repository, &keys, "otp/hotp", &mut provider)?
            .expose()
            .windows(9)
            .any(|part| part == b"counter=0")
    );
    assert_eq!(service.uri("otp/hotp", &mut provider)?.counter(), Some(0));
    assert_eq!(repository.read_entry(&hotp)?, hotp_before);

    let duplicate_path = EntryPath::parse("otp/duplicate")?;
    let duplicate_plaintext = SecretBytes::new(
        b"password\notpauth://totp/one?secret=JBSWY3DPEHPK3PXP\notpauth://totp/two?secret=JBSWY3DPEHPK3PXP\n"
            .to_vec(),
    );
    let recipients =
        RecipientPolicyManager::new(&repository, &keys).resolve_for_entry(&duplicate_path, None)?;
    repository.write_entry(
        &duplicate_path,
        &keys.encrypt(duplicate_plaintext, recipients.recipients())?,
    )?;
    let duplicate_before = repository.read_entry(&duplicate_path)?;
    let duplicate_request = OtpAppendRequest {
        entry: "otp/duplicate".to_owned(),
        force: true,
        echo: true,
        source: OtpInputSource::Uri,
    };
    assert!(matches!(
        service.begin_append(&duplicate_request, &mut provider),
        Err(OtpError::AmbiguousUri { .. })
    ));
    assert_eq!(repository.read_entry(&duplicate_path)?, duplicate_before);
    Ok(())
}

fn token(secret: &[u8], algorithm: &str) -> Result<OtpUri, OtpError> {
    OtpUri::parse_str(&format!(
        "otpauth://totp/RFC6238?secret={}&algorithm={algorithm}&digits=8&period=30",
        BASE32_NOPAD.encode(secret)
    ))
}

fn decrypt(
    repository: &Repository,
    keys: &KeyStore,
    path: &str,
    provider: &mut impl SecretProvider,
) -> Result<SecretBytes, Box<dyn std::error::Error>> {
    Ok(keys.decrypt(&repository.read_entry(&EntryPath::parse(path)?)?, provider)?)
}