//! Two-factor helpers: TOTP (RFC 6238) over HMAC-SHA1 with 6 digits and a 30 second step,
//! RFC 4648 base32 for the shared secret, and recovery codes.
//!
//! Codes are checked at the current step and one step either side. The caller records the
//! matched step so that a code is accepted only once (see `Store::accept_totp_step`).

use hmac::{Hmac, KeyInit, Mac};
use sha1::Sha1;
use sha2::{Digest, Sha256};

use crate::error::{Error, Result};
use crate::token::{constant_time_eq, random_readable};

/// Shared secret length in bytes (160 bits, the RFC 4226 recommendation).
pub const SECRET_LEN: usize = 20;
/// Seconds per time step.
pub const STEP_SECS: u64 = 30;
/// Digits in a code.
pub const DIGITS: usize = 6;
/// Steps accepted on each side of the current one.
pub const SKEW_STEPS: i64 = 1;

const BASE32_ALPHABET: &[u8; 32] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZ234567";

type HmacSha1 = Hmac<Sha1>;

/// A fresh random secret.
pub fn new_secret() -> [u8; SECRET_LEN] {
    let mut secret = [0u8; SECRET_LEN];
    rand::fill(&mut secret);
    secret
}

/// HOTP (RFC 4226) with HMAC-SHA1, reduced to [`DIGITS`] digits.
pub fn hotp(secret: &[u8], counter: u64) -> u32 {
    // HMAC accepts keys of any length, so construction cannot fail.
    let mut mac = <HmacSha1 as KeyInit>::new_from_slice(secret).expect("HMAC takes any key length");
    mac.update(&counter.to_be_bytes());
    let tag = mac.finalize().into_bytes();
    let offset = usize::from(tag[tag.len() - 1] & 0x0f);
    let binary = u32::from_be_bytes([
        tag[offset] & 0x7f,
        tag[offset + 1],
        tag[offset + 2],
        tag[offset + 3],
    ]);
    binary % 10u32.pow(DIGITS as u32)
}

/// The time step containing `unix_secs`.
pub fn step_for(unix_secs: u64) -> u64 {
    unix_secs / STEP_SECS
}

/// The zero-padded code for one step.
pub fn code_at(secret: &[u8], step: u64) -> String {
    format!("{:0width$}", hotp(secret, step), width = DIGITS)
}

/// Checks `code` against the steps around `now_step`. Returns the matched step when it is
/// later than `last_step`, which is the replay guard. Any malformed code returns `None`.
pub fn verify(secret: &[u8], code: &str, now_step: u64, last_step: u64) -> Option<u64> {
    if code.len() != DIGITS || !code.bytes().all(|b| b.is_ascii_digit()) {
        return None;
    }
    let now = i64::try_from(now_step).ok()?;
    let mut matched = None;
    // Every candidate is compared, so timing does not reveal which step matched.
    for delta in -SKEW_STEPS..=SKEW_STEPS {
        let Some(step) = now.checked_add(delta).and_then(|s| u64::try_from(s).ok()) else {
            continue;
        };
        let expected = code_at(secret, step);
        if constant_time_eq(expected.as_bytes(), code.as_bytes()) && matched.is_none() {
            matched = Some(step);
        }
    }
    matched.filter(|&step| step > last_step)
}

/// RFC 4648 base32, uppercase, without padding. This is the form authenticator apps expect.
pub fn base32_encode(data: &[u8]) -> String {
    let mut out = String::with_capacity(data.len().div_ceil(5) * 8);
    let mut buffer: u64 = 0;
    let mut bits: u32 = 0;
    for &byte in data {
        buffer = (buffer << 8) | u64::from(byte);
        bits += 8;
        while bits >= 5 {
            bits -= 5;
            out.push(char::from(
                BASE32_ALPHABET[((buffer >> bits) & 31) as usize],
            ));
        }
    }
    if bits > 0 {
        out.push(char::from(
            BASE32_ALPHABET[((buffer << (5 - bits)) & 31) as usize],
        ));
    }
    out
}

/// Decodes base32 case-insensitively. Spaces and `=` padding are ignored.
pub fn base32_decode(input: &str) -> Result<Vec<u8>> {
    let mut out = Vec::with_capacity(input.len() * 5 / 8);
    let mut buffer: u64 = 0;
    let mut bits: u32 = 0;
    for c in input.chars() {
        if c.is_whitespace() || c == '=' {
            continue;
        }
        let upper = c.to_ascii_uppercase();
        let value = BASE32_ALPHABET
            .iter()
            .position(|&a| char::from(a) == upper)
            .ok_or_else(|| Error::Invalid("secret is not valid base32".into()))?;
        buffer = (buffer << 5) | value as u64;
        bits += 5;
        if bits >= 8 {
            bits -= 8;
            out.push(((buffer >> bits) & 0xff) as u8);
        }
    }
    Ok(out)
}

