cryptopals/src/rsa.rs

use num_bigint::BigUint;
use num_bigint::RandBigInt;
use openssl::bn::BigNum;
use openssl::bn::BigNumContext;
use openssl::error::ErrorStack;

#[derive(Clone)]
pub struct PublicKey(pub BigUint);

#[derive(Clone)]
pub struct PrivateKey(pub BigUint);

#[derive(Clone)]
pub struct Keypair {
    pub private: PrivateKey,
    pub public: PublicKey,
}

impl Keypair {
    pub fn make(p: &BigUint, g: &BigUint) -> Self {
        let mut rng = rand::thread_rng();
        let private = rng.gen_biguint_below(p);
        let public = g.modpow(&private, p);
        Self {
            private: PrivateKey(private),
            public: PublicKey(public),
        }
    }
}

pub struct RsaPublicKey {
    pub e: BigNum,
    pub n: BigNum,
}

pub struct RsaPrivateKey {
    d: BigNum,
    n: BigNum,
}

pub fn rsa_gen_keys() -> Result<(RsaPublicKey, RsaPrivateKey), ErrorStack> {
    let mut ctx = BigNumContext::new()?;

    // Generate 2 random primes. We'll use small numbers to start, so you can just pick them out of a prime table. Call them "p" and "q".
    let mut p = BigNum::new()?;
    let mut q = BigNum::new()?;
    p.generate_prime(256, true, None, None)?;
    q.generate_prime(256, true, None, None)?;

    // Let n be p * q. Your RSA math is modulo n.
    let mut n = BigNum::new()?;
    n.checked_mul(&p, &q, &mut ctx)?;
    // This is stupid but I couldn't figure out how to clone a bignum so we do this.
    let mut n2 = BigNum::new()?;
    n2.checked_mul(&p, &q, &mut ctx)?;

    // Let et be (p-1)*(q-1) (the "totient"). You need this value only for keygen.
    let mut et = BigNum::new()?;
    q.sub_word(1)?;
    p.sub_word(1)?;
    et.checked_mul(&p, &q, &mut ctx)?;

    // Let e be 3.
    // Compute d = invmod(e, et). invmod(17, 3120) is 2753.
    let e = BigNum::from_u32(3)?;
    let d = invmod(&e, &et)?;

    // Your public key is [e, n]. Your private key is [d, n].
    Ok((RsaPublicKey { e, n }, RsaPrivateKey { d, n: n2 }))
}

pub fn invmod(a: &BigNum, n: &BigNum) -> Result<BigNum, ErrorStack> {
    fn extended_gcd(a: BigNum, b: BigNum) -> Result<(BigNum, BigNum, BigNum), ErrorStack> {
        // credit: https://www.dcode.fr/extended-gcd
        //  r1 = b, r2 = a, u1 = 0, v1 = 1, u2 = 1, v2 = 0
        let mut r1 = b;
        let mut r2 = a;
        let mut u1 = BigNum::from_u32(0)?;
        let mut v1 = BigNum::from_u32(1)?;
        let mut u2 = BigNum::from_u32(1)?;
        let mut v2 = BigNum::from_u32(0)?;

        // while (r2! = 0) do
        while r2.num_bits() != 0 {
            // q = r1 ÷ r2 (integer division)
            let q = &r1 / &r2;

            // r3 = r1, u3 = u1, v3 = v1,
            let r3 = r1;
            let u3 = u1;
            let v3 = v1;

            // r1 = r2, u1 = u2, v1 = v2,
            r1 = r2;
            u1 = u2;
            v1 = v2;

            //   r2 = r3 - q * r2, u2 = u3 - q * u2, v2 = v3 - q * v2
            r2 = &r3 - &(&q * &r1);
            u2 = &u3 - &(&q * &u1);
            v2 = &v3 - &(&q * &v1);
        }

        //  return (r1, u1, v1) (r1 natural integer and u1, v1 rational integers)
        Ok((r1, u1, v1))
    }

    // No, couldn't think of a worse way to do that.
    let a_cloned = BigNum::from_hex_str(&a.to_hex_str()?)?;
    let n_cloned = BigNum::from_hex_str(&n.to_hex_str()?)?;

    let (_, u1, _) = extended_gcd(a_cloned, n_cloned)?;
    let r_manual = &(&(&u1 % n) + n) % n;

    let mut ctx = BigNumContext::new()?;
    let mut r = BigNum::new()?;
    r.mod_inverse(&a, &n, &mut ctx)?;
    assert_eq!(r_manual, r, "invmod implementation incorrect");
    Ok(r_manual)
}

pub fn rsa_encrypt(m: &BigNum, p: &RsaPublicKey) -> Result<BigNum, ErrorStack> {
    let mut ctx = BigNumContext::new()?;
    let mut c = BigNum::new()?;
    // To encrypt: c = m**e%n.
    c.mod_exp(&m, &p.e, &p.n, &mut ctx)?;
    Ok(c)
}

pub fn rsa_decrypt(c: &BigNum, p: &RsaPrivateKey) -> Result<BigNum, ErrorStack> {
    let mut ctx = BigNumContext::new()?;
    let mut m = BigNum::new()?;
    // To decrypt: m = c**d%n.
    m.mod_exp(&c, &p.d, &p.n, &mut ctx)?;
    Ok(m)
}

pub fn rsa_encrypt_str(m: &str, p: &RsaPublicKey) -> Result<BigNum, ErrorStack> {
    //  Finally, to encrypt a string, do something cheesy.
    let m = BigNum::from_slice(&m.as_bytes())?;
    assert!(m < p.n);
    rsa_encrypt(&m, p)
}

pub fn rsa_decrypt_str(c: &BigNum, p: &RsaPrivateKey) -> Result<String, ErrorStack> {
    let m = rsa_decrypt(c, p)?;
    Ok(String::from(std::str::from_utf8(&m.to_vec()).unwrap()))
}