#![cfg(target_os = "windows")]
use clap::{Parser, Subcommand};
use rsa::pkcs8::{EncodePublicKey, LineEnding};
use win_tpm::Tpm;

#[derive(Parser, Debug)]
#[command(author, version)]
struct Args {
    #[command(subcommand)]
    cmd: Cmd,
}

#[derive(Debug, Subcommand)]
enum Cmd {
    /// Delete the key with the given name from the tpm
    Deletekey { name: String },
    /// Create a new rsa key in the tpm
    Createkey { name: String },
    /// Export the public key
    Exportkey { name: String },
    /// Enumerate all keys in the tpm
    Enumkeys,
    /// Decrypts the hex string
    Decrypt { key_name: String, hex_data: String },
    /// Encrypts the hex string
    Encrypt {
        key_name: String,
        input: String,
        /// Treat input string as hex data
        #[clap(long)]
        hex: bool,
    },
}

fn real_main() -> Result<(), Box<dyn std::error::Error>> {
    let args = Args::parse();

    let tpm = Tpm::open()?;

    match args.cmd {
        Cmd::Deletekey { name } => {
            tpm.open_rsa_key(&name)?.delete();
            println!("Successfully deleted key \"{name}\"")
        }
        Cmd::Createkey { name } => {
            tpm.create_rsa_key(&name)?;
            println!("Successfully created key \"{name}\"");
        }
        Cmd::Exportkey { name } => {
            let key = tpm.open_rsa_key(&name)?;
            println!("{}", key.public_key()?.to_public_key_pem(LineEnding::LF)?);
        }
        Cmd::Enumkeys => {
            for key in tpm.enum_keys() {
                println!("Key \"{}\" with algorithm \"{}\"", key.name, key.algorithm);
            }
        }
        Cmd::Decrypt { key_name, hex_data } => {
            let hex_data = hex::decode(hex_data)?;
            let decrypted = tpm.open_rsa_key(&key_name)?.decrypt(&hex_data)?;
            match String::from_utf8(decrypted) {
                Ok(s) => println!("{s}"),
                Err(v) => println!("{}", hex::encode(v.as_bytes())),
            }
        }
        Cmd::Encrypt {
            key_name,
            input,
            hex,
        } => {
            let data = if hex {
                hex::decode(input)?
            } else {
                input.into_bytes()
            };

            let encrypted = tpm.open_rsa_key(&key_name)?.encrypt(&data)?;
            println!("{}", hex::encode(encrypted))
        }
    }

    Ok(())
}

fn main() {
    if let Err(e) = real_main() {
        eprintln!("Error: {e}");
        if let Some(source) = e.source() {
            eprintln!("\t{source}");
        }
    }
}