libfreemkv 0.3.1

//! AACS bus authentication handshake — ECDH key agreement + bus key derivation.
//!
//! Implements the AACS SCSI authentication protocol to obtain:
//!   - Volume ID (VID) — needed for VUK derivation
//!   - Read Data Key — needed for AACS 2.0 (UHD) bus decryption
//!
//! Flow:
//!   1. Invalidate AGIDs → allocate fresh AGID
//!   2. Send host certificate + nonce
//!   3. Receive drive certificate + nonce
//!   4. Receive drive key point + signature, verify
//!   5. Sign host key point, send
//!   6. ECDH: host_priv × drive_key_point → bus key (low 128 bits of x)
//!   7. Read VID or Read Data Keys (encrypted with bus key)
//!
//! Supports:
//!   - AACS 1.0: custom 160-bit curve, SHA-1, 20-byte keys
//!   - AACS 2.0: drives accept AACS 1.0 host certs for backward compatibility
//!     (full P-256/SHA-256 AACS 2.0 handshake prepared but rarely needed)

use crate::error::{Error, Result};
use crate::drive::DriveSession;
use crate::scsi::DataDirection;
use num_bigint::BigUint;
use num_traits::{One, Zero};
use sha1::{Sha1, Digest};

/// Execute a SCSI command that reads data from the device.
fn scsi_read(session: &mut DriveSession, cdb: &[u8], len: usize) -> Result<Vec<u8>> {
    let mut buf = vec![0u8; len];
    session.scsi_execute(cdb, DataDirection::FromDevice, &mut buf, 5_000)?;
    Ok(buf)
}

/// Execute a SCSI command that writes data to the device.
fn scsi_write(session: &mut DriveSession, cdb: &[u8], data: &[u8]) -> Result<()> {
    let mut buf = data.to_vec();
    session.scsi_execute(cdb, DataDirection::ToDevice, &mut buf, 5_000)?;
    Ok(())
}

// ── AACS 1.0 elliptic curve parameters (160-bit) ───────────────────────────

const EC_P: [u8; 20] = [
    0x9D, 0xC9, 0xD8, 0x13, 0x55, 0xEC, 0xCE, 0xB5, 0x60, 0xBD,
    0xB0, 0x9E, 0xF9, 0xEA, 0xE7, 0xC4, 0x79, 0xA7, 0xD7, 0xDF,
];
const EC_A: [u8; 20] = [
    0x9D, 0xC9, 0xD8, 0x13, 0x55, 0xEC, 0xCE, 0xB5, 0x60, 0xBD,
    0xB0, 0x9E, 0xF9, 0xEA, 0xE7, 0xC4, 0x79, 0xA7, 0xD7, 0xDC,
];
#[cfg(test)]
const EC_B: [u8; 20] = [
    0x40, 0x2D, 0xAD, 0x3E, 0xC1, 0xCB, 0xCD, 0x16, 0x52, 0x48,
    0xD6, 0x8E, 0x12, 0x45, 0xE0, 0xC4, 0xDA, 0xAC, 0xB1, 0xD8,
];
const EC_N: [u8; 20] = [
    0x9D, 0xC9, 0xD8, 0x13, 0x55, 0xEC, 0xCE, 0xB5, 0x60, 0xBD,
    0xC4, 0x4F, 0x54, 0x81, 0x7B, 0x2C, 0x7F, 0x5A, 0xB0, 0x17,
];
const EC_GX: [u8; 20] = [
    0x2E, 0x64, 0xFC, 0x22, 0x57, 0x83, 0x51, 0xE6, 0xF4, 0xCC,
    0xA7, 0xEB, 0x81, 0xD0, 0xA4, 0xBD, 0xC5, 0x4C, 0xCE, 0xC6,
];
const EC_GY: [u8; 20] = [
    0x09, 0x14, 0xA2, 0x5D, 0xD0, 0x54, 0x42, 0x88, 0x9D, 0xB4,
    0x55, 0xC7, 0xF2, 0x3C, 0x9A, 0x07, 0x07, 0xF5, 0xCB, 0xB9,
];

// ── AACS LA (Licensing Administrator) public key for cert verification ──────

