From 2bd3314a30700eee8b977cb4914f435189f1da27 Mon Sep 17 00:00:00 2001
From: lionarius <main@lionarius.ru>
Date: Thu, 7 Nov 2024 15:27:01 +0300
Subject: [PATCH] .

---
 src/main.rs | 63 ++++++++++++++++++++++++++++++++++++++++++-----------
 1 file changed, 50 insertions(+), 13 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 98c777c..81ab198 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -2,7 +2,7 @@ use core::str;
 use std::fmt::Display;
 
 use hex::{FromHex, ToHex};
-use inquire::{validator::Validation, CustomType, Select, Text};
+use inquire::{validator::Validation, Confirm, CustomType, Select, Text};
 
 mod pha;
 mod rsa;
@@ -26,6 +26,20 @@ impl Display for Mode {
     }
 }
 
+enum InputMode {
+    Utf8,
+    Hex,
+}
+
+impl Display for InputMode {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            InputMode::Utf8 => write!(f, "UTF-8"),
+            InputMode::Hex => write!(f, "Hex"),
+        }
+    }
+}
+
 fn main() -> anyhow::Result<()> {
     let mode = Select::new(
         "Select a mode",
@@ -48,6 +62,21 @@ fn main() -> anyhow::Result<()> {
     Ok(())
 }
 
+fn input_data(prompt: &str) -> anyhow::Result<Vec<u8>> {
+    let input_mode =
+        Select::new("Select input mode", vec![InputMode::Utf8, InputMode::Hex]).prompt()?;
+
+    let data = match input_mode {
+        InputMode::Utf8 => Text::new(prompt).prompt()?.into(),
+        InputMode::Hex => {
+            let data = Text::new("Enter encrypted data").prompt()?;
+            Vec::<u8>::from_hex(data)?
+        }
+    };
+
+    Ok(data)
+}
+
 fn generate_key() -> anyhow::Result<()> {
     let bits = CustomType::new("Enter the key size in bits")
         .with_default(2048)
@@ -96,20 +125,28 @@ fn encrypt() -> anyhow::Result<()> {
     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 data = input_data("Enter data")?;
+
+    let chunkify = Confirm::new("Do you want to chunk data?")
+        .with_default(true)
+        .prompt()?;
 
     let public_key = pha::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();
+    let chunk_size = 1.max((public_key.bits() / 8) as usize);
 
-    if chunk_size < 2 {
-        println!(
-            "This key cannot encrypt byte data! Key size {} bits is less than 16 bits.",
-            public_key.bits()
-        );
-        return Ok(());
-    }
+    let safe_chunk_size = if chunkify {
+        if chunk_size < 2 {
+            println!(
+                "This key cannot encrypt byte data! Key size {} bits is less than 16 bits.",
+                public_key.bits()
+            );
+            return Ok(());
+        }
+
+        chunk_size - 1
+    } else {
+        1
+    };
 
     println!();
     println!("data (raw) = {:?}", data);
@@ -139,7 +176,7 @@ fn decrypt() -> anyhow::Result<()> {
     let encrypted = Text::new("Enter encrypted data").prompt()?;
 
     let private_key = pha::PrivateKey::new(n, d);
-    let chunk_size = (private_key.bits() / 8) as usize;
+    let chunk_size = 1.max((private_key.bits() / 8) as usize);
     let encrypted = Vec::<u8>::from_hex(encrypted)?;
 
     println!();