main.rs

#![deny(unused_results)]

#[inline]
fn u64_usize(n: u64) -> usize {
    n.try_into().expect("FATAL: u64 length to usize error")
}

#[inline]
fn usize_u64(n: usize) -> u64 {
    n.try_into().expect("FATAL: usize length to u64 error")
}

use cshake::{CShake, CShakeCustom, cshake_customs, Absorb, Squeeze, SqueezeXor, SqueezeSkip};

cshake_customs! {
    CIPHER_CUSTOM -> "__shakenc__file-stream-cipher"
    HASH_CUSTOM -> "__shakenc__file-hash"
    RAND_CUSTOM -> "__shakenc__random-generator"
    OS_RAND_MASK_CUSTOM -> "__shakenc__os-rand-mask"
}

struct Context {
    cipher: CShake<CIPHER_CUSTOM>,
    ihash: Option<CShake<HASH_CUSTOM>>,
    ohash: Option<CShake<HASH_CUSTOM>>,
}

trait OptionExec<T> {
    fn exec<F: FnOnce(&mut T)>(&mut self, f: F);
}

impl<T> OptionExec<T> for Option<T> {
    fn exec<F: FnOnce(&mut T)>(&mut self, f: F) {
        if let Some(x) = self {
            f(x);
        }
    }
}

impl Context {
    fn init(key: &[u8], ihash: bool, ohash: bool) -> Self {
        Self {
            cipher: CIPHER_CUSTOM.create().chain_absorb(key),
            ihash: ihash.then(|| HASH_CUSTOM.create()),
            ohash: ohash.then(|| HASH_CUSTOM.create()),
        }
    }

    fn next(&mut self, buf: &mut [u8]) {
        self.ihash.exec(|ctx| ctx.absorb(buf));
        self.cipher.squeeze_xor(buf);
        self.ohash.exec(|ctx| ctx.absorb(buf));
    }

    fn finish<const N: usize>(self) -> HashResult<N> {
        HashResult {
            ihash: self.ihash.map(|mut ctx| ctx.squeeze_to_array()),
            ohash: self.ohash.map(|mut ctx| ctx.squeeze_to_array()),
        }
    }
}

struct HashResult<const N: usize> {
    ihash: Option<[u8; N]>,
    ohash: Option<[u8; N]>,
}

impl<const N: usize> std::fmt::Display for HashResult<N> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        if let Some(ihash) = self.ihash {
            f.write_str("input  hash: ")?;
            f.write_str(&hex::encode(ihash))?;
            f.write_str("\n")?;
        };
        if let Some(ohash) = self.ohash {
            f.write_str("output hash: ")?;
            f.write_str(&hex::encode(ohash))?;
            f.write_str("\n")?;
        };
        Ok(())
    }
}

use std::{num::NonZeroUsize, fs::OpenOptions, io::{Read, Write, self}, path::PathBuf, sync::mpsc};
use zeroize::Zeroizing;
use indicatif::{ProgressBar, ProgressStyle, HumanBytes};

#[derive(argh::FromArgs)]
/// shakenc: cSHAKE256 encrypt & random generating
struct Args {
    /// key (if not provided in arguments, you will need to enter them later)
    #[argh(option, short = 'k')]
    key: Option<String>,
    /// treat input key as hex (SHOULD use with files generated by --rand-key)
    #[argh(switch)]
    hex_key: bool,
    /// use random 256bit HEX key generated by OS (vaild when --key or -k not provided)
    #[argh(switch)]
    rand_key: bool,
    /// buffer size (MiB, default 16MiB, will take this size of runtime memory)
    #[argh(option)]
    buf: Option<NonZeroUsize>,
    /// allow overwrite existing files
    #[argh(switch)]
    overwrite: bool,
    #[argh(subcommand)]
    sub: Commands,
}