const AACS_LA_PUB_X: [u8; 20] = [
    0x01, 0xF3, 0x5D, 0xAB, 0xD8, 0xAE, 0x5F, 0x40, 0x56, 0x5E,
    0x30, 0xC8, 0x8A, 0x60, 0x42, 0x82, 0x07, 0x61, 0xDF, 0x93,
];
const AACS_LA_PUB_Y: [u8; 20] = [
    0x44, 0x87, 0xB5, 0xAC, 0x07, 0x10, 0x8D, 0x10, 0x5B, 0xA5,
    0xB9, 0xE3, 0x2F, 0x3B, 0xBB, 0xFC, 0x0C, 0x2C, 0xBC, 0xD1,
];

// ── Elliptic curve arithmetic over GF(p) ───────────────────────────────────

#[derive(Clone, Debug)]
struct EcPoint {
    x: BigUint,
    y: BigUint,
    infinity: bool,
}

impl EcPoint {
    fn infinity() -> Self {
        EcPoint { x: BigUint::zero(), y: BigUint::zero(), infinity: true }
    }

    fn new(x: BigUint, y: BigUint) -> Self {
        EcPoint { x, y, infinity: false }
    }

    fn from_bytes(x_bytes: &[u8], y_bytes: &[u8]) -> Self {
        EcPoint::new(BigUint::from_bytes_be(x_bytes), BigUint::from_bytes_be(y_bytes))
    }
}

/// Modular inverse using extended Euclidean algorithm.
fn mod_inv(a: &BigUint, m: &BigUint) -> Option<BigUint> {
    use num_bigint::BigInt;
    use num_traits::Signed;

    let a = BigInt::from(a.clone());
    let m = BigInt::from(m.clone());

    let (mut old_r, mut r) = (a, m.clone());
    let (mut old_s, mut s) = (BigInt::one(), BigInt::zero());

    while !r.is_zero() {
        let q = &old_r / &r;
        let temp_r = r.clone();
        r = old_r - &q * &r;
        old_r = temp_r;
        let temp_s = s.clone();
        s = old_s - &q * &s;
        old_s = temp_s;
    }

    if old_r != BigInt::one() {
        return None;
    }

    if old_s.is_negative() {
        old_s += &m;
    }
    Some(old_s.to_biguint().unwrap())
}

/// EC point addition on curve y² = x³ + ax + b (mod p).
fn ec_add(p1: &EcPoint, p2: &EcPoint, a: &BigUint, p: &BigUint) -> EcPoint {
    if p1.infinity { return p2.clone(); }
    if p2.infinity { return p1.clone(); }

    if p1.x == p2.x {
        if p1.y == p2.y && !p1.y.is_zero() {
            return ec_double(p1, a, p);
        }
        return EcPoint::infinity();
    }

    // λ = (y2 - y1) / (x2 - x1) mod p
    let dy = if p2.y >= p1.y {
        (&p2.y - &p1.y) % p
    } else {
        (p - (&p1.y - &p2.y) % p) % p
    };
    let dx = if p2.x >= p1.x {
        (&p2.x - &p1.x) % p
    } else {
        (p - (&p1.x - &p2.x) % p) % p
    };

    let dx_inv = mod_inv(&dx, p).unwrap();
    let lam = (&dy * &dx_inv) % p;

    // x3 = λ² - x1 - x2 mod p
    let x3 = {
        let lam2 = (&lam * &lam) % p;
        let sum = (&p1.x + &p2.x) % p;
        if lam2 >= sum {
            (lam2 - sum) % p
        } else {
            (p - (sum - lam2) % p) % p
        }
    };

    // y3 = λ(x1 - x3) - y1 mod p
    let y3 = {
        let diff = if p1.x >= x3 {
            (&p1.x - &x3) % p
        } else {
            (p - (&x3 - &p1.x) % p) % p
        };
        let prod = (&lam * &diff) % p;
        if prod >= p1.y {
            (prod - &p1.y) % p
        } else {
            (p - (&p1.y - prod) % p) % p
        }
    };

    EcPoint::new(x3, y3)
}

