-
Notifications
You must be signed in to change notification settings - Fork 14
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
1 parent
bb30290
commit 3765371
Showing
2 changed files
with
174 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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 | ||
} |