From 612c13dc9076df25704bbcae23448c7b4c5a611b Mon Sep 17 00:00:00 2001
From: lionarius <main@lionarius.ru>
Date: Tue, 22 Oct 2024 14:05:01 +0300
Subject: [PATCH] add calculate key

---
 src/main.rs | 61 +++++++++++++++++++++++++++--------------------------
 src/rsa.rs  | 12 +++++++++++
 2 files changed, 43 insertions(+), 30 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 2c1ca16..da943fc 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -9,6 +9,7 @@ mod utils;
 
 enum Mode {
     GenerateKey,
+    CalculateKey,
     Encrypt,
     Decrypt,
 }
@@ -17,6 +18,7 @@ impl Display for Mode {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
             Mode::GenerateKey => write!(f, "Generate key"),
+            Mode::CalculateKey => write!(f, "Calculate key"),
             Mode::Encrypt => write!(f, "Encrypt"),
             Mode::Decrypt => write!(f, "Decrypt"),
         }
@@ -26,12 +28,18 @@ impl Display for Mode {
 fn main() -> anyhow::Result<()> {
     let mode = Select::new(
         "Select a mode",
-        vec![Mode::GenerateKey, Mode::Encrypt, Mode::Decrypt],
+        vec![
+            Mode::GenerateKey,
+            Mode::CalculateKey,
+            Mode::Encrypt,
+            Mode::Decrypt,
+        ],
     )
     .prompt()?;
 
     match mode {
         Mode::GenerateKey => generate_key()?,
+        Mode::CalculateKey => calculate_key()?,
         Mode::Encrypt => encrypt()?,
         Mode::Decrypt => decrypt()?,
     }
@@ -55,10 +63,29 @@ fn generate_key() -> anyhow::Result<()> {
 
     let (public_key, private_key) = rsa::keygen(bits);
 
+    println!();
     println!("Generated key with {} bits", private_key.bits());
     println!();
     println!("Public key:");
     println!("{}", public_key);
+    println!();
+    println!("Private key:");
+    println!("{}", private_key);
+
+    Ok(())
+}
+
+fn calculate_key() -> anyhow::Result<()> {
+    let p = CustomType::new("Enter prime p").prompt()?;
+    let q = CustomType::new("Enter prime q").prompt()?;
+    let e = CustomType::new("Enter public exponent e").prompt()?;
+
+    let (public_key, private_key) = rsa::calculate_key(p, q, e);
+
+    println!();
+    println!("Public key:");
+    println!("{}", public_key);
+    println!();
     println!("Private key:");
     println!("{}", private_key);
 
@@ -66,27 +93,14 @@ fn generate_key() -> anyhow::Result<()> {
 }
 
 fn encrypt() -> anyhow::Result<()> {
-    let key_size = CustomType::new("Enter the key size in bits")
-        .with_default(2048u64)
-        .with_validator(|input: &u64| {
-            if input < &16 {
-                Ok(Validation::Invalid(
-                    "Key size must be at least 16 bits".into(),
-                ))
-            } else {
-                Ok(Validation::Valid)
-            }
-        })
-        .prompt()?;
-
     let n = CustomType::new("Enter public key's N").prompt()?;
     let e = CustomType::new("Enter public key's E").prompt()?;
 
     let data = Text::new("Enter data").prompt()?;
 
-    let chunk_size = (key_size / 8) as usize;
-    let safe_chunk_size = chunk_size - 1;
     let public_key = rsa::PublicKey::new(n, e);
+    let chunk_size = (public_key.bits() / 8) as usize;
+    let safe_chunk_size = chunk_size - 1;
     let data = data.as_bytes();
 
     println!();
@@ -111,26 +125,13 @@ fn encrypt() -> anyhow::Result<()> {
 }
 
 fn decrypt() -> anyhow::Result<()> {
-    let key_size = CustomType::new("Enter the key size in bits")
-        .with_default(2048u64)
-        .with_validator(|input: &u64| {
-            if input < &8 {
-                Ok(Validation::Invalid(
-                    "Key size must be at least 8 bits".into(),
-                ))
-            } else {
-                Ok(Validation::Valid)
-            }
-        })
-        .prompt()?;
-
     let n = CustomType::new("Enter private key's N").prompt()?;
     let d = CustomType::new("Enter private key's D").prompt()?;
 
     let encrypted = Text::new("Enter encrypted data").prompt()?;
 
-    let chunk_size = (key_size / 8) as usize;
     let private_key = rsa::PrivateKey::new(n, d);
+    let chunk_size = (private_key.bits() / 8) as usize;
     let encrypted = Vec::<u8>::from_hex(encrypted)?;
 
     println!();
diff --git a/src/rsa.rs b/src/rsa.rs
index f034250..2dc7d11 100644
--- a/src/rsa.rs
+++ b/src/rsa.rs
@@ -90,3 +90,15 @@ pub fn keygen(bits: u64) -> (PublicKey, PrivateKey) {
 
     (public_key, private_key)
 }
+
+pub fn calculate_key(p: BigUint, q: BigUint, e: BigUint) -> (PublicKey, PrivateKey) {
+    let n = &p * &q;
+    let phi = (p - 1u8) * (q - 1u8);
+
+    let d = e.modinv(&phi).unwrap();
+
+    let public_key = PublicKey::new(n.clone(), e);
+    let private_key = PrivateKey::new(n, d);
+
+    (public_key, private_key)
+}