/// EC point doubling.
fn ec_double(pt: &EcPoint, a: &BigUint, p: &BigUint) -> EcPoint {
    if pt.infinity || pt.y.is_zero() {
        return EcPoint::infinity();
    }

    // λ = (3x² + a) / (2y) mod p
    let three = BigUint::from(3u32);
    let two = BigUint::from(2u32);

    let numerator = (&three * &pt.x * &pt.x + a) % p;
    let denominator = (&two * &pt.y) % p;
    let denom_inv = mod_inv(&denominator, p).unwrap();
    let lam = (&numerator * &denom_inv) % p;

    // x3 = λ² - 2x mod p
    let x3 = {
        let lam2 = (&lam * &lam) % p;
        let two_x = (&two * &pt.x) % p;
        if lam2 >= two_x {
            (lam2 - two_x) % p
        } else {
            (p - (two_x - lam2) % p) % p
        }
    };

    // y3 = λ(x - x3) - y mod p
    let y3 = {
        let diff = if pt.x >= x3 {
            (&pt.x - &x3) % p
        } else {
            (p - (&x3 - &pt.x) % p) % p
        };
        let prod = (&lam * &diff) % p;
        if prod >= pt.y {
            (prod - &pt.y) % p
        } else {
            (p - (&pt.y - prod) % p) % p
        }
    };

    EcPoint::new(x3, y3)
}

/// Scalar multiplication using double-and-add.
fn ec_mul(k: &BigUint, pt: &EcPoint, a: &BigUint, p: &BigUint) -> EcPoint {
    if k.is_zero() {
        return EcPoint::infinity();
    }

    let mut result = EcPoint::infinity();
    let mut base = pt.clone();
    let mut scalar = k.clone();

    while !scalar.is_zero() {
        if scalar.bit(0) {
            result = ec_add(&result, &base, a, p);
        }
        base = ec_double(&base, a, p);
        scalar >>= 1;
    }

    result
}

/// Convert BigUint to fixed-size big-endian bytes, zero-padded.
fn to_bytes_be_padded(n: &BigUint, len: usize) -> Vec<u8> {
    let bytes = n.to_bytes_be();
    if bytes.len() >= len {
        bytes[bytes.len() - len..].to_vec()
    } else {
        let mut padded = vec![0u8; len - bytes.len()];
        padded.extend_from_slice(&bytes);
        padded
    }
}

// ── ECDSA ───────────────────────────────────────────────────────────────────

/// ECDSA sign: sign SHA-1(data) with private key on AACS curve.
/// Returns (r, s) each 20 bytes.
fn ecdsa_sign(priv_key: &[u8; 20], data: &[u8]) -> ([u8; 20], [u8; 20]) {
    let p = BigUint::from_bytes_be(&EC_P);
    let a = BigUint::from_bytes_be(&EC_A);
    let n = BigUint::from_bytes_be(&EC_N);
    let g = EcPoint::from_bytes(&EC_GX, &EC_GY);
    let d = BigUint::from_bytes_be(priv_key);

    // Hash the data
    let hash = Sha1::digest(data);
    let z = BigUint::from_bytes_be(&hash);

    loop {
        // Generate random k
        let mut k_bytes = [0u8; 20];
        use rand::RngCore;
        rand::thread_rng().fill_bytes(&mut k_bytes);
        let k = BigUint::from_bytes_be(&k_bytes) % &n;
        if k.is_zero() { continue; }

        // R = k × G
        let r_point = ec_mul(&k, &g, &a, &p);
        let r = &r_point.x % &n;
        if r.is_zero() { continue; }

        // s = k⁻¹(z + r·d) mod n
        let k_inv = match mod_inv(&k, &n) {
            Some(v) => v,
            None => continue,
        };
        let s = (&k_inv * ((&z + &r * &d) % &n)) % &n;
        if s.is_zero() { continue; }

        let r_bytes = to_bytes_be_padded(&r, 20);
        let s_bytes = to_bytes_be_padded(&s, 20);

        let mut r_out = [0u8; 20];
        let mut s_out = [0u8; 20];
        r_out.copy_from_slice(&r_bytes);
        s_out.copy_from_slice(&s_bytes);

        return (r_out, s_out);
    }
}