#[derive(argh::FromArgs)]
#[argh(subcommand)]
enum Commands {
    Crypt(Crypt),
    Rng(Rng),
    Rnv(Rnv),
    Flip(Flip),
}

#[derive(argh::FromArgs)]
#[argh(subcommand, name = "crypt")]
/// cSHAKE256 as a stream cipher for file encrypt/decrypt
/// 
/// IMPORTANT WARNING: One of the purposes of shakenc is to simply show how cryptography works, so it takes the key
/// input directly as input to the cshake256 function. In other words, in the case of us using the cshake256 function
/// as a cipher, the key does not go through a key derivation function (KDF) and is not salted with a non-fixed salt
/// (I don't think a fixed cshake custom counts as a salt). For using rand-key option or other sources of random
/// entropy as a key, this is generally fine. However, if the key input has a low entropy (e.g. a string), then the
/// encryption can be broken by brute force, and is therefore not suitable for serious encryption situations.
struct Crypt {
    /// input file path
    #[argh(option, short = 'i')]
    input: PathBuf,
    /// output file path (none for verify by hash)
    #[argh(option, short = 'o')]
    output: Option<PathBuf>,
    /// hash input file
    #[argh(switch)]
    ih: bool,
    /// hash output file
    #[argh(switch)]
    oh: bool,
}

#[derive(argh::FromArgs)]
#[argh(subcommand, name = "rng")]
/// cSHAKE256 as a reproduceable random generator (generate files with random bits or test the speed)
struct Rng {
    /// output file path (none for test generating speed)
    #[argh(option, short = 'o')]
    output: Option<PathBuf>,
    /// output file length (MiB)
    #[argh(option, short = 'l')]
    len: u64,
}

#[derive(argh::FromArgs)]
#[argh(subcommand, name = "rnv")]
/// cSHAKE256 as a reproduceable random generator (verify generated files)
struct Rnv {
    /// input file path
    #[argh(option, short = 'i')]
    input: PathBuf,
    /// count error bytes after first error occurred rather than stop immediately
    #[argh(switch)]
    count_err: bool,
}

#[derive(argh::FromArgs)]
#[argh(subcommand, name = "flip")]
/// cSHAKE256 as a reproduceable random generator to flip a coin
struct Flip {}

enum KeyInput {
    ArgString(Zeroizing<Vec<u8>>),
    ArgHex(Zeroizing<Vec<u8>>),
    PromptString,
    PromptHex,
    GetFromOS,
}

impl KeyInput {
    fn from_args(key: Option<Zeroizing<Vec<u8>>>, rand_key: bool, hex_key: bool) -> Self {
        use KeyInput::*;
        match (key, rand_key, hex_key) {
            (Some(key), _, false) => ArgString(key),
            (Some(key), _, true) => ArgHex(key),
            (None, false, false) => PromptString,
            (None, false, true) => PromptHex,
            (None, true, _) => GetFromOS,
        }
    }

    fn process(self) -> Zeroizing<Vec<u8>> {
        #[inline]
        fn prompt_key(prompt: &str) -> Zeroizing<String> {
            match secprompt::prompt_password(prompt) {
                Ok(val) => val,
                Err(err) => {
                    if err.kind() == io::ErrorKind::UnexpectedEof {
                        std::process::exit(0);
                    } else {
                        panic!("fatal: {:?}", err)
                    }
                }
            }
        }

        #[inline]
        fn hex_decode(key: &[u8]) -> Zeroizing<Vec<u8>> {
            hex::decode(key).expect("decode key hex error").into()
        }

        match self {
            KeyInput::ArgString(key) => key,
            KeyInput::ArgHex(key) => hex_decode(&key),
            // TODO: Zeroizing::map https://github.com/RustCrypto/utils/issues/947
            KeyInput::PromptString => prompt_key("key: ").as_bytes().to_owned().into(),
            KeyInput::PromptHex => hex_decode(prompt_key("key (hex): ").as_bytes()),
            KeyInput::GetFromOS => {
                let mut oskey: Zeroizing<[u8; 32]> = Default::default();
                getrandom::getrandom(oskey.as_mut()).expect("failed to get random key generated by OS");
                let mut key: Zeroizing<Vec<u8>> = vec![0; 32].into();
                OS_RAND_MASK_CUSTOM.once(oskey.as_ref(), key.as_mut());
                let encoded = Zeroizing::new(hex::encode(&key));
                println!("key (hex): {}", encoded.as_str());
                key
            },
        }
    }
}