/// The `otpauth://` URI that authenticator apps import. `account` is the username.
pub fn otpauth_uri(account: &str, secret: &[u8]) -> String {
    format!(
        "otpauth://totp/forge:{account}?secret={}&issuer=forge",
        base32_encode(secret)
    )
}

/// A new recovery code, shown as `xxxx-xxxx`.
pub fn new_recovery_code() -> String {
    let raw = random_readable(8);
    format!("{}-{}", &raw[..4], &raw[4..])
}

/// Hex sha256 of a recovery code after normalization: lowercase, with dashes and spaces removed.
/// Used both when storing and when checking, so the typed form does not matter.
pub fn recovery_code_hash(code: &str) -> String {
    let normalized: String = code
        .chars()
        .filter(|c| !c.is_whitespace() && *c != '-')
        .map(|c| c.to_ascii_lowercase())
        .collect();
    hex::encode(Sha256::digest(normalized.as_bytes()))
}

#[cfg(test)]
mod tests {
    use super::*;

    /// RFC 6238 appendix B, SHA-1 row. The 8-digit codes reduce to the 6-digit ones below.
    const RFC_KEY: &[u8] = b"12345678901234567890";

    #[test]
    fn rfc6238_sha1_test_vectors() {
        let vectors: [(u64, &str); 6] = [
            (59, "287082"),
            (1_111_111_109, "081804"),
            (1_111_111_111, "050471"),
            (1_234_567_890, "005924"),
            (2_000_000_000, "279037"),
            (20_000_000_000, "353130"),
        ];
        for (time, expected) in vectors {
            assert_eq!(code_at(RFC_KEY, time / STEP_SECS), expected, "time {time}");
        }
    }

    #[test]
    fn rfc4226_hotp_test_vectors() {
        let expected = [
            755_224, 287_082, 359_152, 969_429, 338_314, 254_676, 287_922, 162_583, 399_871,
            520_489,
        ];
        for (counter, code) in expected.into_iter().enumerate() {
            assert_eq!(hotp(RFC_KEY, counter as u64), code, "counter {counter}");
        }
    }

    #[test]
    fn base32_rfc4648_vectors() {
        let cases = [
            ("", ""),
            ("f", "MY"),
            ("fo", "MZXQ"),
            ("foo", "MZXW6"),
            ("foob", "MZXW6YQ"),
            ("fooba", "MZXW6YTB"),
            ("foobar", "MZXW6YTBOI"),
        ];
        for (plain, encoded) in cases {
            assert_eq!(base32_encode(plain.as_bytes()), encoded);
            assert_eq!(base32_decode(encoded).unwrap(), plain.as_bytes());
            assert_eq!(
                base32_decode(&encoded.to_lowercase()).unwrap(),
                plain.as_bytes()
            );
        }
        assert!(base32_decode("MZ1Q").is_err());
    }

    #[test]
    fn verify_accepts_the_skew_window_and_returns_the_step() {
        let now = 1_000;
        for step in [now - 1, now, now + 1] {
            let code = code_at(RFC_KEY, step);
            assert_eq!(verify(RFC_KEY, &code, now, 0), Some(step));
        }
        let far = code_at(RFC_KEY, now + 2);
        assert_eq!(verify(RFC_KEY, &far, now, 0), None);
    }

    #[test]
    fn verify_rejects_steps_at_or_before_the_last_used_one() {
        let now = 1_000;
        let code = code_at(RFC_KEY, now);
        assert_eq!(verify(RFC_KEY, &code, now, now), None);
        assert_eq!(verify(RFC_KEY, &code, now, now + 1), None);
        assert_eq!(verify(RFC_KEY, &code, now, now - 1), Some(now));
    }

    #[test]
    fn verify_rejects_malformed_codes() {
        for bad in ["", "12345", "1234567", "12345a", " 28708", "-87082"] {
            assert_eq!(verify(RFC_KEY, bad, 1_000, 0), None, "input {bad:?}");
        }
    }

    #[test]
    fn otpauth_uri_carries_the_base32_secret() {
        let uri = otpauth_uri("alice", b"foobar");
        assert_eq!(
            uri,
            "otpauth://totp/forge:alice?secret=MZXW6YTBOI&issuer=forge"
        );
    }

    #[test]
    fn recovery_codes_are_dashed_and_hash_normalized() {
        let code = new_recovery_code();
        assert_eq!(code.len(), 9);
        assert_eq!(code.as_bytes()[4], b'-');
        assert_eq!(
            recovery_code_hash(&code),
            recovery_code_hash(&code.to_uppercase().replace('-', " "))
        );
        assert_ne!(
            recovery_code_hash("abcd-efgh"),
            recovery_code_hash("abcd-efgj")
        );
    }
}