/// ECDSA verify: verify signature (r, s) against SHA-1(data) using public key.
fn ecdsa_verify(pub_x: &[u8; 20], pub_y: &[u8; 20], sig_r: &[u8; 20], sig_s: &[u8; 20], data: &[u8]) -> bool {
    let p = BigUint::from_bytes_be(&EC_P);
    let a = BigUint::from_bytes_be(&EC_A);
    let n = BigUint::from_bytes_be(&EC_N);
    let g = EcPoint::from_bytes(&EC_GX, &EC_GY);
    let q = EcPoint::from_bytes(pub_x, pub_y);

    let r = BigUint::from_bytes_be(sig_r);
    let s = BigUint::from_bytes_be(sig_s);

    if r.is_zero() || r >= n || s.is_zero() || s >= n {
        return false;
    }

    let hash = Sha1::digest(data);
    let z = BigUint::from_bytes_be(&hash);

    let s_inv = match mod_inv(&s, &n) {
        Some(v) => v,
        None => return false,
    };

    let u1 = (&z * &s_inv) % &n;
    let u2 = (&r * &s_inv) % &n;

    let p1 = ec_mul(&u1, &g, &a, &p);
    let p2 = ec_mul(&u2, &q, &a, &p);
    let r_point = ec_add(&p1, &p2, &a, &p);

    if r_point.infinity {
        return false;
    }

    &r_point.x % &n == r
}

// ── AACS certificate handling ───────────────────────────────────────────────

/// Verify an AACS certificate (92 bytes) against the AACS LA public key.
fn verify_cert(cert: &[u8]) -> bool {
    if cert.len() < 92 { return false; }
    // Certificate format: type(1) + flags(1) + padding(2) + serial(6) + pub_x(20) + pub_y(20) + sig_r(20) + sig_s(20)
    // Signature is over the first 52 bytes
    let mut sig_r = [0u8; 20];
    let mut sig_s = [0u8; 20];
    sig_r.copy_from_slice(&cert[52..72]);
    sig_s.copy_from_slice(&cert[72..92]);

    ecdsa_verify(&AACS_LA_PUB_X, &AACS_LA_PUB_Y, &sig_r, &sig_s, &cert[..52])
}

/// Extract public key from certificate.
fn cert_pub_key(cert: &[u8]) -> ([u8; 20], [u8; 20]) {
    let mut x = [0u8; 20];
    let mut y = [0u8; 20];
    x.copy_from_slice(&cert[12..32]);
    y.copy_from_slice(&cert[32..52]);
    (x, y)
}

// ── Bus key derivation (ECDH) ───────────────────────────────────────────────

/// Compute bus key via ECDH: bus_key = low 128 bits of (host_priv × drive_key_point).x
fn compute_bus_key(host_priv: &[u8; 20], drive_key_point_x: &[u8; 20], drive_key_point_y: &[u8; 20]) -> [u8; 16] {
    let p = BigUint::from_bytes_be(&EC_P);
    let a = BigUint::from_bytes_be(&EC_A);

    let d = BigUint::from_bytes_be(host_priv);
    let dkp = EcPoint::from_bytes(drive_key_point_x, drive_key_point_y);

    let shared = ec_mul(&d, &dkp, &a, &p);

    // Bus key = lowest 128 bits (last 16 bytes) of x-coordinate
    let x_bytes = to_bytes_be_padded(&shared.x, 20);
    let mut bus_key = [0u8; 16];
    bus_key.copy_from_slice(&x_bytes[4..20]); // last 16 of 20
    bus_key
}

/// Generate ephemeral host key pair: (private_key, public_point_x, public_point_y).
fn generate_host_key_pair() -> ([u8; 20], [u8; 20], [u8; 20]) {
    let p_mod = BigUint::from_bytes_be(&EC_P);
    let a = BigUint::from_bytes_be(&EC_A);
    let g = EcPoint::from_bytes(&EC_GX, &EC_GY);

    let mut priv_bytes = [0u8; 20];
    use rand::RngCore;
    rand::thread_rng().fill_bytes(&mut priv_bytes);
    let d = BigUint::from_bytes_be(&priv_bytes);

    let q = ec_mul(&d, &g, &a, &p_mod);

    let qx = to_bytes_be_padded(&q.x, 20);
    let qy = to_bytes_be_padded(&q.y, 20);

    let mut pub_x = [0u8; 20];
    let mut pub_y = [0u8; 20];
    pub_x.copy_from_slice(&qx);
    pub_y.copy_from_slice(&qy);

    (priv_bytes, pub_x, pub_y)
}

