umskt-rs/src/bink2002.rs

use std::fmt::{Display, Formatter};

use anyhow::{bail, Result};
use bitreader::BitReader;
use openssl::{
    bn::{BigNum, BigNumContext, MsbOption},
    ec::{EcGroup, EcPoint},
    rand::rand_bytes,
    sha::sha1,
};

use crate::{
    crypto::{EllipticCurve, PrivateKey},
    key::{base24_decode, base24_encode, strip_key},
    math::{bitmask, by_dword, next_sn_bits},
};

const FIELD_BITS: i32 = 512;
const FIELD_BYTES: usize = 64;
const SHA_MSG_LENGTH: usize = 3 + 2 * FIELD_BYTES;

const SIGNATURE_LENGTH_BITS: u8 = 62;
const HASH_LENGTH_BITS: u8 = 31;
const CHANNEL_ID_LENGTH_BITS: u8 = 10;
const UPGRADE_LENGTH_BITS: u8 = 1;
const EVERYTHING_ELSE: u8 =
    SIGNATURE_LENGTH_BITS + HASH_LENGTH_BITS + CHANNEL_ID_LENGTH_BITS + UPGRADE_LENGTH_BITS;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct ProductKey {
    upgrade: bool,
    channel_id: u32,
    hash: u32,
    signature: u64,
    auth_info: u32,
}

impl ProductKey {
    pub fn new(
        curve: &EllipticCurve,
        private_key: &PrivateKey,
        channel_id: u32,
        sequence: Option<u32>,
        auth_info: Option<u32>,
        upgrade: Option<bool>,
    ) -> Result<Self> {
        // Generate random sequence if none supplied
        let sequence = match sequence {
            Some(serial) => serial,
            None => {
                let mut bn_rand = BigNum::new()?;
                bn_rand.rand(19, MsbOption::MAYBE_ZERO, false)?;
                let o_raw = u32::from_be_bytes(bn_rand.to_vec_padded(4)?.try_into().unwrap());
                o_raw % 999999
            }
        };

        // Generate random auth info if none supplied
        let auth_info = match auth_info {
            Some(auth_info) => auth_info,
            None => {
                let mut p_auth_info_bytes = [0_u8; 4];
                rand_bytes(&mut p_auth_info_bytes)?;
                u32::from_ne_bytes(p_auth_info_bytes) & ((1 << 10) - 1)
            }
        };

        // Default to upgrade=false
        let upgrade = upgrade.unwrap_or(false);

        // Generate a new random key
        let product_key = Self::generate(
            &curve.curve,
            &curve.gen_point,
            &private_key.gen_order,
            &private_key.private_key,
            channel_id * 1_000_000 + sequence,
            auth_info,
            upgrade,
        )?;

        // Make sure the key is valid
        product_key.verify(&curve.curve, &curve.gen_point, &curve.pub_point)?;

        // Ship it
        Ok(product_key)
    }

    pub fn from_key(curve: &EllipticCurve, key: &str) -> Result<Self> {
        let key = strip_key(key)?;
        let Ok(p_raw) = base24_decode(&key) else {
            bail!("Product key is in an incorrect format!")
        };
        let product_key = Self::from_packed(&p_raw)?;
        let verified = product_key.verify(&curve.curve, &curve.gen_point, &curve.pub_point)?;
        if !verified {
            bail!("Product key is invalid! Wrong BINK ID?");
        }
        Ok(product_key)
    }

