diff --git a/internal/aes/keygen.go b/internal/aes/keygen.go index 17dc563..3d288b5 100644 --- a/internal/aes/keygen.go +++ b/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 diff --git a/internal/aes/keygen_test.go b/internal/aes/keygen_test.go index f3671d1..0a0632e 100644 --- a/internal/aes/keygen_test.go +++ b/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) diff --git a/internal/hash/sha.go b/internal/hash/sha.go index 266ecbb..ed25b3a 100644 --- a/internal/hash/sha.go +++ b/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