//! LSX session crypto.
//!
//! Transcription of the crypto block of `fifa17-recon/tools/lsx_responder_v2.py`,
//! which carries the note:
//!
//! > (verbatim from v1 -- verified end-to-end by decrypting captured
//! >  captures/lsx/lsx_raw/C1_ENC-IN_*.bin.  DO NOT TOUCH.)
//!
//! So the Python is the specification and this file is a byte-for-byte port. The
//! block cipher itself comes from the `aes` crate; the framing around it
//! (PKCS7 pad to 16, lowercase hex, NUL terminator) is *protocol*, not
//! cryptography, and is therefore spelled out here rather than delegated to a
//! mode/padding helper.
//!
//! Wire recap (reversed from stp-origin_emu.dll @ 0x6ffffc930000):
//!   handshake: server sends `<Challenge key="...">` in PLAINTEXT; the client
//!              answers plaintext with `response=`/`key=`; the server answers
//!              `<ChallengeAccepted response="H">` where
//!              `H = hex(AES128_ECB(K_FIXED, clientKeyAscii))` + a fixed 3rd block.
//!   session:   every later frame, both directions, Responses AND Events, is
//!              `hex_lower(AES128_ECB(SESSION_KEY, pkcs7pad16(xml))) + b"\0"`,
//!              with SESSION_KEY derived from `H` through the MSVCR srand/rand LCG.

use std::fmt;
use std::sync::LazyLock;

use aes::cipher::generic_array::GenericArray;
use aes::cipher::{BlockDecrypt, BlockEncrypt, KeyInit};
use aes::Aes128;

/// A 128-bit LSX key (the fixed handshake key or a derived session key).
pub type Key = [u8; 16];

/// Recovered from stp-origin_emu.dll .rdata @ VA 0x6ffffc935038.
/// `bytes(range(16))` == `000102030405060708090a0b0c0d0e0f`.
pub const K_FIXED: Key = [
    0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f,
];

/// Failure modes of the LSX codec. The variants mirror the exceptions the Python
/// raises at the same points, because they surface through the same log lines
/// (`decrypt fail: ...`, `connection error: ...`).
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum CryptoError {
    /// ECB has no padding of its own: pycryptodome raises ValueError here.
    NotBlockAligned(usize),
    /// `bytes.fromhex` rejected the payload.
    BadHex,
    /// Nothing before the first NUL, so there is no block to decrypt (the Python
    /// dies on `raw[-1]` with IndexError).
    Empty,
    /// The client's own `response=` carried a 3rd block that is not the emu's
    /// constant. The Python asserts here on purpose: a future client that
    /// randomises block 3 must fail LOUDLY, not silently.
    UnexpectedTail(String),
}

impl fmt::Display for CryptoError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            // Wording kept close to pycryptodome's so existing log greps still hit.
            CryptoError::NotBlockAligned(n) => write!(
                f,
                "Data must be padded to 16 byte boundary in ECB mode (got {n} bytes)"
            ),
            CryptoError::BadHex => f.write_str("Non-hexadecimal digit found"),
            CryptoError::Empty => f.write_str("empty ciphertext"),
            CryptoError::UnexpectedTail(t) => {
                write!(f, "unexpected ChallengeResponse tail {t:?}")
            }
        }
    }
}

impl std::error::Error for CryptoError {}

/// MSVCR120 `srand`/`rand` LCG (verified: `srand(7); rand() == 61`).
pub struct MsvcrRand {
    state: u32,
}

impl MsvcrRand {
    /// `srand(seed)`.
    pub fn new(seed: u32) -> Self {
        Self { state: seed }
    }

    /// `rand()` -- 15 bits, exactly as MSVCR120 returns them.
    pub fn next_u15(&mut self) -> u16 {
        self.state = self.state.wrapping_mul(214013).wrapping_add(2531011);
        ((self.state >> 16) & 0x7FFF) as u16
    }
}

/// Reimplementation of emu `sub_0x6ffffc931f10` tail (0x9320bf-0x932101).
///
/// ```text
/// srand(7); r0 = rand()                       -> r0 == 61
/// bx  = (u16)((resp[0] << 8) + resp[1])       (16-bit wrap)
/// srand(bx + r0)
/// key[i] = (uint8_t)rand()   for i in 0..15
/// ```
///
/// `resp[0]`/`resp[1]` are the first two ASCII *characters* of the hex response,
/// not the first decoded byte.
pub fn derive_session_key(resp_hex: &str) -> Key {
    let r0 = u32::from(MsvcrRand::new(7).next_u15()); // == 61
    let b = resp_hex.as_bytes();
    assert!(
        b.len() >= 2,
        "derive_session_key needs the first two response chars, got {:?}",
        resp_hex
    );
    let bx = ((u32::from(b[0]) << 8) + u32::from(b[1])) & 0xFFFF;
    let mut g = MsvcrRand::new(bx.wrapping_add(r0));
    let mut key = [0u8; 16];
    for byte in key.iter_mut() {
        *byte = (g.next_u15() & 0xFF) as u8;
    }
    key
}

