From 15296c40de2e0a7b5e51f13af4027558ec6d0fe0 Mon Sep 17 00:00:00 2001
From: Avril <flanchan@cumallover.me>
Date: Sun, 21 Mar 2021 16:28:06 +0000
Subject: [PATCH] implemented binary stream encrypt

---
 src/key.rs  | 65 ++++++++++++++++++++++++++++++++++++++++++++++++-----
 src/main.rs | 54 +++++++++++++++++++++++++++++++-------------
 test        |  1 +
 3 files changed, 98 insertions(+), 22 deletions(-)
 create mode 100644 test

diff --git a/src/key.rs b/src/key.rs
index 0074eca..a9920f7 100644
--- a/src/key.rs
+++ b/src/key.rs
@@ -1,22 +1,25 @@
 use getrandom::getrandom;
-use std::fmt;
+use std::{fmt, str};
 use crate::cha::{
     KEY_SIZE,
     IV_SIZE,
 };
 use crate::ext::*;
 
-#[derive(Debug, Clone, PartialEq, Eq, Hash, Copy)]
+#[derive(Debug, Clone, PartialEq, Eq, Hash, Copy, Default)]
 #[repr(transparent)]
 pub struct Key([u8; KEY_SIZE]);
 
-#[derive(Debug, Clone, PartialEq, Eq, Hash, Copy)]
+#[derive(Debug, Clone, PartialEq, Eq, Hash, Copy, Default)]
 #[repr(transparent)]
 pub struct IV([u8; IV_SIZE]);
 
-
 impl Key
 {
+    #[inline] pub fn from_bytes(k: [u8; KEY_SIZE]) -> Self
+    {
+	Self(k)
+    }
     pub fn new() -> Self
     {
 	let mut output = [0u8; KEY_SIZE];
@@ -27,6 +30,11 @@ impl Key
 
 impl IV
 {
+    
+    #[inline] pub fn from_bytes(k: [u8; IV_SIZE]) -> Self
+    {
+	Self(k)
+    }
     pub fn new() -> Self
     {
 	let mut output = [0u8; IV_SIZE];
@@ -35,6 +43,23 @@ impl IV
     }
 }
 
+impl From<[u8; KEY_SIZE]> for Key
+{
+    #[inline] fn from(from: [u8; KEY_SIZE]) -> Self
+    {
+	Self(from)
+    }
+}
+
+impl From<[u8; IV_SIZE]> for IV
+{
+    fn from(from: [u8; IV_SIZE]) -> Self
+    {
+	Self(from)
+    }
+}
+
+
 impl AsRef<[u8]> for Key
 {
     fn as_ref(&self) -> &[u8]
@@ -85,7 +110,7 @@ impl fmt::Display for Key
 {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result
     {
-	write!(f, "Key({})", self.0.iter().copied().into_hex())
+	write!(f, "{}", self.0.iter().copied().into_hex())
     }
 }
 
@@ -93,6 +118,34 @@ impl fmt::Display for IV
 {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result
     {
-	write!(f, "Key({})", self.0.iter().copied().into_hex())
+	write!(f, "{}", self.0.iter().copied().into_hex())
+    }
+}
+
+impl str::FromStr for Key
+{
+    type Err = base64::DecodeError;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+	let mut buffer = Vec::with_capacity(KEY_SIZE);
+	base64::decode_config_buf(s.as_bytes(), base64::STANDARD, &mut buffer)?;
+
+	let mut this = Self::default();
+	this.0.copy_from_slice(&buffer[..]);
+	Ok(this)
+    }
+}
+
+impl str::FromStr for IV
+{
+    type Err = base64::DecodeError;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+	let mut buffer = Vec::with_capacity(IV_SIZE);
+	base64::decode_config_buf(s.as_bytes(), base64::STANDARD, &mut buffer)?;
+
+	let mut this = Self::default();
+	this.0.copy_from_slice(&buffer[..]);
+	Ok(this)
     }
 }
diff --git a/src/main.rs b/src/main.rs
index 9edcf1b..1075444 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -8,7 +8,7 @@ mod cha;
 mod stream;
 
 use key::{Key, IV};
-
+/*
 fn encrypt((key, iv): &(Key, IV), input: impl AsRef<[u8]>) -> Result<String, openssl::error::ErrorStack>
 {
     let input = input.as_ref();
@@ -42,27 +42,49 @@ fn decrypt((key, iv): &(Key, IV), input: impl AsRef<str>) -> Result<Vec<u8>, ope
 
     println!(">> {}", (&output[..n]).hex());
     assert!(dec.finalize(&mut output[..n])? == 0);
-//    assert!(dec.finalize(&mut output[..n])? == 0);
+    //    assert!(dec.finalize(&mut output[..n])? == 0);
     println!(">> {}", (&output[..n]).hex());
 
     output.truncate(n);
     Ok(output)
+}*/
+
+fn keys() -> Result<(Key, IV), base64::DecodeError>
+{
+    let mut args = std::env::args().skip(1);
+    
+    let key = match args.next() {
+	Some(key) => key.parse()?,
+	None => {
+	    let key = Key::new();
+	    eprintln!("{}", base64::encode(&key));
+	    key
+	},
+    };
+    let iv = match args.next() {
+	Some(iv) => iv.parse()?,
+	None => {
+	    let key = IV::new();
+	    eprintln!("{}", base64::encode(&key));
+	    key
+	},
+    };
+
+    Ok((key, iv))
 }
 
 fn main() {
-    let input = std::env::args().nth(1).unwrap_or({
-	let mut input = [0u8; 16];
-	getrandom::getrandom(&mut input[..]).expect("rng fatal");
-	khash::generate(&Default::default(), input).expect("kana-hash fatal")
-	//input.hex().into()
-    });
+
+    let (key, iv) = keys().expect("Failed to read keys from argv (base64)");
     
-    let key = cha::keygen();
-    let enc = encrypt(&key, &input).expect("encrypt");
-    println!("{}", enc);
-    let dec = decrypt(&key, enc).expect("decrypt");
-
-    let output = std::str::from_utf8(&dec[..]).unwrap();
-    println!("{:?}", output);
-    assert_eq!(output, &input[..]);
+    let stdout = std::io::stdout();
+    let input = std::io::stdin();
+
+    // Encryption
+    {
+	use std::io::Write;
+	let mut output = stream::Sink::encrypt(stdout.lock(), key, iv).expect("Failed to create encrypter");
+	std::io::copy(&mut input.lock(), &mut output).expect("Failed to encrypt");
+	output.flush().expect("Failed to flush stdout");
+    }
 }
diff --git a/test b/test
new file mode 100644
index 0000000..8503b3e
--- /dev/null
+++ b/test
@@ -0,0 +1 @@
+ÝÑðhbÑ Ü~,
\ No newline at end of file