mod ctor_migrate_util;
mod key_util;

use crate::hsc::ctor_migrate_util::{
    CTOR_KEYSTORE_ID, CTOR_KEYSTORE1_PATH, CTOR_KEYSTORE2_PATH, CTOR_KEYSTORE3_PATH,
    CTOR_KEYSTORE4_PATH, CTOR_KEYSTORE5_PATH, CTorMigrateCmd, ONION_ADDR_SERVICE_1,
    ONION_ADDR_SERVICE_2,
};

use crate::hsc::key_util::{ADDR_LEN, ArtiHscKeyCmd, KeyCmdBuilder, ONION_ADDR};

/// Error message displayed by `key get` when the service discovery key is not found and the
/// `--generate=no` flag is provided.
const ERR_MSG: &str = "Service discovery key not found. Rerun with --generate=if-needed to generate a new service discovery keypair";

#[test]
fn gen_key() {
    let state_dir = tempfile::TempDir::new().unwrap();
    let state_dir = state_dir.path();

    let cmd = KeyCmdBuilder::default()
        .subcommand(ArtiHscKeyCmd::Get { generate: true })
        .state_dir(state_dir.to_path_buf())
        .build()
        .unwrap();

    let output = cmd.run().unwrap();

    assert!(output.status.success());
    assert!(
        String::from_utf8(output.stdout)
            .unwrap()
            .contains("descriptor:x25519:")
    );

    // Assert new private key has been generated
    assert!(cmd.keystore_contains_priv_key(ONION_ADDR));
}

/// The `arti hsc key get` command fails if no valid keys exist for the given onion address
/// and the `--generate=no` flag is specified.
#[test]
fn get_fails_if_no_gen_and_no_key() {
    let state_dir = tempfile::TempDir::new().unwrap();
    let state_dir = state_dir.path();

    let cmd = KeyCmdBuilder::default()
        .subcommand(ArtiHscKeyCmd::Get { generate: false })
        .state_dir(state_dir.to_path_buf())
        .build()
        .unwrap();

    let output = cmd.run().unwrap();
    assert!(!output.status.success());
    assert!(String::from_utf8(output.stderr).unwrap().contains(ERR_MSG));
}

#[test]
fn generate_then_rotate() {
    let state_dir = tempfile::TempDir::new().unwrap();
    let state_dir = state_dir.path();

    let cmd = KeyCmdBuilder::default()
        .subcommand(ArtiHscKeyCmd::Get { generate: true })
        .state_dir(state_dir.to_path_buf())
        .build()
        .unwrap();
    let output = cmd.run().unwrap();

    assert!(output.status.success());
    let descriptor = String::from_utf8(output.stdout).unwrap();
    assert!(descriptor.contains("descriptor:x25519:"));

    let cmd = KeyCmdBuilder::default()
        .subcommand(ArtiHscKeyCmd::Rotate)
        .state_dir(state_dir.to_path_buf())
        .build()
        .unwrap();
    let output = cmd.run().unwrap();

    assert!(output.status.success());
    let rotated_descriptor = String::from_utf8(output.stdout).unwrap();
    assert!(rotated_descriptor.contains("descriptor:x25519:"));

    // Assert key has been rotated
    assert_ne!(descriptor, rotated_descriptor);

    // Assert new private key has been generated
    assert!(cmd.keystore_contains_priv_key(ONION_ADDR));
}

#[test]
fn generate_then_remove() {
    let state_dir = tempfile::TempDir::new().unwrap();
    let state_dir = state_dir.path();

    let cmd = KeyCmdBuilder::default()
        .subcommand(ArtiHscKeyCmd::Get { generate: true })
        .state_dir(state_dir.to_path_buf())
        .build()
        .unwrap();

    let output = cmd.run().unwrap();
    assert!(output.status.success());
    assert!(
        String::from_utf8(output.stdout)
            .unwrap()
            .contains("descriptor:x25519:")
    );

    let cmd = KeyCmdBuilder::default()
        .subcommand(ArtiHscKeyCmd::Remove)
        .state_dir(state_dir.to_path_buf())
        .build()
        .unwrap();
    let output = cmd.run().unwrap();
    assert!(output.status.success());

    let keystore_path = state_dir
        .join("keystore/client")
        .join(&ONION_ADDR[..ADDR_LEN]);
    let entries = keystore_path.read_dir().unwrap().flatten();
    // Assert key has been removed
    assert_eq!(entries.count(), 0);
}

