encrypt, decrypt, keygen all work

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"
 

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());
-
-    output.truncate(n);
-    Ok(output)
-}*/
+    let mut args = std::env::args();
+    let prog_name = args.next().unwrap();
 
-fn keys() -> Result<(Key, IV), base64::DecodeError>
+    let mode = match args.next()
+	.map(|x| x.chars().next().map(|x| x.to_ascii_lowercase()))
+	.flatten()
     {
-    let mut args = std::env::args().skip(1);
+	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 
+    {
+	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));
+	},
     }
 }