feat: ml-kem generate

src/mlkemutil.rs

@@ -1,40 +1,151 @@
-use crate::util::base64_encode;
+use crate::rsautil::rsa_public_key_to_jwk;
+use crate::util::{base64_encode, to_pem};
+use ecdsa::elliptic_curve::pkcs8::LineEnding;
 use ml_kem::kem::{Decapsulate, Encapsulate};
-use ml_kem::{EncodedSizeUser, KemCore, MlKem768};
+use ml_kem::{EncodedSizeUser, KemCore, MlKem1024, MlKem512, MlKem768};
 use rust_util::XResult;
 use std::convert::TryInto;
 
-// #[test]
-pub fn generate_ml_kem_768() -> XResult<()> {
-    let mut rng = rand::thread_rng();
-    let (dk, ek) = <MlKem768 as KemCore>::generate(&mut rng);
-    println!("dk: {}", base64_encode(&dk.as_bytes().0.to_vec()));
-    println!("ek: {}", base64_encode(ek.as_bytes().0.to_vec()));
-
-    let ek_bytes = dk.as_bytes().0.to_vec();
-    let dk = <MlKem768 as KemCore>::DecapsulationKey::from_bytes(&opt_result!(
-        ek_bytes.as_slice().try_into(),
-        "Parse decapsulation key failed: {}"
-    ));
-
-    let (encoded_ciphertext, shared_key) = opt_result!(
-        ek.encapsulate(&mut rng),
-        "Encapsulation key encapsulate failed: {:?}"
-    );
-    println!(
-        "encoded_ciphertext: {}",
-        base64_encode(&encoded_ciphertext.0.to_vec())
-    );
-    println!("shared_key: {}", base64_encode(&shared_key.0.to_vec()));
-
-    let k_bytes = encoded_ciphertext.0.to_vec();
-    let shared_key_2 = opt_result!(
-        dk.decapsulate(opt_result!(
-            &k_bytes.as_slice().try_into(),
-            "Parse encoded ciphertext failed: {}"
-        )),
-        "Decapsulation key decapsulate failed: {:?}"
-    );
-    println!("shared_key2: {}", base64_encode(&shared_key_2.0.to_vec()));
-    Ok(())
+pub fn generate_mlkem_keypair(len: usize) -> XResult<(String, String, String, Vec<u8>, String)> {
+    let (dk_private, ek_public) = match len {
+        512 => generate_ml_kem_512(),
+        768 => generate_ml_kem_768(),
+        1024 => generate_ml_kem_1024(),
+        _ => return simple_error!("Invalid ML-KEM={}", len),
+    };
+    let secret_key_der_base64 = base64_encode(&dk_private);
+    let secret_key_pem = to_pem(&format!("ML-KEM-{} PRIVATE KEY", len), &dk_private);
+    let public_key_pem = to_pem(&format!("ML-KEM-{} PUBLIC KEY", len), &ek_public);
+    let public_key_der = ek_public;
+    let jwk_ec_key = "".to_string();
+    Ok((
+        secret_key_der_base64,
+        secret_key_pem,
+        public_key_pem,
+        public_key_der,
+        jwk_ec_key,
+    ))
+}
+
+pub fn generate_ml_kem_512() -> (Vec<u8>, Vec<u8>) {
+    let mut rng = rand::thread_rng();
+    let (dk_private, ek_public) = <MlKem512 as KemCore>::generate(&mut rng);
+    (
+        dk_private.as_bytes().0.to_vec(),
+        ek_public.as_bytes().0.to_vec(),
+    )
+}
+
+pub fn generate_ml_kem_768() -> (Vec<u8>, Vec<u8>) {
+    let mut rng = rand::thread_rng();
+    let (dk_private, ek_public) = <MlKem768 as KemCore>::generate(&mut rng);
+    (
+        dk_private.as_bytes().0.to_vec(),
+        ek_public.as_bytes().0.to_vec(),
+    )
+}
+
+pub fn generate_ml_kem_1024() -> (Vec<u8>, Vec<u8>) {
+    let mut rng = rand::thread_rng();
+    let (dk_private, ek_public) = <MlKem1024 as KemCore>::generate(&mut rng);
+    (
+        dk_private.as_bytes().0.to_vec(),
+        ek_public.as_bytes().0.to_vec(),
+    )
+}
+
+pub fn parse_encapsulation_key_512_public_then_encapsulate(
+    bytes: &[u8],
+) -> XResult<(Vec<u8>, Vec<u8>)> {
+    let ek = <MlKem512 as KemCore>::EncapsulationKey::from_bytes(&opt_result!(
+        bytes.try_into(),
+        "Parse encapsulation key 512 failed: {}"
+    ));
+    let (ciphertext, shared_key) = opt_result!(
+        ek.encapsulate(&mut rand::thread_rng()),
+        "Encapsulation key 512 encapsulate failed: {:?}"
+    );
+    Ok((ciphertext.0.to_vec(), shared_key.0.to_vec()))
+}
+
+pub fn parse_encapsulation_key_768_public_then_encapsulate(
+    bytes: &[u8],
+) -> XResult<(Vec<u8>, Vec<u8>)> {
+    let ek = <MlKem768 as KemCore>::EncapsulationKey::from_bytes(&opt_result!(
+        bytes.try_into(),
+        "Parse encapsulation key 768 failed: {}"
+    ));
+    let (ciphertext, shared_key) = opt_result!(
+        ek.encapsulate(&mut rand::thread_rng()),
+        "Encapsulation key 768 encapsulate failed: {:?}"
+    );
+    Ok((ciphertext.0.to_vec(), shared_key.0.to_vec()))
+}
+
+pub fn parse_encapsulation_key_1024_public_then_encapsulate(
+    bytes: &[u8],
+) -> XResult<(Vec<u8>, Vec<u8>)> {
+    let ek = <MlKem1024 as KemCore>::EncapsulationKey::from_bytes(&opt_result!(
+        bytes.try_into(),
+        "Parse encapsulation key 1024 failed: {}"
+    ));
+    let (ciphertext, shared_key) = opt_result!(
+        ek.encapsulate(&mut rand::thread_rng()),
+        "Encapsulation key 1024 encapsulate failed: {:?}"
+    );
+    Ok((ciphertext.0.to_vec(), shared_key.0.to_vec()))
+}
+
+pub fn parse_decapsulate_key_512_private_then_decapsulate(
+    key_bytes: &[u8],
+    ciphertext_bytes: &[u8],
+) -> XResult<Vec<u8>> {
+    let dk = <MlKem512 as KemCore>::DecapsulationKey::from_bytes(&opt_result!(
+        key_bytes.try_into(),
+        "Parse decapsulation key 512 failed: {}"
+    ));
+    let shared_key = opt_result!(
+        dk.decapsulate(opt_result!(
+            ciphertext_bytes.try_into(),
+            "Parse encoded ciphertext 512 failed: {}"
+        )),
+        "Decapsulation key 512 decapsulate failed: {:?}"
+    );
+    Ok(shared_key.0.to_vec())
+}
+
+pub fn parse_decapsulate_key_768_private_then_decapsulate(
+    key_bytes: &[u8],
+    ciphertext_bytes: &[u8],
+) -> XResult<Vec<u8>> {
+    let dk = <MlKem768 as KemCore>::DecapsulationKey::from_bytes(&opt_result!(
+        key_bytes.try_into(),
+        "Parse decapsulation key 768 failed: {}"
+    ));
+    let shared_key = opt_result!(
+        dk.decapsulate(opt_result!(
+            ciphertext_bytes.try_into(),
+            "Parse encoded ciphertext 768 failed: {}"
+        )),
+        "Decapsulation key 768 decapsulate failed: {:?}"
+    );
+    Ok(shared_key.0.to_vec())
+}
+
+pub fn parse_decapsulate_key_1024_private_then_decapsulate(
+    key_bytes: &[u8],
+    ciphertext_bytes: &[u8],
+) -> XResult<Vec<u8>> {
+    let dk = <MlKem768 as KemCore>::DecapsulationKey::from_bytes(&opt_result!(
+        key_bytes.try_into(),
+        "Parse decapsulation key 1024 failed: {}"
+    ));
+    let shared_key = opt_result!(
+        dk.decapsulate(opt_result!(
+            ciphertext_bytes.try_into(),
+            "Parse encoded ciphertext 1024 failed: {}"
+        )),
+        "Decapsulation key 1024 decapsulate failed: {:?}"
+    );
+    Ok(shared_key.0.to_vec())
 }