    fn generate(
        e_curve: &EcGroup,
        base_point: &EcPoint,
        gen_order: &BigNum,
        private_key: &BigNum,
        p_channel_id: u32,
        p_auth_info: u32,
        p_upgrade: bool,
    ) -> Result<Self> {
        let mut num_context = BigNumContext::new().unwrap();

        let mut c = BigNum::new()?;
        let mut x = BigNum::new()?;
        let mut y = BigNum::new()?;

        let p_data = p_channel_id << 1 | p_upgrade as u32;

        let mut no_square = false;
        let key = loop {
            let mut r = EcPoint::new(e_curve)?;

            c.rand(FIELD_BITS, MsbOption::MAYBE_ZERO, false)?;

            r.mul(e_curve, base_point, &c, &num_context)?;

            r.affine_coordinates(e_curve, &mut x, &mut y, &mut num_context)?;

            let mut msg_buffer: [u8; SHA_MSG_LENGTH] = [0; SHA_MSG_LENGTH];

            let mut x_bin = x.to_vec_padded(FIELD_BYTES as i32)?;
            x_bin.reverse();
            let mut y_bin = y.to_vec_padded(FIELD_BYTES as i32)?;
            y_bin.reverse();

            msg_buffer[0x00] = 0x79;
            msg_buffer[0x01] = (p_data & 0x00FF) as u8;
            msg_buffer[0x02] = ((p_data & 0xFF00) >> 8) as u8;

            msg_buffer[3..3 + FIELD_BYTES].copy_from_slice(&x_bin);
            msg_buffer[3 + FIELD_BYTES..3 + FIELD_BYTES * 2].copy_from_slice(&y_bin);

            let msg_digest = sha1(&msg_buffer);

            let p_hash: u32 = by_dword(&msg_digest[0..4]) & bitmask(31) as u32;

            msg_buffer[0x00] = 0x5D;
            msg_buffer[0x01] = (p_data & 0x00FF) as u8;
            msg_buffer[0x02] = ((p_data & 0xFF00) >> 8) as u8;
            msg_buffer[0x03] = (p_hash & 0x000000FF) as u8;
            msg_buffer[0x04] = ((p_hash & 0x0000FF00) >> 8) as u8;
            msg_buffer[0x05] = ((p_hash & 0x00FF0000) >> 16) as u8;
            msg_buffer[0x06] = ((p_hash & 0xFF000000) >> 24) as u8;
            msg_buffer[0x07] = (p_auth_info & 0x00FF) as u8;
            msg_buffer[0x08] = ((p_auth_info & 0xFF00) >> 8) as u8;
            msg_buffer[0x09] = 0x00;
            msg_buffer[0x0A] = 0x00;

            let msg_digest = sha1(&msg_buffer[..=0x0A]);

            let i_signature = next_sn_bits(by_dword(&msg_digest[4..8]) as u64, 30, 2) << 32
                | by_dword(&msg_digest[0..4]) as u64;

            let mut e = BigNum::from_slice(&i_signature.to_be_bytes())?;

            let e_2 = e.to_owned()?;
            e.mod_mul(&e_2, private_key, gen_order, &mut num_context)?;

            let mut s = e.to_owned()?;

            let s_2 = s.to_owned()?;
            s.mod_sqr(&s_2, gen_order, &mut num_context)?;

            let c_2 = c.to_owned()?;
            c.lshift(&c_2, 2)?;

            s = &s + &c;

            let s_2 = s.to_owned()?;
            if s.mod_sqrt(&s_2, gen_order, &mut num_context).is_err() {
                no_square = true;
            };

            let s_2 = s.to_owned()?;
            s.mod_sub(&s_2, &e, gen_order, &mut num_context)?;

            if s.is_bit_set(0) {
                s = &s + gen_order;
            }

            let s_2 = s.to_owned()?;
            s.rshift1(&s_2)?;

            let p_signature = u64::from_be_bytes(s.to_vec_padded(8)?.try_into().unwrap());

            let product_key = Self {
                upgrade: p_upgrade,
                channel_id: p_channel_id,
                hash: p_hash,
                signature: p_signature,
                auth_info: p_auth_info,
            };

            if p_signature <= bitmask(62) && !no_square {
                break product_key;
            }

            no_square = false;
        };

        Ok(key)
    }

    fn verify(
        &self,
        e_curve: &EcGroup,
        base_point: &EcPoint,
        public_key: &EcPoint,
    ) -> Result<bool> {
        let mut num_context = BigNumContext::new()?;

        let p_data = self.channel_id << 1 | self.upgrade as u32;

        let mut msg_buffer: [u8; SHA_MSG_LENGTH] = [0; SHA_MSG_LENGTH];

        msg_buffer[0x00] = 0x5D;
        msg_buffer[0x01] = (p_data & 0x00FF) as u8;
        msg_buffer[0x02] = ((p_data & 0xFF00) >> 8) as u8;
        msg_buffer[0x03] = (self.hash & 0x000000FF) as u8;
        msg_buffer[0x04] = ((self.hash & 0x0000FF00) >> 8) as u8;
        msg_buffer[0x05] = ((self.hash & 0x00FF0000) >> 16) as u8;
        msg_buffer[0x06] = ((self.hash & 0xFF000000) >> 24) as u8;
        msg_buffer[0x07] = (self.auth_info & 0x00FF) as u8;
        msg_buffer[0x08] = ((self.auth_info & 0xFF00) >> 8) as u8;
        msg_buffer[0x09] = 0x00;
        msg_buffer[0x0A] = 0x00;

        let msg_digest = sha1(&msg_buffer[..=0x0A]);

        let i_signature = next_sn_bits(by_dword(&msg_digest[4..8]) as u64, 30, 2) << 32
            | by_dword(&msg_digest[0..4]) as u64;

        let e = BigNum::from_slice(&i_signature.to_be_bytes())?;
        let s = BigNum::from_slice(&self.signature.to_be_bytes())?;

        let mut x = BigNum::new()?;
        let mut y = BigNum::new()?;

        let mut p = EcPoint::new(e_curve)?;
        let mut t = EcPoint::new(e_curve)?;

        t.mul(e_curve, base_point, &s, &num_context)?;
        p.mul(e_curve, public_key, &e, &num_context)?;
        let p_2 = p.to_owned(e_curve)?;

        p.add(e_curve, &t, &p_2, &mut num_context)?;
        let p_2 = p.to_owned(e_curve)?;

        p.mul(e_curve, &p_2, &s, &num_context)?;

        p.affine_coordinates(e_curve, &mut x, &mut y, &mut num_context)?;

        let mut x_bin = x.to_vec_padded(FIELD_BYTES as i32)?;
        x_bin.reverse();
        let mut y_bin = y.to_vec_padded(FIELD_BYTES as i32)?;
        y_bin.reverse();

        msg_buffer[0x00] = 0x79;
        msg_buffer[0x01] = (p_data & 0x00FF) as u8;
        msg_buffer[0x02] = ((p_data & 0xFF00) >> 8) as u8;

        msg_buffer[3..3 + FIELD_BYTES].copy_from_slice(&x_bin);
        msg_buffer[3 + FIELD_BYTES..3 + FIELD_BYTES * 2].copy_from_slice(&y_bin);

        let msg_digest = sha1(&msg_buffer);

        let hash: u32 = by_dword(&msg_digest[0..4]) & bitmask(31) as u32;

        Ok(hash == self.hash)
    }

