umskt-rs/src/bin/xpkey.rs

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

use anyhow::{anyhow, Result};
use clap::Parser;
use serde_json::{from_reader, from_str};

use umskt::{
    bink1998, bink2002, confid,
    crypto::{EllipticCurve, PrivateKey},
};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Mode {
    Bink1998Generate,
    Bink2002Generate,
    Bink1998Validate,
    Bink2002Validate,
    ConfirmationId,
}

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

#[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: u64,

    /// Specify which keys file to load
    #[arg(short = 'f', long = "file")]
    keys_filename: Option<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,

    /// Product key to validate signature
    #[arg(short = 'V', long = "validate")]
    key_to_check: Option<String>,

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

fn main() -> Result<()> {
    let mut options = parse_command_line();
    let keys = 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 = bink["priv"].as_str().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 = bink["n"].as_str().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();

    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: {gen_order}");
        println!(" k: {private_key}");
        println!();
    }

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

    match options.application_mode {
        Mode::Bink1998Generate => {
            let private_key = PrivateKey::new(gen_order, private_key)?;
            bink1998_generate(&options, &curve, &private_key)?;
        }
        Mode::Bink2002Generate => {
            let private_key = PrivateKey::new(gen_order, private_key)?;
            bink2002_generate(&options, &curve, &private_key)?;
        }
        Mode::Bink1998Validate => bink1998_validate(&options, &curve)?,
        Mode::Bink2002Validate => bink2002_validate(&options, &curve)?,
        Mode::ConfirmationId => confirmation_id(&options)?,
    }

    Ok(())
}

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> {
    let keys = {
        if let Some(filename) = &options.keys_filename {
            if options.verbose {
                println!("Loading keys file {}", filename);
            }

            let keys = load_json(filename)?;

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

            keys
        } else {
            from_str(std::include_str!("../../keys.json"))?
        }
    };

    if options.list {
        let products = keys["Products"]
            .as_object()
            .ok_or(anyhow!("`Products` object not found in keys",))?;
        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 options.key_to_check.is_some() {
        options.application_mode = Mode::Bink1998Validate;
    }

    if bink_id >= 0x40 {
        if options.key_to_check.is_some() {
            options.application_mode = Mode::Bink2002Validate;
        } else {
            options.application_mode = Mode::Bink2002Generate;
        }
    }

    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)
}

fn bink1998_generate(
    options: &Options,
    curve: &EllipticCurve,
    private_key: &PrivateKey,
) -> Result<()> {
    for _ in 0..options.num_keys {
        let product_key =
            bink1998::ProductKey::new(curve, private_key, options.channel_id, None, None)?;
        if options.verbose {
            println!("{:?}", product_key);
        }
        println!("{product_key}");
    }
    Ok(())
}

fn bink2002_generate(
    options: &Options,
    curve: &EllipticCurve,
    private_key: &PrivateKey,
) -> Result<()> {
    for _ in 0..options.num_keys {
        let product_key =
            bink2002::ProductKey::new(curve, private_key, options.channel_id, None, None)?;
        if options.verbose {
            println!("{:?}", product_key);
        }
        println!("{product_key}");
    }
    Ok(())
}

fn bink1998_validate(options: &Options, curve: &EllipticCurve) -> Result<()> {
    let product_key =
        bink1998::ProductKey::from_key(curve, options.key_to_check.as_ref().unwrap())?;
    if options.verbose {
        println!("{:?}", product_key);
    }
    println!("{product_key}");
    println!("Key validated successfully!");
    Ok(())
}

fn bink2002_validate(options: &Options, curve: &EllipticCurve) -> Result<()> {
    let product_key =
        bink2002::ProductKey::from_key(curve, options.key_to_check.as_ref().unwrap())?;
    if options.verbose {
        println!("{:?}", product_key);
    }
    println!("{product_key}");
    println!("Key validated successfully!");
    Ok(())
}

fn confirmation_id(options: &Options) -> Result<()> {
    if let Some(instid) = &options.instid {
        let confirmation_id = confid::generate(instid)?;
        println!("Confirmation ID: {confirmation_id}");
    };
    Ok(())
}