Change keygen behaviour for scripting

cmd/keygen/keygen.go

@@ -3,6 +3,7 @@ package main
 import (
 	"crypto/rand"
 	"encoding/base64"
+	"flag"
 	"fmt"
 	ecdh "github.com/cbeuw/go-ecdh"
 )
@@ -10,38 +11,23 @@ import (
 var b64 = base64.StdEncoding.EncodeToString
 
 func main() {
-	for {
-		fmt.Println("1 to generate UID, 2 to generate a key pair")
+	var isUID *bool
+	var isKeypair *bool
+	isUID = flag.Bool("u", false, "Generate UID")
+	isKeypair = flag.Bool("k", false, "Generate a key pair")
+	flag.Parse()
 
-		var sel int
-		_, err := fmt.Scanln(&sel)
-		if err != nil {
-			fmt.Println("Please enter a number")
-			continue
-		}
-		if sel != 1 && sel != 2 {
-			fmt.Println("Please enter 1 or 2")
-			continue
-		}
+	if *isUID {
+		UID := make([]byte, 32)
+		rand.Read(UID)
+		fmt.Printf(b64(UID))
+	} else if *isKeypair {
+		ec := ecdh.NewCurve25519ECDH()
+		staticPv, staticPub, _ := ec.GenerateKey(rand.Reader)
+		marshPub := ec.Marshal(staticPub)
+		marshPv := staticPv.(*[32]byte)[:]
 
-		if sel == 1 {
-			UID := make([]byte, 32)
-			rand.Read(UID)
-			fmt.Printf("\"UID\":\"%v\"\n", b64(UID))
-		} else if sel == 2 {
+		fmt.Printf("%v,%v", b64(marshPub), b64(marshPv))
 
-			ec := ecdh.NewCurve25519ECDH()
-			staticPv, staticPub, _ := ec.GenerateKey(rand.Reader)
-			marshPub := ec.Marshal(staticPub)
-			marshPv := staticPv.(*[32]byte)[:]
-
-			fmt.Printf("USER: \n")
-			fmt.Printf("\"PublicKey\":\"%v\"\n", b64(marshPub))
-
-			fmt.Println("=========================================")
-
-			fmt.Printf("SERVER: \n")
-			fmt.Printf("\"PrivateKey\":\"%v\"\n", b64(marshPv))
-		}
 	}
 }