/// Tests whether `ctor-migrate` can successfully migrate a valid C Tor client keystore
/// to an Arti keystore that does not already have discovery keys for the services
/// the C Tor client is configured with.
#[test]
fn simple_ctor_migration() {
    let migrate_cmd = CTorMigrateCmd::new();

    let assert_key_is_missing = |svc: &str| {
        let err = migrate_cmd.keystore_contains_client_key(svc).unwrap_err();
        assert!(err.to_string().contains(ERR_MSG));
    };

    // The client keystore doesn't have keys for either of the two services before the migration
    assert_key_is_missing(ONION_ADDR_SERVICE_1);
    assert_key_is_missing(ONION_ADDR_SERVICE_2);

    let output = migrate_cmd.output(CTOR_KEYSTORE1_PATH).unwrap();
    assert!(output.status.success());
    assert!(
        migrate_cmd
            .keystore_contains_client_key(ONION_ADDR_SERVICE_1)
            .is_ok()
    );
    assert!(
        migrate_cmd
            .keystore_contains_client_key(ONION_ADDR_SERVICE_2)
            .is_ok()
    );
}

/// Tests whether `ctor-migrate` fails if multiple client keys for the same service are present in
/// the C Tor keystore. A C Tor keystore must not contain multiple client keys for the same
/// service.
#[test]
fn migrate_duplicate_ctor_entries() {
    let migrate_cmd = CTorMigrateCmd::new();
    let output = migrate_cmd.output(CTOR_KEYSTORE2_PATH).unwrap();
    assert!(!output.status.success());
    assert!(
        String::from_utf8(output.stderr)
            .unwrap()
            .contains("Invalid C Tor keystore (multiple keys exist for service")
    );
}

/// Tests whether `ctor-migrate` succeeds when client keys for the same service are present in both
/// the Arti primary keystore and in the C Tor keystore being migrated. It also verifies that the
/// keys in the Arti primary keystore differ from the original ones after migration.
#[test]
fn forced_migration_overwrites_arti_keys() {
    let migrate_cmd = CTorMigrateCmd::new();
    migrate_cmd.populate_state_dir();
    // Ensure the keystore is now populated with the expected keys
    let output_service_1_prev = migrate_cmd
        .keystore_contains_client_key(ONION_ADDR_SERVICE_1)
        .unwrap();
    let output_service_2_prev = migrate_cmd
        .keystore_contains_client_key(ONION_ADDR_SERVICE_2)
        .unwrap();

    assert!(
        migrate_cmd
            .output(CTOR_KEYSTORE1_PATH)
            .unwrap()
            .status
            .success()
    );

    // Check whether the current keys are different from the original ones
    let output_service_1_current = migrate_cmd
        .keystore_contains_client_key(ONION_ADDR_SERVICE_1)
        .unwrap();
    let output_service_2_current = migrate_cmd
        .keystore_contains_client_key(ONION_ADDR_SERVICE_2)
        .unwrap();
    assert_ne!(output_service_1_prev, output_service_1_current);
    assert_ne!(output_service_2_prev, output_service_2_current);
}

/// Tests that `ctor-migrate` fails when there are no valid entries in the
/// registered C Tor keystore, then test `ctor-migrate` fails when there
/// are no valid keys in the registered C Tor keystore.
#[test]
fn migrate_invalid_ctor_keystore() {
    let assert_cmd_fails = |path: &str| {
        let migrate_cmd = CTorMigrateCmd::new();

        let output = migrate_cmd.output(path).unwrap();
        assert!(!output.status.success());
        let error = String::from_utf8(output.stderr).unwrap();
        assert!(error.contains("No CTor client keys found in keystore"));
        assert!(error.contains(CTOR_KEYSTORE_ID));
    };

    // NOTE: with CTOR_KEYSTORE3_PATH the keystore has a valid entry containing an invalid key.
    // With CTOR_KEYSTORE4_PATH the keystore has an invalid entry containing a valid key.
    // Both cases should fail.
    assert_cmd_fails(CTOR_KEYSTORE3_PATH);
    assert_cmd_fails(CTOR_KEYSTORE4_PATH);
}

/// Tests whether `ctor-migrate` succeeds when both valid and invalid entries are present in the
/// registered C Tor keystore.
#[test]
fn migrate_skips_invalid_ctor_entries() {
    let migrate_cmd = CTorMigrateCmd::new();

    let output = migrate_cmd.output(CTOR_KEYSTORE5_PATH).unwrap();
    assert!(output.status.success());

    assert!(
        migrate_cmd
            .keystore_contains_client_key(ONION_ADDR_SERVICE_1)
            .is_ok()
    );
}