/// `AES128-ECB(K_FIXED, 0x10 * 16)` -- the constant the emu appends as the 3rd
/// hex block (== the PKCS7 pad block of an aligned 32-byte key). See
/// REPACK_INTEL.md sec.0-B.
pub fn tail_const() -> &'static str {
    static TAIL: LazyLock<String> = LazyLock::new(|| {
        let mut block = [0x10u8; 16];
        ecb_encrypt_blocks(&K_FIXED, &mut block);
        hex_lower(&block)
    });
    TAIL.as_str()
}

/// Emu-exact `ChallengeAccepted.response` (stp-origin_emu.dll 0x180001f10).
///
/// The emu computes only TWO AES blocks from the 32-ASCII client key, then
/// `strcat_s`'s the client's OWN `response[64:]` verbatim (@0x1800020a9) -> 96
/// hex. The older 3-block PKCS7 form is numerically identical *while the client
/// PKCS7-pads its 3rd block* (REPACK_INTEL.md sec.0-A/0-B, workflow-confirmed
/// byte-exact). We reproduce the emu exactly and, when the client's `response=`
/// is available, echo its tail and assert the constant so a future client that
/// randomises block 3 fails LOUDLY instead of silently.
pub fn challenge_response(
    client_key_ascii: &str,
    client_response_attr: &str,
) -> Result<String, CryptoError> {
    let mut buf = client_key_ascii.as_bytes().to_vec();
    if buf.is_empty() || buf.len() % 16 != 0 {
        // ECB cannot pad: the Python's AES.encrypt raises here too.
        return Err(CryptoError::NotBlockAligned(buf.len()));
    }
    ecb_encrypt_blocks(&K_FIXED, &mut buf);
    let two = hex_lower(&buf);

    // Python slices `client_response_attr[64:]`, i.e. by characters.
    if let Some((idx, _)) = client_response_attr.char_indices().nth(64) {
        let tail = &client_response_attr[idx..];
        if tail != tail_const() {
            return Err(CryptoError::UnexpectedTail(tail.to_string()));
        }
        return Ok(two + tail);
    }
    if client_response_attr.chars().count() == 64 {
        // len(attr) >= 64 with an empty tail: the assert compares "" against the
        // constant and fails, exactly as it would here.
        return Err(CryptoError::UnexpectedTail(String::new()));
    }
    Ok(two + tail_const())
}

/// pkcs7-pad to 16, AES-128-ECB, lowercase hex, NUL-terminated.
pub fn lsx_encrypt(xml: &str, key: &Key) -> Vec<u8> {
    let body = xml.as_bytes();
    let pad = 16 - (body.len() % 16); // emu always pads (pad==16 when aligned)
    let mut buf = Vec::with_capacity(body.len() + pad);
    buf.extend_from_slice(body);
    buf.resize(body.len() + pad, pad as u8);
    ecb_encrypt_blocks(key, &mut buf);

    // hex + NUL, sized up front: this runs on every frame we ever send.
    let mut out = Vec::with_capacity(buf.len() * 2 + 1);
    for b in &buf {
        out.push(HEX_DIGITS[usize::from(*b >> 4)]);
        out.push(HEX_DIGITS[usize::from(*b & 0x0f)]);
    }
    out.push(0);
    out
}

/// Hex up to the first NUL, decrypt, then take up to the first NUL of the
/// plaintext (the PKCS7 tail is stripped when it is well-formed, matching the
/// Python's tolerant check).
pub fn lsx_decrypt(data: &[u8], key: &Key) -> Result<String, CryptoError> {
    let head = match data.iter().position(|b| *b == 0) {
        Some(i) => &data[..i],
        None => data,
    };
    let mut raw = hex_decode(trim_ascii_ws(head))?;
    if raw.is_empty() {
        return Err(CryptoError::Empty);
    }
    if raw.len() % 16 != 0 {
        return Err(CryptoError::NotBlockAligned(raw.len()));
    }
    ecb_decrypt_blocks(key, &mut raw);

    let pad = usize::from(raw[raw.len() - 1]);
    if pad > 0 && pad <= 16 && raw[raw.len() - pad..].iter().all(|c| usize::from(*c) == pad) {
        raw.truncate(raw.len() - pad);
    }
    if let Some(i) = raw.iter().position(|b| *b == 0) {
        raw.truncate(i);
    }
    // Python decodes with errors="replace".
    Ok(String::from_utf8_lossy(&raw).into_owned())
}