#[cfg(target_arch = "aarch64")]
cpufeatures::new!(armv8_sha3_intrinsics, "sha3");

// TODO(upstream): template codegen & bar max width
const PROGRESS_TEMPLATE: &str = "{bar:60} {percent}% {bytes}/{total_bytes} {elapsed_precise}/{duration_precise} {bytes_per_sec} ETA={eta}";

fn main() {
    let Args { key, rand_key, hex_key, buf: buf_len, overwrite, sub } = argh::from_env();
    let key = key.map(|key| key.into_bytes().into());

    let (close_tx, close_rx) = mpsc::sync_channel::<()>(0);
    ctrlc::set_handler(move || close_tx.send(()).expect("fatal")).expect("fatal");
    let progress_style = ProgressStyle::with_template(PROGRESS_TEMPLATE).expect("fatal");

    #[cfg(target_arch = "aarch64")]
    if armv8_sha3_intrinsics::get() {
        eprintln!("ARMv8 SHA3 feature detected");
    }

    let key = KeyInput::from_args(key, rand_key, hex_key).process();
    let buf_len = buf_len.map(NonZeroUsize::get).unwrap_or(16) * 1048576;
    // TODO: add mode-specific determination? (https://users.rust-lang.org/t/99470)
    let mut buf = vec![0u8; buf_len];

    if let Ok(()) = close_rx.try_recv() {
        eprintln!("aborted");
        return;
    }

    match sub {
        Commands::Crypt(Crypt { input, output, ih: ihash, oh: ohash }) => {
            let mut ctx = Context::init(&key, ihash, ohash);
            let mut input = OpenOptions::new().read(true).open(input).expect("failed to open input file");
            let mut output = output.map(|output| OpenOptions::new().create_new(!overwrite).write(true).open(output).expect("failed to open output file"));
            let len = input.metadata().expect("fatal").len();
            let mut progress = 0;
            let progress_bar = ProgressBar::new(len);
            progress_bar.set_style(progress_style);

            loop {
                macro_rules! ioop {
                    ($op:expr, $f:expr) => {
                        match $op {
                            Ok(val) => val,
                            Err(err) => {
                                eprintln!("aborted at byte {} ({}) by {} file error {:?}", progress, HumanBytes(progress), $f, err);
                                break;
                            },
                        }
                    };
                }
                if let Ok(()) = close_rx.try_recv() {
                    eprintln!("aborted at byte {} ({})", progress, HumanBytes(progress));
                    break;
                }
                let read_len = ioop!(input.read(&mut buf), "input");
                if read_len != 0 {
                    // buf == buf[..read_len] when buf_len == read_len
                    let buf = &mut buf[..read_len];
                    ctx.next(buf);
                    if let Some(output) = output.as_mut() {
                        ioop!(output.write_all(buf), "output");
                    }
                    progress += usize_u64(read_len);
                    progress_bar.inc(usize_u64(read_len));
                } else {
                    // must be EOF beacuse buf_len != 0
                    assert_eq!(progress, len);
                    eprintln!("finished");
                    println!("{}", ctx.finish::<32>());
                    break;
                }
            }
        },

        Commands::Rng(Rng { output, len }) => {
            let mut ctx = RAND_CUSTOM.create().chain_absorb(&key);
            let mut output = output.map(|output| OpenOptions::new().create_new(!overwrite).write(true).open(output).expect("failed to open output file"));
            let len = len * 1048576;
            let mut progress = 0;
            let progress_bar = ProgressBar::new(len);
            progress_bar.set_style(progress_style);

            loop {
                macro_rules! ioop {
                    ($op:expr) => {
                        match $op {
                            Ok(val) => val,
                            Err(err) => {
                                eprintln!("aborted at byte {} ({}) by file error {:?}", progress, HumanBytes(progress), err);
                                break;
                            },
                        }
                    };
                }
                if let Ok(()) = close_rx.try_recv() {
                    eprintln!("aborted at byte {} ({})", progress, HumanBytes(progress));
                    break;
                }
                if (len - progress) != 0 {
                    let write_len = buf_len.min(u64_usize(len - progress));
                    let buf = &mut buf[..write_len];
                    if let Some(output) = output.as_mut() {
                        ctx.squeeze(buf);
                        ioop!(output.write_all(buf));
                    } else {
                        ctx.squeeze_skip(write_len);
                    }
                    progress += usize_u64(write_len);
                    progress_bar.inc(usize_u64(write_len));
                } else {
                    assert_eq!(progress, len);
                    eprintln!("finished");
                    break;
                }
            }
        },

        Commands::Rnv(Rnv { input, count_err }) => {
            let mut ctx = RAND_CUSTOM.create().chain_absorb(&key);
            let mut input = OpenOptions::new().read(true).open(input).expect("failed to open input file");
            let len = input.metadata().expect("fatal").len();
            let mut progress = 0;
            let mut err: u64 = 0;
            let progress_bar = ProgressBar::new(len);
            progress_bar.set_style(progress_style);

            macro_rules! eprintln_werr {
                ($($arg:tt)*) => {
                    eprint!($($arg)*);
                    if count_err & (err != 0) {
                        eprintln!(" with {} byte(s) ({}) error", err, HumanBytes(err));
                    } else {
                        eprintln!();
                    }
                };
            }

            'main: loop {
                macro_rules! ioop {
                    ($op:expr) => {
                        match $op {
                            Ok(val) => val,
                            Err(e) => {
                                eprintln_werr!("aborted at byte {} ({}) by file error {:?}", progress, HumanBytes(progress), e);
                                break;
                            },
                        }
                    };
                }
                if let Ok(()) = close_rx.try_recv() {
                    eprintln_werr!("aborted at byte {} ({})", progress, HumanBytes(progress));
                    break;
                }
                let read_len = ioop!(input.read(&mut buf));
                if read_len != 0 {
                    // buf == buf[..read_len] when buf_len == read_len
                    let buf = &mut buf[..read_len];
                    ctx.squeeze_xor(buf);
                    // TODO: if err != 0 then not enumerate()? necessary?
                    for (pos, b) in buf.iter_mut().enumerate() {
                        if *b != 0 {
                            let actual_pos = progress + usize_u64(pos);
                            if count_err {
                                if err == 0 {
                                    eprintln!("error occurred at byte {} ({})", actual_pos, HumanBytes(actual_pos));
                                }
                                err += 1;
                            } else {
                                eprintln!("aborted by error occurred at byte {} ({})", actual_pos, HumanBytes(actual_pos));
                                break 'main;
                            }
                        }
                    }
                    progress += usize_u64(read_len);
                    progress_bar.inc(usize_u64(read_len));
                } else {
                    // must be EOF beacuse buf_len != 0
                    assert_eq!(progress, len);
                    eprintln_werr!("finished");
                    break;
                }
            }
        },
        Commands::Flip(Flip { }) => {
            let mut ctx = RAND_CUSTOM.create().chain_absorb(&key);
            let mut out = [0; 1];
            ctx.squeeze(&mut out);
            let [out] = out;
            if out <= 127 {
                println!("HEADS");
            } else {
                println!("TAILS");
            }
        },
    }
}