From 4481822068bbb98bfb16751c3f13c4ba124a1e74 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=BA=90=E6=96=87=E9=9B=A8?=
 <41315874+fumiama@users.noreply.github.com>
Date: Tue, 12 Apr 2022 19:01:31 +0800
Subject: [PATCH] add commandline tool

---
 README.md    |  5 ++++
 blake2b.go   |  8 +++---
 main/main.go | 79 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 88 insertions(+), 4 deletions(-)
 create mode 100644 main/main.go

diff --git a/README.md b/README.md
index 31fcbf7..2a88293 100644
--- a/README.md
+++ b/README.md
@@ -10,6 +10,11 @@ This package was initially based on the pure go [BLAKE2b](https://github.com/dch
 
 In addition to AVX there is also support for AVX2 as well as SSE. Best performance is obtained with AVX2 which gives roughly a **4X** performance increase approaching hashing speeds of **1GB/sec** on a single core.
 
+There is also a commandline tool for blake2b, see `main` folder for more details.
+```bash
+go run main/main.go
+```
+
 Benchmarks
 ----------
 
diff --git a/blake2b.go b/blake2b.go
index b7d5a97..ac65156 100644
--- a/blake2b.go
+++ b/blake2b.go
@@ -135,8 +135,8 @@ func (d *digest) initialize(c *Config) {
 	if c.Tree != nil {
 		p[2] = c.Tree.Fanout
 		p[3] = c.Tree.MaxDepth
-		binary.LittleEndian.PutUint32(p[4:], c.Tree.LeafSize)
-		binary.LittleEndian.PutUint64(p[8:], c.Tree.NodeOffset)
+		binary.LittleEndian.PutUint32(p[4:8], c.Tree.LeafSize)
+		binary.LittleEndian.PutUint64(p[8:16], c.Tree.NodeOffset)
 		p[16] = c.Tree.NodeDepth
 		p[17] = c.Tree.InnerHashSize
 	} else {
@@ -147,7 +147,7 @@ func (d *digest) initialize(c *Config) {
 	// Initialize.
 	d.size = c.Size
 	for i := 0; i < 8; i++ {
-		d.h[i] = iv[i] ^ binary.LittleEndian.Uint64(p[i*8:])
+		d.h[i] = iv[i] ^ binary.LittleEndian.Uint64(p[i*8:(i+1)*8])
 	}
 	if c.Tree != nil && c.Tree.IsLastNode {
 		d.isLastNode = true
@@ -156,7 +156,7 @@ func (d *digest) initialize(c *Config) {
 	// Process key.
 	if c.Key != nil {
 		copy(d.paddedKey[:], c.Key)
-		d.Write(d.paddedKey[:])
+		d.Write(c.Key)
 		d.isKeyed = true
 	}
 	// Save a copy of initialized state.
diff --git a/main/main.go b/main/main.go
new file mode 100644
index 0000000..5435e00
--- /dev/null
+++ b/main/main.go
@@ -0,0 +1,79 @@
+package main
+
+import (
+	"encoding/hex"
+	"flag"
+	"fmt"
+	"io"
+	"os"
+
+	blake2b "github.com/fumiama/blake2b-simd"
+)
+
+func main() {
+	help := flag.Bool("h", false, "display this help")
+	size := flag.Int("s", 32, "digest output bytes size")
+	key := flag.String("k", "", "hex-encoded key for prefix-MAC")
+	var keyb []byte
+	var err error
+	if *key != "" {
+		keyb, err = hex.DecodeString(*key)
+		if err != nil {
+			panic(err)
+		}
+	}
+	salt := flag.String("t", "", "hex-encoded 16 bytes salt (if < 16 bytes, padded with zeros)")
+	var saltb []byte
+	if *salt != "" {
+		saltb, err = hex.DecodeString(*salt)
+		if err != nil {
+			panic(err)
+		}
+	}
+	pers := flag.String("p", "", "hex-encoded 16 bytes personalization (if < 16 bytes, padded with zeros)")
+	var persb []byte
+	if *pers != "" {
+		persb, err = hex.DecodeString(*pers)
+		if err != nil {
+			panic(err)
+		}
+	}
+	flag.Parse()
+	if *help {
+		fmt.Println(os.Args[0], "[commands] [file1 file2...(null for stdin)]")
+		flag.PrintDefaults()
+		return
+	}
+	h, err := blake2b.New(&blake2b.Config{
+		Size:   uint8(*size),
+		Key:    keyb,
+		Salt:   saltb,
+		Person: persb,
+	})
+	if err != nil {
+		panic(err)
+	}
+	file := flag.Args()
+	var f []io.Reader
+	if len(file) == 0 {
+		f = append(f, os.Stdin)
+		file = append(file, "stdin")
+	} else {
+		for _, name := range file {
+			fi, err := os.Open(name)
+			if err != nil {
+				panic(err)
+			}
+			f = append(f, fi)
+		}
+	}
+	for i, fi := range f {
+		_, err = io.Copy(h, fi)
+		if err != nil {
+			panic(err)
+		}
+		sum := make([]byte, 0, *size)
+		fmt.Println(file[i], hex.EncodeToString(h.Sum(sum)))
+		h.Reset()
+	}
+}