happy-cracking 0.3.0

use anyhow::{Context, Result};
use clap::Subcommand;
use num_bigint::BigUint;
use num_traits::{One, Zero};
use std::collections::HashMap;

#[derive(Subcommand)]
pub enum DhAction {
    #[command(about = "Compute Diffie-Hellman public key (g^a mod p)")]
    Pubkey {
        #[arg(long, help = "Generator g")]
        g: String,
        #[arg(long, help = "Prime modulus p")]
        p: String,
        #[arg(long, help = "Private key a")]
        a: String,
    },
    #[command(about = "Compute shared secret (B^a mod p)")]
    SharedSecret {
        #[arg(long, help = "Other party's public key")]
        public_key: String,
        #[arg(long, help = "Prime modulus p")]
        p: String,
        #[arg(long, help = "Your private key a")]
        a: String,
    },
    #[command(about = "Baby-step Giant-step discrete log (find x such that g^x = target mod p)")]
    Dlog {
        #[arg(long, help = "Generator g")]
        g: String,
        #[arg(long, help = "Prime modulus p")]
        p: String,
        #[arg(long, help = "Target value")]
        target: String,
        #[arg(long, help = "Order of generator (for BSGS bound)")]
        order: String,
    },
}

pub fn run(action: DhAction) -> Result<()> {
    match action {
        DhAction::Pubkey { g, p, a } => {
            let g = g.parse::<BigUint>().context("Invalid number for g")?;
            let p = p.parse::<BigUint>().context("Invalid number for p")?;
            let a = a.parse::<BigUint>().context("Invalid number for a")?;
            let pubkey = compute_pubkey(&g, &a, &p)?;
            println!("{}", pubkey);
        }
        DhAction::SharedSecret { public_key, p, a } => {
            let pk = public_key
                .parse::<BigUint>()
                .context("Invalid number for public-key")?;
            let p = p.parse::<BigUint>().context("Invalid number for p")?;
            let a = a.parse::<BigUint>().context("Invalid number for a")?;
            let secret = compute_shared_secret(&pk, &a, &p)?;
            println!("{}", secret);
        }
        DhAction::Dlog {
            g,
            p,
            target,
            order,
        } => {
            let g = g.parse::<BigUint>().context("Invalid number for g")?;
            let p = p.parse::<BigUint>().context("Invalid number for p")?;
            let target = target
                .parse::<BigUint>()
                .context("Invalid number for target")?;
            let order = order
                .parse::<BigUint>()
                .context("Invalid number for order")?;
            let x = baby_step_giant_step(&g, &target, &p, &order)?;
            println!("{}", x);
        }
    }
    Ok(())
}

// Compute Diffie-Hellman public key: g^a mod p.
pub fn compute_pubkey(g: &BigUint, a: &BigUint, p: &BigUint) -> Result<BigUint> {
    if p.is_zero() {
        anyhow::bail!("Modulus p must be non-zero");
    }
    Ok(g.modpow(a, p))
}

// Compute Diffie-Hellman shared secret: public_key^a mod p.
pub fn compute_shared_secret(public_key: &BigUint, a: &BigUint, p: &BigUint) -> Result<BigUint> {
    if p.is_zero() {
        anyhow::bail!("Modulus p must be non-zero");
    }
    Ok(public_key.modpow(a, p))
}

// Safety limit for BSGS algorithm (approx 4 million entries in hash map).
// Prevents OOM when users provide large prime orders.
// 2^22 = 4,194,304.
const MAX_BSGS_ITERATIONS: u64 = 1 << 22;

// Baby-step Giant-step algorithm for computing discrete logarithm.
// Finds x such that g^x ≡ target (mod p), where g has the given order.
// Time and space complexity: O(sqrt(order)).
pub fn baby_step_giant_step(
    g: &BigUint,
    target: &BigUint,
    p: &BigUint,
    order: &BigUint,
) -> Result<BigUint> {
    if p.is_zero() {
        anyhow::bail!("Modulus p must be non-zero");
    }
    if order.is_zero() {
        anyhow::bail!("Order must be non-zero");
    }

    let m = order.sqrt() + BigUint::one();

    if m > BigUint::from(MAX_BSGS_ITERATIONS) {
        anyhow::bail!(
            "Order too large for BSGS algorithm (limit: sqrt(order) <= {})",
            MAX_BSGS_ITERATIONS
        );
    }

    // Baby step: build table of j -> g^j mod p for j in [0, m)
    let mut table: HashMap<BigUint, BigUint> = HashMap::new();
    let mut g_j = BigUint::one(); // g^0 = 1
    let mut j = BigUint::zero();
    while j < m {
        table.insert(g_j.clone(), j.clone());
        g_j = (&g_j * g) % p;
        j += BigUint::one();
    }

    // Giant step factor: g^(-m) mod p
    // g^(-m) = (g^(-1))^m = (g^(order-1))^m mod p (since g^order = 1)
    let g_inv = g.modpow(&(order - BigUint::one()), p);
    let g_neg_m = g_inv.modpow(&m, p);

    // Giant step: compute target * (g^(-m))^i for i in [0, m)
    let mut gamma = target % p;
    let mut i = BigUint::zero();
    while i < m {
        if let Some(j_val) = table.get(&gamma) {
            // x = i*m + j
            let x = (&i * &m + j_val) % order;
            return Ok(x);
        }
        gamma = (&gamma * &g_neg_m) % p;
        i += BigUint::one();
    }

    anyhow::bail!("Discrete log not found (target may not be in the group generated by g)")
}

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

    #[test]
    fn test_pubkey_basic() {
        let g = BigUint::from(2u32);
        let p = BigUint::from(23u32);
        let a = BigUint::from(6u32);
        // 2^6 mod 23 = 64 mod 23 = 18
        assert_eq!(compute_pubkey(&g, &a, &p).unwrap(), BigUint::from(18u32));
    }

    #[test]
    fn test_shared_secret_basic() {
        let g = BigUint::from(5u32);
        let p = BigUint::from(23u32);
        let a = BigUint::from(6u32);
        let b = BigUint::from(15u32);
        let pub_a = compute_pubkey(&g, &a, &p).unwrap();
        let pub_b = compute_pubkey(&g, &b, &p).unwrap();
        let secret_a = compute_shared_secret(&pub_b, &a, &p).unwrap();
        let secret_b = compute_shared_secret(&pub_a, &b, &p).unwrap();
        assert_eq!(secret_a, secret_b);
    }

    #[test]
    fn test_error_on_zero_modulus() {
        let g = BigUint::from(2u32);
        let p = BigUint::from(0u32);
        let a = BigUint::from(10u32);
        assert!(compute_pubkey(&g, &a, &p).is_err());
    }
}