// ── AES-CMAC (for MAC verification) ────────────────────────────────────────

/// AES-128-CMAC over 16 bytes of data.
fn aes_cmac_16(data: &[u8; 16], key: &[u8; 16]) -> [u8; 16] {
    use aes::Aes128;
    use aes::cipher::{BlockEncrypt, KeyInit, generic_array::GenericArray};

    let cipher = Aes128::new(GenericArray::from_slice(key));

    // For single-block CMAC:
    // 1. Generate subkey K1
    let mut l = GenericArray::clone_from_slice(&[0u8; 16]);
    cipher.encrypt_block(&mut l);

    let mut k1 = [0u8; 16];
    let carry = (l[0] >> 7) & 1;
    for i in 0..15 {
        k1[i] = (l[i] << 1) | (l[i + 1] >> 7);
    }
    k1[15] = l[15] << 1;
    if carry == 1 {
        k1[15] ^= 0x87; // Rb for AES-128
    }

    // 2. XOR data with K1, encrypt
    let mut block = [0u8; 16];
    for i in 0..16 {
        block[i] = data[i] ^ k1[i];
    }
    let mut ga = GenericArray::clone_from_slice(&block);
    cipher.encrypt_block(&mut ga);

    let mut mac = [0u8; 16];
    mac.copy_from_slice(&ga);
    mac
}

// ── SCSI command builders ───────────────────────────────────────────────────

/// Build REPORT KEY CDB (0xA4).
fn cdb_report_key(agid: u8, format: u8, len: u16) -> [u8; 12] {
    let mut cdb = [0u8; 12];
    cdb[0] = crate::scsi::SCSI_REPORT_KEY;
    cdb[7] = crate::scsi::AACS_KEY_CLASS;
    cdb[8] = (len >> 8) as u8;
    cdb[9] = (len & 0xFF) as u8;
    cdb[10] = (agid << 6) | (format & 0x3F);
    cdb
}

/// Build SEND KEY CDB (0xA3).
fn cdb_send_key(agid: u8, format: u8, len: u16) -> [u8; 12] {
    let mut cdb = [0u8; 12];
    cdb[0] = crate::scsi::SCSI_SEND_KEY;
    cdb[7] = crate::scsi::AACS_KEY_CLASS;
    cdb[8] = (len >> 8) as u8;
    cdb[9] = (len & 0xFF) as u8;
    cdb[10] = (agid << 6) | (format & 0x3F);
    cdb
}

/// Build REPORT DISC STRUCTURE CDB (0xAD).
fn cdb_report_disc_structure(agid: u8, format: u8, len: u16) -> [u8; 12] {
    let mut cdb = [0u8; 12];
    cdb[0] = crate::scsi::SCSI_READ_DISC_STRUCTURE;
    cdb[1] = 0x01; // Blu-ray
    cdb[7] = format;
    cdb[8] = (len >> 8) as u8;
    cdb[9] = (len & 0xFF) as u8;
    cdb[10] = agid << 6;
    cdb
}

// ── High-level handshake ────────────────────────────────────────────────────

/// Result of a successful AACS authentication handshake.
#[derive(Debug)]
pub struct AacsAuth {
    /// Bus key (16 bytes) — derived from ECDH
    pub bus_key: [u8; 16],
    /// AGID used for this session
    pub agid: u8,
    /// Volume ID (16 bytes) — read after auth
    pub volume_id: Option<[u8; 16]>,
    /// Read data key (16 bytes) — for AACS 2.0 bus decryption
    pub read_data_key: Option<[u8; 16]>,
    /// Drive certificate (92 bytes)
    pub drive_cert: [u8; 92],
}

