ml-kem cli example (#239)

Cargo.toml

@@ -64,6 +64,7 @@ quickcheck_macros = "1"
 serde_json = { version = "1.0" }
 serde = { version = "1.0", features = ["derive"] }
 hex = { version = "0.4.3", features = ["serde"] }
+clap = { version = "4.5", features = ["derive"] }
 
 [target.'cfg(target_arch = "wasm32")'.dev-dependencies]
 wasm-bindgen-test = "0.3"

examples/mlkem.rs

@@ -0,0 +1,173 @@
+use std::{
+    fs::File,
+    io::{Read, Write},
+};
+
+use libcrux::{drbg::Drbg, kem};
+
+use clap::{Parser, Subcommand};
+
+#[derive(Subcommand)]
+enum GenerateCli {
+    GenerateKey {
+        /// Optionally, a file name to store the keys into.
+        ///
+        /// The keys will be store into `$out.pub` and `$out.priv` when this is used.
+        #[arg(short, long)]
+        out: Option<String>,
+    },
+    Encaps {
+        /// Public key input file to encrypt to.
+        #[arg(short, long)]
+        key: String,
+
+        /// Optionally, a file name to store the encapsulation output.
+        ///
+        /// This defaults to `mlkem.ct`.
+        #[arg(short, long)]
+        ct: Option<String>,
+
+        /// Optionally, a file name to store the shared secret output.
+        ///
+        /// This defaults to `mlkem-encapsulated.ss`.
+        #[arg(short, long)]
+        ss: Option<String>,
+    },
+    Decaps {
+        /// Private key input file to decapsulate with.
+        #[arg(short, long)]
+        key: String,
+
+        /// Encapsulated secret.
+        #[arg(short, long)]
+        ct: String,
+
+        /// Optionally, a file name to store the shared secret.
+        ///
+        /// This defaults to `mlkem-decapsulated.ss``.
+        #[arg(short, long)]
+        ss: Option<String>,
+    },
+}
+
+#[derive(Parser)]
+struct Cli {
+    /// Generate a key pair and write it out.
+    ///
+    /// When no output files are given, `mlkem.pub` and `mlkem` are used.
+    #[command(subcommand)]
+    cmd: GenerateCli,
+
+    /// The key length, 512, 768, or 1024
+    ///
+    /// Defaults to 768
+    #[arg(short, long)]
+    algorithm: Option<u16>,
+}
+
+fn main() {
+    pretty_env_logger::init();
+
+    let cli = Cli::parse();
+
+    let alg = if let Some(l) = cli.algorithm {
+        match l {
+            512 => kem::Algorithm::MlKem512,
+            768 => kem::Algorithm::MlKem768,
+            1024 => kem::Algorithm::MlKem1024,
+            _ => {
+                eprintln!("Invalid algorithm variant {l}");
+                return;
+            }
+        }
+    } else {
+        kem::Algorithm::MlKem768
+    };
+
+    let mut rng = Drbg::new(libcrux::digest::Algorithm::Sha256).unwrap();
+
+    match cli.cmd {
+        GenerateCli::GenerateKey { out: file } => {
+            // Generate a key pair and write it out.
+            let (sk_name, pk_name) = match file {
+                Some(n) => (format!("{n}.priv"), format!("{n}.pub")),
+                None => ("mlkem.priv".to_owned(), "mlkem.pub".to_owned()),
+            };
+
+            let (secret_key, public_key) = kem::key_gen(alg, &mut rng).unwrap();
+
+            println!("Writing private key to {sk_name}");
+            File::create(sk_name.clone())
+                .expect(&format!("Can not create file {sk_name}"))
+                .write_all(&secret_key.encode())
+                .expect("Error writing private key");
+            println!("Writing public key to {pk_name}");
+            File::create(pk_name.clone())
+                .expect(&format!("Can not create file {pk_name}"))
+                .write_all(&public_key.encode())
+                .expect("Error writing public key");
+        }
+
+        GenerateCli::Encaps { key, ct: out, ss } => {
+            let pk = bytes_from_file(key);
+            let pk = kem::PublicKey::decode(alg, &pk).expect("Error decoding public key");
+
+            let (shared_secret, ciphertext) =
+                pk.encapsulate(&mut rng).expect("Error encapsulating");
+
+            let ct_out = match out {
+                Some(n) => n,
+                None => "mlkem.ct".to_owned(),
+            };
+            let ss_out = match ss {
+                Some(n) => n,
+                None => "mlkem-encapsulated.ss".to_owned(),
+            };
+
+            println!("Writing ciphertext to {ct_out}");
+            let mut out_file =
+                File::create(ct_out.clone()).expect(&format!("Can not create file {ct_out}"));
+            out_file
+                .write_all(&ciphertext.encode())
+                .expect("Error writing public key");
+
+            println!("Writing shared secret to {ss_out}");
+            let mut out_file =
+                File::create(ss_out.clone()).expect(&format!("Can not create file {ss_out}"));
+            out_file
+                .write_all(&shared_secret.encode())
+                .expect("Error writing public key");
+        }
+        GenerateCli::Decaps { key, ss: out, ct } => {
+            let sk = bytes_from_file(key);
+            let sk = kem::PrivateKey::decode(alg, &sk).expect("Error decoding private key");
+
+            let ct = bytes_from_file(ct);
+
+            let ct = kem::Ct::decode(alg, &ct).expect("Error decoding ct.");
+            let shared_secret = ct.decapsulate(&sk).expect("Error decapsulating.");
+
+            let out = match out {
+                Some(n) => n,
+                None => "mlkem-decapsulated.ss".to_owned(),
+            };
+
+            println!("Writing shared secret to {out}");
+            let mut out_file =
+                File::create(out.clone()).expect(&format!("Can not create file {out}"));
+            out_file
+                .write_all(&shared_secret.encode())
+                .expect("Error writing public key");
+        }
+    }
+}
+
+fn bytes_from_file(file: String) -> Vec<u8> {
+    println!("Reading file {file}");
+    let mut bytes = Vec::new();
+    File::open(file.clone())
+        .expect(&format!("Error opening file {file}"))
+        .read_to_end(&mut bytes)
+        .expect(&format!("Error reading file {file}."));
+    bytes
+}