internal/aes: Add -bin flag to keygen

internal/aes/keygen.go

@@ -1,7 +1,6 @@
 package aes
 
 import (
-	"bytes"
 	"crypto/rand"
 	"encoding/hex"
 	"errors"
@@ -15,7 +14,7 @@ import (
 func KeyGen(args []string) error {
 	fset := flag.NewFlagSet("aes-keygen", flag.ContinueOnError)
 	fset.Usage = func() {
-		fmt.Fprint(os.Stderr, `Usage: pocryp aes-keygen [-out OUTPUT] NUM_BITS
+		fmt.Fprint(os.Stderr, `Usage: pocryp aes-keygen [-out OUTPUT] [-bin] NUM_BITS
 
 Generate AES key.
 Valid NUM_BITS: 128, 192, 256.
@@ -28,6 +27,7 @@ Options:
 	}
 
 	fOutput := fset.String("out", "", "Write the result to the file at path OUTPUT.")
+	fBin := fset.Bool("bin", false, "Write output as binary not hex.")
 
 	if err := fset.Parse(args); err != nil {
 		return err
@@ -47,27 +47,31 @@ Options:
 		fset.Usage()
 		return errors.New("invalid num bits requested")
 	}
-
 	numBits /= 8
 
-	output := make([]byte, numBits)
-	if _, err := rand.Read(output); err != nil {
-		return err
-	}
-
+	var w io.Writer
 	if *fOutput == "" {
-		fmt.Println(hex.EncodeToString(output))
-		return nil
+		w = os.Stdout
+	} else {
+		f, err := os.Create(*fOutput)
+		if err != nil {
+			return err
+		}
+		defer f.Close()
+		w = f
 	}
 
-	f, err := os.Create(*fOutput)
-	if err != nil {
+	output := make([]byte, numBits)
+	if _, err := rand.Read(output); err != nil {
 		return err
 	}
-	defer f.Close()
 
-	if _, err := io.Copy(f, bytes.NewBuffer(output)); err != nil {
-		return err
+	if *fBin {
+		if _, err := w.Write(output); err != nil {
+			return err
+		}
+	} else {
+		fmt.Fprintln(w, hex.EncodeToString(output))
 	}
 
 	return nil

internal/aes/keygen_test.go

@@ -31,7 +31,7 @@ func TestKeyGen(t *testing.T) {
 	for _, numBits := range tests {
 		t.Run(numBits, func(t *testing.T) {
 			outPath := filepath.Join(tmp, "out"+numBits)
-			if err := KeyGen([]string{"-out", outPath, numBits}); err != nil {
+			if err := KeyGen([]string{"-bin", "-out", outPath, numBits}); err != nil {
 				t.Fatalf("%s: %v", numBits, err)
 			}
 			numBitsInt, _ := strconv.Atoi(numBits)

internal/hash/sha.go

@@ -83,8 +83,7 @@ Options:
 			return err
 		}
 	} else {
-		output := hex.EncodeToString(digest)
-		fmt.Fprintln(w, output)
+		fmt.Fprintln(w, hex.EncodeToString(digest))
 	}
 
 	return nil