//! Key rotation tasks of the relay.

use anyhow::Context;
use std::{
    sync::Arc,
    time::{Duration, SystemTime},
};
use tor_basic_utils::rand_hostname;
use tor_cert::x509::TlsKeyAndCert;
use tor_chanmgr::ChanMgr;
use tor_proto::RelayIdentities;

use tor_key_forge::ToEncodableCert;
use tor_keymgr::{
    KeyMgr, KeyPath, KeySpecifier, KeySpecifierPattern, Keygen, KeystoreSelector, ToEncodableKey,
};
use tor_relay_crypto::{
    gen_link_cert, gen_signing_cert, gen_tls_cert,
    pk::{
        RelayIdentityKeypair, RelayIdentityKeypairSpecifier, RelayIdentityRsaKeypair,
        RelayIdentityRsaKeypairSpecifier, RelayLinkSigningKeypair,
        RelayLinkSigningKeypairSpecifier, RelayLinkSigningKeypairSpecifierPattern,
        RelaySigningKeypair, RelaySigningKeypairSpecifier, RelaySigningKeypairSpecifierPattern,
        Timestamp,
    },
};
use tor_rtcompat::{Runtime, SleepProviderExt};

/// Buffer time before key expiry to trigger rotation. This ensures we rotate slightly before the
/// key actually expires rather than right at or after expiry.
///
/// C-tor uses 3 hours for the link/auth key and 1 day for the signing key. Let's use 3 hours here,
/// it should be plenty to make it happen even if hiccups happen.
const KEY_ROTATION_EXPIRE_BUFFER: Duration = Duration::from_secs(3 * 60 * 60);

/// Key lifefime duration of 2 days
const KEY_DURATION_2DAYS: Duration = Duration::from_secs(2 * 24 * 60 * 60);
/// Key lifefime duration of 30 days
const KEY_DURATION_30DAYS: Duration = Duration::from_secs(30 * 24 * 60 * 60);
/// Key lifefime duration of 6 months
const KEY_DURATION_6MONTHS: Duration = Duration::from_secs(6 * 30 * 24 * 60 * 60);

/// Trait to help us specify what we need for key rotation. This allows us to have the generic
/// function `rotate_key()`.
trait RotatableKeySpec {
    /// Key specifier type.
    type Specifier: KeySpecifier;
    /// Key specifier pattern (for the ArtiPath).
    type Pattern: KeySpecifierPattern;

    /// Build a new specifier.
    fn key_specifier() -> Self::Specifier;
    /// For logs.
    fn label() -> &'static str;
    /// Build a specifier from a [`KeyPath`]
    fn spec_from_keypath(keypath: &KeyPath) -> Result<Self::Specifier, tor_keymgr::KeyPathError>;
    /// The `valid_until` of the given key specifier.
    fn valid_until_from_spec(spec: &Self::Specifier) -> Timestamp;
}

impl RotatableKeySpec for RelaySigningKeypair {
    type Specifier = RelaySigningKeypairSpecifier;
    type Pattern = RelaySigningKeypairSpecifierPattern;

    fn key_specifier() -> Self::Specifier {
        let valid_until = Timestamp::from(SystemTime::now() + Duration::from_secs(30 * 86400));
        Self::Specifier::new(valid_until)
    }
    fn label() -> &'static str {
        "KP_relaysign_ed"
    }
    fn spec_from_keypath(keypath: &KeyPath) -> Result<Self::Specifier, tor_keymgr::KeyPathError> {
        keypath.try_into()
    }
    fn valid_until_from_spec(spec: &Self::Specifier) -> Timestamp {
        spec.valid_until()
    }
}

impl RotatableKeySpec for RelayLinkSigningKeypair {
    type Specifier = RelayLinkSigningKeypairSpecifier;
    type Pattern = RelayLinkSigningKeypairSpecifierPattern;

    fn key_specifier() -> Self::Specifier {
        let valid_until = Timestamp::from(SystemTime::now() + Duration::from_secs(2 * 86400));
        Self::Specifier::new(valid_until)
    }
    fn label() -> &'static str {
        "KP_link_ed"
    }
    fn spec_from_keypath(keypath: &KeyPath) -> Result<Self::Specifier, tor_keymgr::KeyPathError> {
        keypath.try_into()
    }
    fn valid_until_from_spec(spec: &Self::Specifier) -> Timestamp {
        spec.valid_until()
    }
}

/// Generate a key `K` directly into the key manager.
///
/// If the key already exists, the error is ignored as this could happen if the system time drifts
/// between the get and the generate.
fn generate_key<K>(keymgr: &KeyMgr, spec: &dyn KeySpecifier) -> Result<(), tor_keymgr::Error>
where
    K: ToEncodableKey,
    K::Key: Keygen,
{
    let mut rng = tor_llcrypto::rng::CautiousRng;

    match keymgr.generate::<K>(spec, KeystoreSelector::default(), &mut rng, false) {
        Ok(_) => {}
        // Key already existing can happen due to wall clock strangeness,
        // so simply ignore it.
        Err(tor_keymgr::Error::KeyAlreadyExists) => tracing::warn!(
            "Failed to generate key at {:?} because one already exists. Clock drift?",
            spec.arti_path(),
        ),
        Err(e) => return Err(e),
    };
    Ok(())
}

