From 37e3bbb70d6ede76d647d0c71796729e799aacb8 Mon Sep 17 00:00:00 2001
From: jacekpoz <jacekpoz@proton.me>
Date: Tue, 14 May 2024 14:37:27 +0200
Subject: [PATCH] add key length as cli param

---
 l4/z1/rsa-rs/src/main.rs | 17 +++++++++--------
 1 file changed, 9 insertions(+), 8 deletions(-)

diff --git a/l4/z1/rsa-rs/src/main.rs b/l4/z1/rsa-rs/src/main.rs
index 9f09f56..86aa91b 100644
--- a/l4/z1/rsa-rs/src/main.rs
+++ b/l4/z1/rsa-rs/src/main.rs
@@ -1,9 +1,9 @@
-use std::str::FromStr;
+use std::env;
 
 use rand::thread_rng;
 use num_bigint::{BigInt, BigUint, RandBigInt, ToBigInt};
 use num_traits::{One, Zero};
-use num_prime::nt_funcs::is_prime;
+use num_prime::{nt_funcs::is_prime, PrimalityTestConfig, RandPrime};
 
 #[derive(Clone, Debug, PartialEq)]
 struct Key {
@@ -140,12 +140,13 @@ fn find_private_key(n: &BigUint, e: &BigUint, d: &BigUint) -> (BigUint, BigUint)
 }
 
 fn main() {
-    let p1 = BigUint::from_str(
-        "974727060224362646764494950343283895192247884131401225431125883626677704823599485452659329462717264789225578738426263942189471610243051112987455612230669642273763764066874828122466915752449939769915907320666994623051644250931060706544475012753421053710000840777305651367720812671715570852037695820669"
-    ).unwrap();
-    let p2 = BigUint::from_str(
-        "405429572483370650711751211583337124241480095100234252252445505276023252806716412598284451359933836905532376631046948097710959823641510851399588137748599199145772088806906054590783739726661129142422308347175545861420111080331585498348237997146513527685781878012545164661089657770453030501423904836319"
-    ).unwrap();
+    let n = env::args().nth(1)
+        .expect(format!("usage: {} <key_length>", env::args().nth(0).unwrap()).as_str())
+        .parse::<usize>()
+        .expect("key length must be usize");
+
+    let p1: BigUint = thread_rng().gen_prime_exact(n, Some(PrimalityTestConfig::strict()));
+    let p2: BigUint = thread_rng().gen_prime_exact(n, Some(PrimalityTestConfig::strict()));
 
     println!("p1 ({b} bits): {p1}", b = p1.bits());
     println!("p2 ({b} bits): {p2}", b = p1.bits());