Support -nHEADmain

1 files changed, 11 insertions, 2 deletions

diff --git a/diceware.go b/diceware.go
index 3716171..78fbbc4 100644
--- a/diceware.go
+++ b/diceware.go
@@ -3,16 +3,25 @@ package main
 import (
 	"crypto/rand"
 	"encoding/binary"
+	"flag"
 	"fmt"
+	"os"
 	"strings"
 
 	"roseh.moe/pkg/wordlist"
 )
 
+var n = flag.Int("n", 10, "number of words to generate")
+
 func main() {
-	buf := make([]byte, 20)
+	flag.Parse()
+	if *n < 0 {
+		fmt.Fprintln(os.Stderr, "What are you doing...")
+		os.Exit(2)
+	}
+	buf := make([]byte, 2**n)
 	rand.Read(buf)
-	words := make([]string, 10)
+	words := make([]string, *n)
 	for i := range words {
 		words[i] = wordlist.Words[binary.NativeEndian.Uint16(buf[2*i:])&0x1fff]
 	}