/// Rotate a key implementing the [`RotatableKeySpec`] trait.
///
/// Rotation is done by listing all keys matching the key specifier pattern and validating the
/// valid_until value of the key store entry. If expired, the key is removed from the key manager.
///
/// Returns a tuple of (rotated, valid_until) where `rotated` indicates if the key was rotated and
/// `valid_until` is the earliest expiry time across all keys of this type.
fn rotate_key<K>(keymgr: &KeyMgr) -> anyhow::Result<(bool, SystemTime)>
where
    K: RotatableKeySpec + ToEncodableKey,
    <K as ToEncodableKey>::Key: Keygen,
{
    // Select all signing keypair in the keystore because we need to inspect the valid_until
    // field and rotate if expired.
    let key_entries = keymgr.list_matching(&K::Pattern::new_any().arti_pattern()?)?;

    let key_specifier = K::key_specifier();

    if key_entries.is_empty() {
        generate_key::<K>(keymgr, &key_specifier)?;
        let valid_until = K::valid_until_from_spec(&key_specifier).into();
        return Ok((true, valid_until));
    }

    let mut have_rotated = false;
    // Smallest valid_until timestamp of all the keys we are about to look at. Start with the
    // biggest value so the first value will change this immediately.
    //
    // NOTE: This is dicy because if the loop below would not run, we would return a sleep time
    // that is massive. The is_empty() above guarantees it won't happen but still. Anyway, this is
    // better than an Option<> and dealing with a None at the end.
    let mut min_valid_until = None;

    for key in key_entries {
        let entry_key_spec: K::Specifier = K::spec_from_keypath(key.key_path())?;
        // Min the entry key valid_until.
        let entry_valid_until = K::valid_until_from_spec(&entry_key_spec);
        min_valid_until = Some(
            min_valid_until.map_or(entry_valid_until, |v: Timestamp| v.min(entry_valid_until)),
        );

        // Account for the buffer time so we rotate before the key actually expires.
        if K::valid_until_from_spec(&entry_key_spec)
            <= Timestamp::from(SystemTime::now() + KEY_ROTATION_EXPIRE_BUFFER)
        {
            tracing::info!(
                "Rotating {} key. Next expiry timestamp {:?}",
                K::label(),
                K::valid_until_from_spec(&key_specifier),
            );
            keymgr.remove_entry(&key)?;
            generate_key::<K>(keymgr, &key_specifier)?;
            have_rotated = true;
            // Min the new key valid_until.
            let new_valid_until = K::valid_until_from_spec(&key_specifier);
            min_valid_until =
                Some(min_valid_until.map_or(new_valid_until, |v| v.min(new_valid_until)));
        }
    }

    Ok((
        have_rotated,
        min_valid_until.expect("valid_until is empty").into(),
    ))
}

/// Attempt to rotate all rotatable keys.
///
/// Returns a tuple of (rotated, next_expiry) where `rotated` indicates if any key was rotated and
/// `next_expiry` is the earliest expiry time across all rotatable keys.
fn try_rotate_keys(keymgr: &KeyMgr) -> anyhow::Result<(bool, SystemTime)> {
    // Attempt to rotate the KP_relaysign_ed.
    let (mut have_rotated, sign_expiry) = rotate_key::<RelaySigningKeypair>(keymgr)?;
    // Attempt to rotate the KP_link_ed.
    let (link_rotated, link_expiry) = rotate_key::<RelayLinkSigningKeypair>(keymgr)?;
    have_rotated |= link_rotated;
    Ok((have_rotated, sign_expiry.min(link_expiry)))
}

