IronStorage/crates/storage/tests/kdbx_import.rs

#![forbid(unsafe_code)]

mod support;

use std::{collections::BTreeMap, fs, io::Cursor};

use ironstorage::{
    crypto::{KeyInfo, KeyStore, SecretProvider, SecretProviderError},
    git::GitIdentity,
    kdbx::{KdbxImportMode, KdbxImportRequest, KdbxImporter},
    repository::{Repository, SecretBytes},
    write::VaultWriter,
};
use keepass::{
    Database, DatabaseKey,
    db::{AutoType, Value, fields},
};
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)
    }
}

#[test]
fn imports_adds_updates_and_quick_adds_without_deleting() -> 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 mut secrets = FixtureSecrets::all(&fixture);
    let source = store.path().join("source.kdbx");
    let unrelated = repository.read_entry(&ironstorage::repository::EntryPath::parse(
        "email/personal",
    )?)?;
    save_database(
        &source,
        "database password",
        "first password",
        "line one\nwww.bahn.de Login:\n\nend",
        b"first attachment",
        None,
    )?;
    let request = KdbxImportRequest::new(&source, None, KdbxImportMode::AddAndUpdate);
    let importer = KdbxImporter::new(&repository, &keys);

    let first = importer.import(
        &request,
        SecretBytes::new(b"database password".to_vec()),
        &mut secrets,
        GitIdentity::ironstorage(),
    )?;

    assert_eq!((first.added(), first.updated()), (2, 0));
    assert_eq!(
        plaintext(&repository, &keys, &mut secrets, "Work/Example")?,
        b"first password\nlogin: alice\nurl: https://example.test\notpauth://totp/Example:alice?secret=JBSWY3DPEHPK3PXP&issuer=Example\ncomments: line one\n www.bahn.de Login:\n \n end\nattachments: manual.txt\nautotype_enabled: true\nautotype_sequence: {USERNAME}{TAB}{PASSWORD}{ENTER}\nCustom_Label: alpha\n beta: still one field\n"
    );
    assert_eq!(
        plaintext(&repository, &keys, &mut secrets, "Work/manual.txt")?,
        b"first attachment"
    );

    let second = importer.import(
        &request,
        SecretBytes::new(b"database password".to_vec()),
        &mut secrets,
        GitIdentity::ironstorage(),
    )?;
    assert_eq!(
        (second.unchanged(), second.added(), second.updated()),
        (2, 0, 0)
    );

    save_database(
        &source,
        "database password",
        "updated password",
        "updated\ncomments",
        b"updated attachment",
        None,
    )?;
    let updated = importer.import(
        &request,
        SecretBytes::new(b"database password".to_vec()),
        &mut secrets,
        GitIdentity::ironstorage(),
    )?;
    assert_eq!((updated.updated(), updated.added()), (2, 0));
    assert!(
        plaintext(&repository, &keys, &mut secrets, "Work/Example")?
            .starts_with(b"updated password\n")
    );

    save_database(
        &source,
        "database password",
        "ignored by quick add",
        "ignored",
        b"ignored",
        None,
    )?;
    let quick = importer.import(
        &KdbxImportRequest::new(&source, None, KdbxImportMode::QuickAdd),
        SecretBytes::new(b"database password".to_vec()),
        &mut secrets,
        GitIdentity::ironstorage(),
    )?;
    assert_eq!((quick.skipped(), quick.added(), quick.updated()), (2, 0, 0));
    assert!(
        plaintext(&repository, &keys, &mut secrets, "Work/Example")?
            .starts_with(b"updated password\n")
    );
    assert_eq!(
        repository.read_entry(&ironstorage::repository::EntryPath::parse(
            "email/personal"
        )?)?,
        unrelated
    );
    Ok(())
}

#[test]
fn password_and_key_file_are_both_required_before_any_store_change() -> 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 mut secrets = FixtureSecrets::all(&fixture);
    let source = store.path().join("protected.kdbx");
    let key_file = store.path().join("database.key");
    fs::write(&key_file, b"key file material")?;
    save_database(
        &source,
        "password",
        "secret",
        "notes",
        b"attachment",
        Some(b"key file material"),
    )?;
    let request = KdbxImportRequest::new(&source, Some(key_file), KdbxImportMode::AddAndUpdate);
    let importer = KdbxImporter::new(&repository, &keys);

    assert!(
        importer
            .import(
                &request,
                SecretBytes::new(b"wrong".to_vec()),
                &mut secrets,
                GitIdentity::ironstorage(),
            )
            .is_err()
    );
    assert!(!VaultWriter::new(&repository, &keys).entry_exists("Work/Example")?);

    let outcome = importer.import(
        &request,
        SecretBytes::new(b"password".to_vec()),
        &mut secrets,
        GitIdentity::ironstorage(),
    )?;
    assert_eq!(outcome.added(), 2);
    Ok(())
}

fn save_database(
    path: &std::path::Path,
    database_password: &str,
    entry_password: &str,
    notes: &str,
    attachment: &[u8],
    key_file: Option<&[u8]>,
) -> TestResult {
    let mut database = Database::new();
    {
        let mut root = database.root_mut();
        let mut group = root.add_group();
        group.name = "Work".to_owned();
        let mut entry = group.add_entry();
        entry.set_unprotected(fields::TITLE, "Example");
        entry.set_unprotected(fields::USERNAME, "alice");
        entry.set_protected(fields::PASSWORD, entry_password);
        entry.set_unprotected(fields::URL, "https://example.test");
        entry.set_unprotected(fields::NOTES, notes);
        entry.set_protected(
            fields::OTP,
            "otpauth://totp/Example:alice?secret=JBSWY3DPEHPK3PXP&issuer=Example",
        );
        entry.set_unprotected("Custom Label", "alpha\nbeta: still one field");
        entry.autotype = Some(AutoType {
            enabled: true,
            default_sequence: Some("{USERNAME}{TAB}{PASSWORD}{ENTER}".to_owned()),
            ..AutoType::default()
        });
        entry.add_attachment("manual.txt", Value::protected(attachment.to_vec()));
    }
    let mut key = DatabaseKey::new().with_password(database_password);
    if let Some(key_file) = key_file {
        key = key.with_keyfile(&mut Cursor::new(key_file))?;
    }
    database.save(&mut fs::File::create(path)?, key)?;
    Ok(())
}

fn plaintext(
    repository: &Repository,
    keys: &KeyStore,
    secrets: &mut impl SecretProvider,
    entry: &str,
) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
    Ok(VaultWriter::new(repository, keys)
        .begin_edit(
            &ironstorage::command::EditRequest {
                entry: entry.to_owned(),
            },
            secrets,
        )?
        .plaintext()
        .expose()
        .to_vec())
}