const HEX_DIGITS: [u8; 16] = *b"0123456789abcdef";

/// Lowercase hex, like `bytes.hex()`.
pub fn hex_lower(bytes: &[u8]) -> String {
    let mut s = String::with_capacity(bytes.len() * 2);
    for b in bytes {
        s.push(char::from(HEX_DIGITS[usize::from(*b >> 4)]));
        s.push(char::from(HEX_DIGITS[usize::from(*b & 0x0f)]));
    }
    s
}

/// `bytes.fromhex`: ASCII whitespace between bytes is skipped, everything else
/// must be a hex digit pair.
fn hex_decode(s: &[u8]) -> Result<Vec<u8>, CryptoError> {
    let mut out = Vec::with_capacity(s.len() / 2);
    let mut hi: Option<u8> = None;
    for c in s {
        if c.is_ascii_whitespace() {
            continue;
        }
        let nib = match c {
            b'0'..=b'9' => c - b'0',
            b'a'..=b'f' => c - b'a' + 10,
            b'A'..=b'F' => c - b'A' + 10,
            _ => return Err(CryptoError::BadHex),
        };
        match hi.take() {
            None => hi = Some(nib),
            Some(h) => out.push((h << 4) | nib),
        }
    }
    if hi.is_some() {
        return Err(CryptoError::BadHex);
    }
    Ok(out)
}

/// `bytes.strip()` -- ASCII whitespace at both ends.
fn trim_ascii_ws(mut b: &[u8]) -> &[u8] {
    while let Some((first, rest)) = b.split_first() {
        if first.is_ascii_whitespace() {
            b = rest;
        } else {
            break;
        }
    }
    while let Some((last, rest)) = b.split_last() {
        if last.is_ascii_whitespace() {
            b = rest;
        } else {
            break;
        }
    }
    b
}

fn ecb_encrypt_blocks(key: &Key, buf: &mut [u8]) {
    debug_assert_eq!(buf.len() % 16, 0);
    let cipher = Aes128::new(GenericArray::from_slice(key));
    for block in buf.chunks_exact_mut(16) {
        cipher.encrypt_block(GenericArray::from_mut_slice(block));
    }
}