/// Build a fresh [`RelayIdentities`] object using a [`KeyMgr`].
///
/// Every single certificate is generated in this function.
///
/// This function assumes that all required keys are in the keymgr.
fn build_proto_identities(keymgr: &KeyMgr) -> anyhow::Result<RelayIdentities> {
    let mut rng = tor_llcrypto::rng::CautiousRng;
    let now = SystemTime::now();

    // Get the identity keypairs.
    let rsa_id_kp: RelayIdentityRsaKeypair = keymgr
        .get(&RelayIdentityRsaKeypairSpecifier::new())
        .context("Failed to get RSA identity from key manager")?
        .context("Missing RSA identity")?;
    let ed_id_kp: RelayIdentityKeypair = keymgr
        .get(&RelayIdentityKeypairSpecifier::new())
        .context("Failed to get Ed25519 identity from key manager")?
        .context("Missing Ed25519 identity")?;
    // We have to list match here because the key specifier here uses a valid_until. We don't know
    // what it is so we list and take the first one.
    let link_sign_kp: RelayLinkSigningKeypair = keymgr
        .get_entry(
            keymgr
                .list_matching(&RelayLinkSigningKeypairSpecifierPattern::new_any().arti_pattern()?)?
                .first()
                .context("No store entry for link authentication key")?,
        )
        .context("Failed to get link authentication key from key manager")?
        .context("Missing link authentication key")?;
    let kp_relaysign_id: RelaySigningKeypair = keymgr
        .get_entry(
            keymgr
                .list_matching(&RelaySigningKeypairSpecifierPattern::new_any().arti_pattern()?)?
                .first()
                .context("No store entry for signing key")?,
        )
        .context("Failed to get signing key from key manager")?
        .context("Missing signing key")?;

    // TLS key and cert. Random hostname like C-tor. We re-use the issuer_hostname for the RSA
    // legacy cert.
    let issuer_hostname = rand_hostname::random_hostname(&mut rng);
    let subject_hostname = rand_hostname::random_hostname(&mut rng);
    let tls_key_and_cert =
        TlsKeyAndCert::create(&mut rng, now, &issuer_hostname, &subject_hostname)
            .context("Failed to create TLS keys and certificates")?;

    // Create the RSA X509 certificate.
    let cert_id_x509_rsa = tor_cert::x509::create_legacy_rsa_id_cert(
        &mut rng,
        SystemTime::now(),
        &issuer_hostname,
        rsa_id_kp.keypair(),
    )
    .context("Failed to create legacy RSA identity certificate")?;

    // The following expiry duration have been taken from C-tor.

    let cert_id_rsa = tor_cert::rsa::EncodedRsaCrosscert::encode_and_sign(
        rsa_id_kp.keypair(),
        &ed_id_kp.to_ed25519_id(),
        now + KEY_DURATION_6MONTHS,
    )?;

    // Create the signing key cert, link cert and tls cert.
    //
    // TODO(relay): We need to check the KeyMgr for the signing cert but for now the KeyMgr API
    // doesn't allow us to get it out. We will do a re-design of the cert API there. This is fine
    // as long as we don't support offline keys.
    let cert_id_sign_ed = gen_signing_cert(&ed_id_kp, &kp_relaysign_id, now + KEY_DURATION_30DAYS)?;
    let cert_sign_link_auth_ed =
        gen_link_cert(&kp_relaysign_id, &link_sign_kp, now + KEY_DURATION_2DAYS)?;
    let cert_sign_tls_ed = gen_tls_cert(
        &kp_relaysign_id,
        *tls_key_and_cert.link_cert_sha256(),
        now + KEY_DURATION_2DAYS,
    )?;

    Ok(RelayIdentities::new(
        &rsa_id_kp.public().into(),
        ed_id_kp.to_ed25519_id(),
        link_sign_kp,
        cert_id_sign_ed.to_encodable_cert(),
        cert_sign_tls_ed,
        cert_sign_link_auth_ed.to_encodable_cert(),
        cert_id_x509_rsa,
        cert_id_rsa,
        tls_key_and_cert,
    ))
}

/// Attempt to generate all keys. The list of keys is:
///
/// * Identity Ed25519 keypair [`RelayIdentityKeypair`].
/// * Identity RSA [`RelayIdentityRsaKeypair`].
/// * Relay signing keypair [`RelaySigningKeypair`].
/// * Relay link signing keypair [`RelayLinkSigningKeypair`].
///
/// This function is only called when our relay bootstraps in order to attempt to generate any
/// missing keys or/and rotate expired keys.
pub(crate) fn try_generate_keys(keymgr: &KeyMgr) -> anyhow::Result<RelayIdentities> {
    // Note that generate_key() won't error if the key already exists.

    // Attempt to generate our identity keys (ed and RSA). Those keys DO NOT rotate.
    generate_key::<RelayIdentityKeypair>(keymgr, &RelayIdentityKeypairSpecifier::new())?;
    generate_key::<RelayIdentityRsaKeypair>(keymgr, &RelayIdentityRsaKeypairSpecifier::new())?;
    // Attempt to rotate the rotatable keys which will generate any missing.
    let _expiry = try_rotate_keys(keymgr)?;

    // Now that we have our up-to-date keys, build the RelayIdentities object.
    build_proto_identities(keymgr)
}

/// Task to rotate keys when they need to be rotated.
pub(crate) async fn rotate_keys_task<R: Runtime>(
    runtime: R,
    keymgr: Arc<KeyMgr>,
    chanmgr: Arc<ChanMgr<R>>,
) -> anyhow::Result<void::Void> {
    loop {
        // Attempt a rotation of all keys.
        let (have_rotated, next_expiry) = try_rotate_keys(&keymgr)?;
        if have_rotated {
            let ids = build_proto_identities(&keymgr)?;
            chanmgr
                .set_relay_identities(Arc::new(ids))
                .context("Failed to set relay identities on ChanMgr")?;
        }

        // Sleep until the earliest key expiry minus buffer so we rotate before it expires.
        // If the subtraction would underflow, wake up immediately to rotate the expired key.
        let next_wake = next_expiry
            .checked_sub(KEY_ROTATION_EXPIRE_BUFFER)
            .unwrap_or(SystemTime::now());
        runtime.sleep_until_wallclock(next_wake).await;
    }
}