/// Perform the full AACS authentication handshake.
///
/// Requires a host private key (20 bytes) and host certificate (92 bytes)
/// from the KEYDB.cfg HC entry.
pub fn aacs_authenticate(
    session: &mut DriveSession,
    host_priv_key: &[u8; 20],
    host_cert: &[u8],
) -> Result<AacsAuth> {
    if host_cert.len() < 92 {
        return Err(Error::AacsError { detail: "host certificate too short".into() });
    }

    // Step 1: Invalidate all AGIDs
    for agid in 0..4u8 {
        let cdb = cdb_report_key(agid, 0x3F, 2);
        let _ = scsi_read(session, &cdb, 2);
    }

    // Step 2: Allocate AGID
    let cdb = cdb_report_key(0, 0x00, 8);
    let response = scsi_read(session, &cdb, 8)
        .map_err(|e| Error::AacsError { detail: format!("failed to allocate AGID: {}", e) })?;
    let agid = (response[7] >> 6) & 0x03;

    // Step 3: Generate host nonce and ephemeral key pair
    let mut host_nonce = [0u8; 20];
    use rand::RngCore;
    rand::thread_rng().fill_bytes(&mut host_nonce);
    let (host_key, host_key_point_x, host_key_point_y) = generate_host_key_pair();

    // Step 4: Send host certificate + nonce (SEND KEY format 0x01)
    let mut send_buf = [0u8; 116];
    send_buf[1] = 0x72; // data length
    send_buf[4..24].copy_from_slice(&host_nonce);
    send_buf[24..116].copy_from_slice(&host_cert[..92]);

    let cdb = cdb_send_key(agid, 0x01, 116);
    scsi_write(session, &cdb, &send_buf)
        .map_err(|_| Error::AacsError { detail: "drive rejected host certificate".into() })?;

    // Step 5: Read drive certificate + nonce (REPORT KEY format 0x01)
    let cdb = cdb_report_key(agid, 0x01, 116);
    let response = scsi_read(session, &cdb, 116)
        .map_err(|_| Error::AacsError { detail: "failed to read drive certificate".into() })?;

    let mut drive_nonce = [0u8; 20];
    let mut drive_cert = [0u8; 92];
    drive_nonce.copy_from_slice(&response[4..24]);
    drive_cert.copy_from_slice(&response[24..116]);

    // Detect AACS 2.0 drive certificate (type 0x11)
    // AACS 2.0 drives use P-256/SHA-256 natively but accept AACS 1.0 host certs
    // for backward compatibility. We proceed with AACS 1.0 handshake.
    if drive_cert[0] == 0x11 {
        // AACS 2.0 drive detected — falling back to AACS 1.0 handshake
        // (full P-256 AACS 2.0 handshake not yet implemented)
        // The drive should still accept our AACS 1.0 host certificate.
    }

    // Verify drive certificate (AACS 1.0 LA signature)
    if drive_cert[0] == 0x01 && !verify_cert(&drive_cert) {
        return Err(Error::AacsError { detail: "drive certificate verification failed".into() });
    }
    // Skip verification for AACS 2.0 certs (different LA key, P-256 curve)

    // Step 6: Read drive key point + signature (REPORT KEY format 0x02)
    let cdb = cdb_report_key(agid, 0x02, 84);
    let response = scsi_read(session, &cdb, 84)
        .map_err(|_| Error::AacsError { detail: "failed to read drive key".into() })?;

    let mut drive_key_point = [0u8; 40];   // x(20) + y(20)
    let mut drive_key_sig = [0u8; 40];     // r(20) + s(20)
    drive_key_point.copy_from_slice(&response[4..44]);
    drive_key_sig.copy_from_slice(&response[44..84]);

    // Verify drive key signature: sign(drive_nonce=host_nonce || drive_key_point)
    let (drive_pub_x, drive_pub_y) = cert_pub_key(&drive_cert);
    let mut verify_data = [0u8; 60];
    verify_data[..20].copy_from_slice(&host_nonce);
    verify_data[20..60].copy_from_slice(&drive_key_point);

    let mut sig_r = [0u8; 20];
    let mut sig_s = [0u8; 20];
    sig_r.copy_from_slice(&drive_key_sig[..20]);
    sig_s.copy_from_slice(&drive_key_sig[20..40]);

    if !ecdsa_verify(&drive_pub_x, &drive_pub_y, &sig_r, &sig_s, &verify_data) {
        return Err(Error::AacsError { detail: "drive key signature verification failed".into() });
    }

    // Step 7: Sign host key point (ECDSA over drive_nonce || host_key_point)
    let mut sign_data = [0u8; 60];
    sign_data[..20].copy_from_slice(&drive_nonce);
    sign_data[20..40].copy_from_slice(&host_key_point_x);
    sign_data[40..60].copy_from_slice(&host_key_point_y);

    let (host_sig_r, host_sig_s) = ecdsa_sign(host_priv_key, &sign_data);

    // Step 8: Send host key point + signature (SEND KEY format 0x02)
    let mut send_buf = [0u8; 84];
    send_buf[1] = 0x52;
    send_buf[4..24].copy_from_slice(&host_key_point_x);
    send_buf[24..44].copy_from_slice(&host_key_point_y);
    send_buf[44..64].copy_from_slice(&host_sig_r);
    send_buf[64..84].copy_from_slice(&host_sig_s);

    let cdb = cdb_send_key(agid, 0x02, 84);
    scsi_write(session, &cdb, &send_buf)
        .map_err(|_| Error::AacsError { detail: "drive rejected host key".into() })?;

    // Step 9: Compute bus key via ECDH
    let mut dkp_x = [0u8; 20];
    let mut dkp_y = [0u8; 20];
    dkp_x.copy_from_slice(&drive_key_point[..20]);
    dkp_y.copy_from_slice(&drive_key_point[20..40]);

    let bus_key = compute_bus_key(&host_key, &dkp_x, &dkp_y);

    Ok(AacsAuth {
        bus_key,
        agid,
        volume_id: None,
        read_data_key: None,
        drive_cert,
    })
}

