umskt-rs/src/cli.rs

use std::{fs::File, io::BufReader, path::Path};

use anyhow::{anyhow, Result};
use clap::Parser;
use openssl::{
    bn::{BigNum, MsbOption},
    ec::{EcGroup, EcPoint},
};
use serde_json::from_reader;

use crate::{bink1998, crypto::initialize_elliptic_curve};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Mode {
    Bink1998,
    Bink2002,
    ConfirmationId,
}

impl Default for Mode {
    fn default() -> Self {
        Self::Bink1998
    }
}

#[derive(Parser, Debug)]
#[command(author, about, long_about = None)]
pub struct Options {
    /// Enable verbose output
    #[arg(short, long)]
    verbose: bool,

    /// Number of keys to generate
    #[arg(short = 'n', long = "number", default_value = "1")]
    num_keys: i32,

    /// Specify which keys file to load
    #[arg(short = 'f', long = "file", default_value = "keys.json")]
    keys_filename: String,

    /// Installation ID used to generate confirmation ID
    #[arg(short, long)]
    instid: Option<String>,

    /// Specify which BINK identifier to load
    #[arg(short, long, default_value = "2E")]
    binkid: String,

    /// Show which products/binks can be loaded
    #[arg(short, long)]
    list: bool,

    /// Specify which Channel Identifier to use
    #[arg(short = 'c', long = "channel", default_value = "640")]
    channel_id: u32,

    #[clap(skip)]
    application_mode: Mode,
}

pub struct Cli {
    options: Options,
    keys: serde_json::Value,
    private_key: BigNum,
    gen_order: BigNum,
    gen_point: EcPoint,
    pub_point: EcPoint,
    e_curve: EcGroup,
    product_key: Option<String>,
    count: u32,
}

impl Cli {
    pub fn new() -> Result<Self> {
        let mut options = Self::parse_command_line();
        let keys = Self::validate_command_line(&mut options)?;

        let bink = &keys["BINK"][&options.binkid];
        // We cannot produce a valid key without knowing the private key k. The reason for this is that
        // we need the result of the function K(x; y) = kG(x; y).
        let private_key = BigNum::from_dec_str(bink["priv"].as_str().unwrap()).unwrap();

        // We can, however, validate any given key using the available public key: {p, a, b, G, K}.
        // genOrder the order of the generator G, a value we have to reverse -> Schoof's Algorithm.
        let gen_order = BigNum::from_dec_str(bink["n"].as_str().unwrap()).unwrap();

        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 n = bink["n"].as_str().unwrap();
        let k = bink["priv"].as_str().unwrap();

        if options.verbose {
            println!("-----------------------------------------------------------");
            println!(
                "Loaded the following elliptic curve parameters: BINK[{}]",
                options.binkid
            );
            println!("-----------------------------------------------------------");
            println!(" P: {p}");
            println!(" a: {a}");
            println!(" b: {b}");
            println!("Gx: {gx}");
            println!("Gy: {gy}");
            println!("Kx: {kx}");
            println!("Ky: {ky}");
            println!(" n: {n}");
            println!(" k: {k}");
            println!();
        }

        let (e_curve, gen_point, pub_point) = initialize_elliptic_curve(p, a, b, gx, gy, kx, ky);

        Ok(Self {
            options,
            keys,
            private_key,
            gen_order,
            gen_point,
            pub_point,
            e_curve,
            product_key: None,
            count: 0,
        })
    }

    fn parse_command_line() -> Options {
        let mut args = Options::parse();
        if args.instid.is_some() {
            args.application_mode = Mode::ConfirmationId;
        }
        args
    }

    fn validate_command_line(options: &mut Options) -> Result<serde_json::Value> {
        if options.verbose {
            println!("Loading keys file {}", options.keys_filename);
        }

        let keys = Self::load_json(&options.keys_filename)?;

        if options.verbose {
            println!("Loaded keys from {} successfully", options.keys_filename);
        }

        if options.list {
            let products = keys["Products"].as_object().ok_or(anyhow!(
                "`Products` object not found in {}",
                options.keys_filename
            ))?;
            for (key, value) in products.iter() {
                println!("{}: {}", key, value["BINK"]);
            }

            println!("\n\n** Please note: any BINK ID other than 2E is considered experimental at this time **\n");
        }

        let bink_id = u32::from_str_radix(&options.binkid, 16)?;

        if bink_id >= 0x40 {
            options.application_mode = Mode::Bink2002;
        }

        if options.channel_id > 999 {
            return Err(anyhow!(
                "Refusing to create a key with a Channel ID greater than 999"
            ));
        }

        Ok(keys)
    }

    fn load_json<P: AsRef<Path>>(path: P) -> Result<serde_json::Value> {
        let file = File::open(path)?;
        let reader = BufReader::new(file);
        let json = from_reader(reader)?;
        Ok(json)
    }

    pub fn run(&mut self) -> Result<()> {
        match self.options.application_mode {
            Mode::Bink1998 => self.bink1998(),
            Mode::Bink2002 => self.bink2002(),
            Mode::ConfirmationId => self.confirmation_id(),
        }
    }

    fn bink1998(&mut self) -> Result<()> {
        let mut n_raw = self.options.channel_id * 1_000_000; // <- change

        let mut bn_rand = BigNum::new()?;
        bn_rand.rand(19, MsbOption::MAYBE_ZERO, false)?;

        let o_raw: u32 = u32::from_be_bytes(bn_rand.to_vec_padded(4)?.try_into().unwrap());
        n_raw += o_raw % 999999;

        if self.options.verbose {
            println!("> PID: {n_raw:09}");
        }

        let private_key = &self.gen_order - &self.private_key;

        let upgrade = false;

        for _ in 0..self.options.num_keys {
            let p_key = bink1998::generate(
                &self.e_curve,
                &self.gen_point,
                &self.gen_order,
                &private_key,
                n_raw,
                upgrade,
            )?;
            Cli::print_key(&p_key);

            if bink1998::verify(&self.e_curve, &self.gen_point, &self.pub_point, &p_key)? {
                self.count += 1;
            }
        }

        println!("Success count: {}/{}", self.count, self.options.num_keys);
        Ok(())
    }

    fn bink2002(&mut self) -> Result<()> {
        todo!()
    }

    fn confirmation_id(&mut self) -> Result<()> {
        todo!()
    }

    fn print_key(pk: &str) {
        assert_eq!(pk.len(), 25);
        println!(
            "{}",
            pk.chars()
                .enumerate()
                .fold(String::new(), |mut acc: String, (i, c)| {
                    if i > 0 && i % 5 == 0 {
                        acc.push('-');
                    }
                    acc.push(c);
                    acc
                })
        );
    }
}