    fn from_packed(p_raw: &[u8]) -> Result<Self> {
        let mut reader = BitReader::new(p_raw);
        let auth_info_length_bits = (p_raw.len() * 8) as u8 - EVERYTHING_ELSE;

        let p_auth_info = reader.read_u32(auth_info_length_bits)?;
        let p_signature = reader.read_u64(SIGNATURE_LENGTH_BITS)?;
        let p_hash = reader.read_u32(HASH_LENGTH_BITS)?;
        let p_channel_id = reader.read_u32(CHANNEL_ID_LENGTH_BITS)?;
        let p_upgrade = reader.read_bool()?;

        Ok(Self {
            upgrade: p_upgrade,
            channel_id: p_channel_id,
            hash: p_hash,
            signature: p_signature,
            auth_info: p_auth_info,
        })
    }

    fn pack(&self) -> Vec<u8> {
        let mut p_raw: u128 = 0;

        p_raw |= (self.auth_info as u128)
            << (SIGNATURE_LENGTH_BITS
                + HASH_LENGTH_BITS
                + CHANNEL_ID_LENGTH_BITS
                + UPGRADE_LENGTH_BITS);
        p_raw |= (self.signature as u128)
            << (HASH_LENGTH_BITS + CHANNEL_ID_LENGTH_BITS + UPGRADE_LENGTH_BITS);
        p_raw |= (self.hash as u128) << (CHANNEL_ID_LENGTH_BITS + UPGRADE_LENGTH_BITS);
        p_raw |= (self.channel_id as u128) << UPGRADE_LENGTH_BITS;
        p_raw |= self.upgrade as u128;

        p_raw
            .to_be_bytes()
            .into_iter()
            .skip_while(|&x| x == 0)
            .collect()
    }
}

impl Display for ProductKey {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        let pk = base24_encode(&self.pack()).unwrap();
        let key = pk
            .chars()
            .enumerate()
            .fold(String::new(), |mut acc: String, (i, c)| {
                if i > 0 && i % 5 == 0 {
                    acc.push('-');
                }
                acc.push(c);
                acc
            });
        write!(f, "{}", key)
    }
}

#[cfg(test)]
mod tests {
    use serde_json::from_reader;
    use std::{fs::File, io::BufReader};

    use crate::crypto::EllipticCurve;

    #[test]
    fn verify_test() {
        // Example product key and its BINK ID
        let product_key = "R882X-YRGC8-4KYTG-C3FCC-JCFDY";
        let bink_id = "54";

        // Load keys.json
        let path = "keys.json";
        let file = File::open(path).unwrap();
        let reader = BufReader::new(file);
        let keys: serde_json::Value = from_reader(reader).unwrap();

        let bink = &keys["BINK"][&bink_id];

        let p = bink["p"].as_str().unwrap();
        let a = bink["a"].as_str().unwrap();
        let b = bink["b"].as_str().unwrap();
        let gx = bink["g"]["x"].as_str().unwrap();
        let gy = bink["g"]["y"].as_str().unwrap();
        let kx = bink["pub"]["x"].as_str().unwrap();
        let ky = bink["pub"]["y"].as_str().unwrap();

        let curve = EllipticCurve::new(p, a, b, gx, gy, kx, ky).unwrap();

        assert!(super::ProductKey::from_key(&curve, product_key).is_ok());
        assert!(super::ProductKey::from_key(&curve, "11111-YRGC8-4KYTG-C3FCC-JCFDY").is_err());
    }
}