/// Read Volume ID after successful authentication.
pub fn read_volume_id(session: &mut DriveSession, auth: &mut AacsAuth) -> Result<[u8; 16]> {
    // REPORT DISC STRUCTURE format 0x80
    let cdb = cdb_report_disc_structure(auth.agid, 0x80, 36);
    let response = scsi_read(session, &cdb, 36)
        .map_err(|_| Error::AacsError { detail: "failed to read Volume ID".into() })?;

    let mut vid = [0u8; 16];
    let mut mac = [0u8; 16];
    vid.copy_from_slice(&response[4..20]);
    mac.copy_from_slice(&response[20..36]);

    // Verify MAC: AES-CMAC(VID, bus_key) should equal mac
    let calc_mac = aes_cmac_16(&vid, &auth.bus_key);
    if calc_mac != mac {
        return Err(Error::AacsError { detail: "VID MAC verification failed".into() });
    }

    auth.volume_id = Some(vid);
    Ok(vid)
}

/// Read data keys after successful authentication (for AACS 2.0 bus encryption).
pub fn read_data_keys(session: &mut DriveSession, auth: &mut AacsAuth) -> Result<([u8; 16], [u8; 16])> {
    // REPORT DISC STRUCTURE format 0x84
    let cdb = cdb_report_disc_structure(auth.agid, 0x84, 36);
    let response = scsi_read(session, &cdb, 36)
        .map_err(|_| Error::AacsError { detail: "failed to read data keys".into() })?;

    let mut enc_rdk = [0u8; 16];
    let mut enc_wdk = [0u8; 16];
    enc_rdk.copy_from_slice(&response[4..20]);
    enc_wdk.copy_from_slice(&response[20..36]);

    // Decrypt with bus key (AES-ECB)
    let read_data_key = super::aes_ecb_decrypt(&auth.bus_key, &enc_rdk);
    let write_data_key = super::aes_ecb_decrypt(&auth.bus_key, &enc_wdk);

    auth.read_data_key = Some(read_data_key);
    Ok((read_data_key, write_data_key))
}