fn ecb_decrypt_blocks(key: &Key, buf: &mut [u8]) {
    debug_assert_eq!(buf.len() % 16, 0);
    let cipher = Aes128::new(GenericArray::from_slice(key));
    for block in buf.chunks_exact_mut(16) {
        cipher.decrypt_block(GenericArray::from_mut_slice(block));
    }
}

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

    /// The captured client key from the 2026-07-30 session the Python's own
    /// selftest pins.
    const CAPTURED_CLIENT_KEY: &str = "18a70055a3541fb27ab8e0f47afad18c";
    const CAPTURED_H: &str = concat!(
        "e4f5166209929e156a2ca47b81cdd6bf",
        "6f2a20371532ed4f968c5a9274899dbf",
        "954f64f2e4e86e9eee82d20216684899"
    );
    const CAPTURED_SESSION_KEY: &str = "6a9da3e78615153cc2f10eec25ae6382";

    #[test]
    fn msvcr_srand7_first_rand_is_61() {
        // The one vector the Python names in its docstring.
        assert_eq!(MsvcrRand::new(7).next_u15(), 61);
    }

    #[test]
    fn tail_const_is_the_emu_constant() {
        // AES128-ECB(K_FIXED, 0x10*16), pinned so a cipher/keying regression is
        // caught without a live client.
        assert_eq!(tail_const(), "954f64f2e4e86e9eee82d20216684899");
    }

    #[test]
    fn k_fixed_is_000102_to_0f() {
        assert_eq!(hex_lower(&K_FIXED), "000102030405060708090a0b0c0d0e0f");
    }

    #[test]
    fn challenge_response_matches_captured_session() {
        let h = challenge_response(CAPTURED_CLIENT_KEY, "").unwrap();
        assert!(h.starts_with("e4f5166209929e15"), "{h}");
        assert_eq!(h, CAPTURED_H);
        assert_eq!(h.len(), 96);
    }

    #[test]
    fn challenge_response_echoes_the_clients_own_tail() {
        // The emu strcat_s's response[64:] verbatim; identical result while the
        // client pads block 3 with the constant.
        let h = challenge_response(CAPTURED_CLIENT_KEY, CAPTURED_H).unwrap();
        assert_eq!(h, CAPTURED_H);
    }

    #[test]
    fn challenge_response_rejects_a_randomised_third_block() {
        let bogus = format!("{}{}", &CAPTURED_H[..64], "00".repeat(16));
        let err = challenge_response(CAPTURED_CLIENT_KEY, &bogus).unwrap_err();
        assert_eq!(err, CryptoError::UnexpectedTail("00".repeat(16)));
    }

    #[test]
    fn challenge_response_rejects_unaligned_client_key() {
        let err = challenge_response("short", "").unwrap_err();
        assert_eq!(err, CryptoError::NotBlockAligned(5));
    }

    #[test]
    fn session_key_derives_from_the_response() {
        let k = derive_session_key(CAPTURED_H);
        assert_eq!(hex_lower(&k), CAPTURED_SESSION_KEY);
    }

    #[test]
    fn encrypt_pins_the_first_login_frame() {
        let k = derive_session_key(CAPTURED_H);
        let frame = r#"<LSX><Event sender=""><Login IsLoggedIn="true"/></Event></LSX>"#;
        let out = lsx_encrypt(frame, &k);
        assert_eq!(*out.last().unwrap(), 0, "frames are NUL-terminated");
        // Pinned against `openssl enc -aes-128-ecb -nopad` over the PKCS7-padded
        // frame under the captured session key: 62 bytes of XML -> pad 2 -> 4 blocks.
        assert_eq!(
            std::str::from_utf8(&out[..out.len() - 1]).unwrap(),
            concat!(
                "ded1180ab8a2ab85b7408cc009eb0191",
                "00b7b4c827fe0b9b4de2ca7834b3ed51",
                "31a0714a8eb66ba3e38d1855724b0a48",
                "09ad8616e71b3a78880c570c9af7ff80"
            )
        );
    }

    #[test]
    fn round_trip() {
        let k = derive_session_key(CAPTURED_H);
        for xml in [
            r#"<LSX><Response id="1" sender=""><ErrorSuccess Code="0" Description=""/></Response></LSX>"#,
            r#"<LSX><Event sender="LOGIN_EVENT"><Login IsLoggedIn="true"/></Event></LSX>"#,
            "",
        ] {
            assert_eq!(lsx_decrypt(&lsx_encrypt(xml, &k), &k).unwrap(), xml);
        }
    }

    #[test]
    fn round_trip_of_block_aligned_plaintext_pads_a_whole_block() {
        let k = derive_session_key(CAPTURED_H);
        let xml = "0123456789abcdef"; // exactly 16 bytes
        let enc = lsx_encrypt(xml, &k);
        // 2 blocks (16 data + 16 pad) -> 64 hex chars + NUL.
        assert_eq!(enc.len(), 65);
        assert_eq!(lsx_decrypt(&enc, &k).unwrap(), xml);
    }

    #[test]
    fn decrypt_ignores_everything_past_the_first_nul() {
        let k = derive_session_key(CAPTURED_H);
        let mut enc = lsx_encrypt("<LSX/>", &k);
        enc.extend_from_slice(b"deadbeef\0trailing");
        assert_eq!(lsx_decrypt(&enc, &k).unwrap(), "<LSX/>");
    }

    #[test]
    fn decrypt_strips_surrounding_whitespace() {
        let k = derive_session_key(CAPTURED_H);
        let enc = lsx_encrypt("<LSX/>", &k);
        let mut padded = b"  ".to_vec();
        padded.extend_from_slice(&enc[..enc.len() - 1]);
        padded.extend_from_slice(b"\r\n\0");
        assert_eq!(lsx_decrypt(&padded, &k).unwrap(), "<LSX/>");
    }

    #[test]
    fn decrypt_rejects_garbage() {
        let k = derive_session_key(CAPTURED_H);
        assert_eq!(lsx_decrypt(b"zz\0", &k).unwrap_err(), CryptoError::BadHex);
        assert_eq!(lsx_decrypt(b"abc\0", &k).unwrap_err(), CryptoError::BadHex);
        assert_eq!(lsx_decrypt(b"\0", &k).unwrap_err(), CryptoError::Empty);
        assert_eq!(
            lsx_decrypt(b"00112233\0", &k).unwrap_err(),
            CryptoError::NotBlockAligned(4)
        );
    }
}
