From 4f89807a0479f8f402acb46f6051e66885da8548 Mon Sep 17 00:00:00 2001
From: Avril <flanchan@cumallover.me>
Date: Sun, 21 Mar 2021 18:57:11 +0000
Subject: [PATCH] encrypt, decrypt, keygen all work

---
 Cargo.toml  |   2 +-
 src/main.rs | 107 ++++++++++++++++++++++++++++------------------------
 2 files changed, 59 insertions(+), 50 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 760a9c4..15491a8 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,7 +1,7 @@
 [package]
 name = "chacha20"
 description = "chacha20_poly1305 encryption tool"
-version = "0.1.0"
+version = "1.0.0"
 authors = ["Avril <flanchan@cumallover.me>"]
 edition = "2018"
 
diff --git a/src/main.rs b/src/main.rs
index 5614484..4638241 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,4 +1,6 @@
 
+#![allow(dead_code)]
+
 mod ext; #[macro_use] use ext::*;
 
 mod key;
@@ -6,50 +8,45 @@ 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();
-    let mut output = vec![0u8; input.len()];
-
-    eprintln!("(enc) Key: {}, IV: {}, Input: ({}, {})", key, iv, input.len(), input.hex());
-    
-    let mut enc = cha::encrypter(key, iv)?;
-
-    let n = enc.update(&input[..], &mut output[..])?;
-    eprintln!("(enc) Written {} bytes", n);
-    
-    println!(">> {}", (&output[..n]).hex());
-    assert!(enc.finalize(&mut output[..n])? == 0);
-    println!(">> {}", (&output[..n]).hex());
 
-    Ok(base64::encode(&output[..n]))
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub enum Mode
+{
+    Encrypt, Decrypt, Keygen
 }
 
-fn decrypt((key, iv): &(Key, IV), input: impl AsRef<str>) -> Result<Vec<u8>, openssl::error::ErrorStack>
+fn keys() -> Result<(Mode, Key, IV), base64::DecodeError>
 {
-    let input = base64::decode(input.as_ref()).expect("invalid base64");
-    let mut output = vec![0u8; input.len()];
-
-    eprintln!("(dec) Key: {}, IV: {}, Input: ({}, {})", key, iv, input.len(), input.hex());
-
-    let mut dec = cha::decrypter(key, iv)?;
-
-    let n = dec.update(&input[..], &mut output[..])?;
-    eprintln!("(dec) Written {} bytes", n);
-
-    println!(">> {}", (&output[..n]).hex());
-    assert!(dec.finalize(&mut output[..n])? == 0);
-    //    assert!(dec.finalize(&mut output[..n])? == 0);
-    println!(">> {}", (&output[..n]).hex());
+    let mut args = std::env::args();
+    let prog_name = args.next().unwrap();
 
-    output.truncate(n);
-    Ok(output)
-}*/
-
-fn keys() -> Result<(Key, IV), base64::DecodeError>
-{
-    let mut args = std::env::args().skip(1);
+    let mode = match args.next()
+	.map(|x| x.chars().next().map(|x| x.to_ascii_lowercase()))
+	.flatten()
+    {
+	Some('e') => Mode::Encrypt,
+	Some('d') => Mode::Decrypt,
+	Some('k') => {
+	    let (key, iv) = cha::keygen();
+	    return Ok((Mode::Keygen, key, iv));
+	},
+	_ => {
+	    eprintln!("{} (v{}) - chacha20_poly1305 command line encryption tool",
+		      env!("CARGO_PKG_NAME"),
+		      env!("CARGO_PKG_VERSION"));
+	    eprintln!(" by {} with <3 (licensed GPL v3.0 or later)", env!("CARGO_PKG_AUTHORS"));
+	    eprintln!("\nStreams stdin to stdout through a chacha20_poly1305 cipher.");
+	    eprintln!();
+	    eprintln!("Usage: {} encrypt [<base64 key>] [<base64 iv>]", prog_name);
+	    eprintln!("Usage: {} decrypt [<base64 key>] [<base64 iv>]", prog_name);
+	    eprintln!("Usage: {} keygen", prog_name);
+	    eprintln!();
+	    eprintln!("(Key size is {}, IV size is {})\n", cha::KEY_SIZE, cha::IV_SIZE);
+	    eprintln!("encrypt/decrypt:\n\tIf key and/or IV are not provided, they are generated randomly and printed to stderr in order on one line each");
+	    eprintln!("keygen:\n\tThe key/iv is printed in the same way as auto-generated keys for the en/decryption modes, but to stdout instead of stderr");
+	    std::process::exit(1)
+	}
+    };
     
     let key = match args.next() {
 	Some(key) => key.parse()?,
@@ -62,27 +59,39 @@ fn keys() -> Result<(Key, IV), base64::DecodeError>
     let iv = match args.next() {
 	Some(iv) => iv.parse()?,
 	None => {
-	    let key = IV::new();
-	    eprintln!("{}", base64::encode(&key));
-	    key
+	    let iv = IV::new();
+	    eprintln!("{}", base64::encode(&iv));
+	    iv
 	},
     };
 
-    Ok((key, iv))
+    Ok((mode, key, iv))
 }
 
 fn main() {
 
-    let (key, iv) = keys().expect("Failed to read keys from argv (base64)");
+    let (mode, key, iv) = keys().expect("Failed to read keys from argv (base64)");
     
     let stdout = std::io::stdout();
     let input = std::io::stdin();
 
-    // Encryption
+    // Streaming
+    use std::io::Write;
+    match mode 
     {
-	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");
+	Mode::Encrypt => {
+	    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");
+	},
+	Mode::Decrypt => {
+	    let mut output = stream::Sink::decrypt(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");
+	},
+	Mode::Keygen => {
+	    println!("{}", base64::encode(&key));
+	    println!("{}", base64::encode(&iv));
+	},
     }
 }