// ── Tests ───────────────────────────────────────────────────────────────────

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

    #[test]
    fn test_ec_curve_generator_on_curve() {
        // Verify G is on the curve: y² = x³ + ax + b (mod p)
        let p = BigUint::from_bytes_be(&EC_P);
        let a = BigUint::from_bytes_be(&EC_A);
        let b = BigUint::from_bytes_be(&EC_B);
        let gx = BigUint::from_bytes_be(&EC_GX);
        let gy = BigUint::from_bytes_be(&EC_GY);

        let lhs = (&gy * &gy) % &p;
        let rhs = (&gx * &gx * &gx + &a * &gx + &b) % &p;
        assert_eq!(lhs, rhs, "Generator point is not on the curve");
    }

    #[test]
    fn test_ec_mul_identity() {
        let p = BigUint::from_bytes_be(&EC_P);
        let a = BigUint::from_bytes_be(&EC_A);
        let g = EcPoint::from_bytes(&EC_GX, &EC_GY);

        // 1 × G = G
        let result = ec_mul(&BigUint::one(), &g, &a, &p);
        assert_eq!(result.x, g.x);
        assert_eq!(result.y, g.y);
    }

    #[test]
    fn test_ec_mul_order() {
        // n × G = O (point at infinity)
        let p = BigUint::from_bytes_be(&EC_P);
        let a = BigUint::from_bytes_be(&EC_A);
        let n = BigUint::from_bytes_be(&EC_N);
        let g = EcPoint::from_bytes(&EC_GX, &EC_GY);

        let result = ec_mul(&n, &g, &a, &p);
        assert!(result.infinity, "n × G should be point at infinity");
    }

    #[test]
    fn test_ecdsa_sign_verify() {
        // Generate a key pair and test sign/verify
        let (priv_key, pub_x, pub_y) = generate_host_key_pair();
        let data = b"test data for AACS ECDSA";

        let (sig_r, sig_s) = ecdsa_sign(&priv_key, data);
        assert!(ecdsa_verify(&pub_x, &pub_y, &sig_r, &sig_s, data),
            "ECDSA signature should verify");

        // Verify with wrong data fails
        assert!(!ecdsa_verify(&pub_x, &pub_y, &sig_r, &sig_s, b"wrong data"),
            "ECDSA should fail with wrong data");
    }

    #[test]
    fn test_ecdh_shared_secret() {
        // Two parties should derive the same shared point
        let p = BigUint::from_bytes_be(&EC_P);
        let a = BigUint::from_bytes_be(&EC_A);
        let g = EcPoint::from_bytes(&EC_GX, &EC_GY);

        let (priv_a, pub_ax, pub_ay) = generate_host_key_pair();
        let (priv_b, pub_bx, pub_by) = generate_host_key_pair();

        // A computes: priv_a × pub_B
        let shared_a = compute_bus_key(&priv_a, &pub_bx, &pub_by);
        // B computes: priv_b × pub_A
        let shared_b = compute_bus_key(&priv_b, &pub_ax, &pub_ay);

        assert_eq!(shared_a, shared_b, "ECDH shared secrets should match");
    }

    #[test]
    fn test_aes_cmac() {
        // Basic CMAC test — at minimum verify it produces consistent output
        let key = [0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6,
                   0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf, 0x4f, 0x3c];
        let data = [0u8; 16];
        let mac1 = aes_cmac_16(&data, &key);
        let mac2 = aes_cmac_16(&data, &key);
        assert_eq!(mac1, mac2);
        assert_ne!(mac1, [0u8; 16]); // shouldn't be all zeros
    }

    #[test]
    fn test_verify_host_cert_from_keydb() {
        // Verify the host cert from our KEYDB
        let keydb_path = match std::env::var("KEYDB_PATH").ok() {
            Some(p) => std::path::PathBuf::from(p),
            None => return, // skip if KEYDB_PATH not set
        };
        if !keydb_path.exists() { return; }

        let db = crate::aacs::KeyDb::load(&keydb_path).unwrap();
        if let Some(hc) = &db.host_cert {
            let valid = verify_cert(&hc.certificate);
            eprintln!("Host cert verification: {}", if valid { "PASS" } else { "FAIL" });
            // Note: our cert is revoked but should still have valid LA signature
            // If it doesn't verify, the LA public key might be wrong
            if !valid {
                eprintln!("  (cert may use different LA key or format)");
            }
        }
    }
}