From df87aac2faf223a205c68fd5d0c62018f08c7eea Mon Sep 17 00:00:00 2001
From: Hatter Jiang <jht5945@gmail.com>
Date: Sun, 24 Nov 2019 13:04:37 +0800
Subject: [PATCH] add BLAKE2S, BLAKE2B

---
 src/main.rs | 18 ++++++++++--------
 src/opt.rs  |  3 +++
 2 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 96dab99..12e096d 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -16,6 +16,8 @@ use crypto::{
     digest::Digest,
     ripemd160::Ripemd160,
     whirlpool::Whirlpool,
+    blake2s::Blake2s,
+    blake2b::Blake2b,
     md5::Md5,
     sha1::Sha1,
     sha2::{
@@ -33,9 +35,7 @@ use indicatif::{
     ProgressStyle
 };
 use libsm::sm3::hash::Sm3Hash;
-use opt::{
-    Options,
-};
+use opt::Options;
 
 const FILE_SIZE_1GB: u64 = 1024 * 1024 * 1024;
 const BUFF_SIZE: usize = 512 * 1024;
@@ -55,7 +55,7 @@ fn main() -> XResult<()> {
 
     let the_algo = options.algorithm.to_uppercase();
     if options.file_name_list.is_empty() {
-        let boxed_digest = get_digest_by_algorithm(the_algo.as_str());
+        let boxed_digest = get_digest_by_algorithm(the_algo.as_str(), &options);
         match boxed_digest {
             Some(mut digest) => {
                 println!("{}  - ({})", calc_digest_stdin(&mut *digest)?, the_algo);
@@ -66,9 +66,9 @@ fn main() -> XResult<()> {
             },
         };
     } else {
-        for f in options.file_name_list {
+        for f in &options.file_name_list {
             let the_fn = f.as_str();
-            let boxed_digest = get_digest_by_algorithm(the_algo.as_str());
+            let boxed_digest = get_digest_by_algorithm(the_algo.as_str(), &options);
             match boxed_digest {
                 Some(mut digest) => {
                     println!("{}  - {} ({})", calc_file_digest(&mut *digest, the_fn)?, the_fn, the_algo);
@@ -95,14 +95,16 @@ Supported algorithms:
 MD5, SHA1, SHA224, SHA256, SHA512, SHA512-224, SHA512-256,
 SHA3-224, SHA3-256, SHA3-384, SHA3-512, SHAKE-128, SHAKE-256,
 KECCAK-224, KECCAK-256, KECCAK-384, KECCAK-512,
-SM3, RIPEMD160, WHIRLPOOL
+SM3, RIPEMD160, WHIRLPOOL, BLAKE2S, BLAKE2B
 "#, VERSION, &GIT_HASH[0..7]);
 }
 
-fn get_digest_by_algorithm(algo: &str) -> Option<Box<dyn Digest>> {
+fn get_digest_by_algorithm(algo: &str, options: &Options) -> Option<Box<dyn Digest>> {
     match algo {
         "RIPEMD160" => Some(Box::new(Ripemd160::new())),
         "WHIRLPOOL" => Some(Box::new(Whirlpool::new())),
+        "BLAKE2S" => Some(Box::new(Blake2s::new(if options.blake_len == 0_usize { 32 }  else { options.blake_len }))),
+        "BLAKE2B" => Some(Box::new(Blake2b::new(if options.blake_len == 0_usize { 64 }  else { options.blake_len }))),
         "MD5" => Some(Box::new(Md5::new())),
         "SHA1" | "SHA-1" => Some(Box::new(Sha1::new())),
         "SHA224" | "SHA-224" => Some(Box::new(Sha224::new())),
diff --git a/src/opt.rs b/src/opt.rs
index eca2ed9..378bcdd 100644
--- a/src/opt.rs
+++ b/src/opt.rs
@@ -5,6 +5,7 @@ pub struct Options {
     pub version: bool,
     pub algorithm: String,
     pub file_name_list: Vec<String>,
+    pub blake_len: usize,
 }
 
 impl Options {
@@ -13,6 +14,7 @@ impl Options {
             version: false,
             algorithm: "SHA256".to_string(),
             file_name_list: Vec::new(),
+            blake_len: 0_usize,
         }
     }
 
@@ -21,6 +23,7 @@ impl Options {
             let mut ap = ArgumentParser::new();
             ap.set_description("digest - command line digest tool.");
             ap.refer(&mut self.algorithm).add_option(&["-a", "--algorithm"], Store, "Algorithm, e.g. SHA256, SM3");
+            ap.refer(&mut self.blake_len).add_option(&["-l", "--blake-len"], Store, "Blake2s/b length, 1~32/64");
             ap.refer(&mut self.version).add_option(&["-v", "--version"], StoreTrue, "Print version");
             ap.refer(&mut self.file_name_list).add_argument("File names", List, "File names to be digested");
             ap.parse_args_or_exit();