From b6efe28a9676f95ff83f38674e3f58e40e8169c5 Mon Sep 17 00:00:00 2001
From: Adrien Delorme <adrien.delorme@icloud.com>
Date: Tue, 13 Oct 2020 12:06:49 +0200
Subject: [PATCH] better cmd

---
 cmd/ssh-keygen/main.go | 77 ++++++++++++++++++++++++++++++------------
 1 file changed, 56 insertions(+), 21 deletions(-)

diff --git a/cmd/ssh-keygen/main.go b/cmd/ssh-keygen/main.go
index eca451ffe..20262d001 100644
--- a/cmd/ssh-keygen/main.go
+++ b/cmd/ssh-keygen/main.go
@@ -2,33 +2,43 @@ package main
 
 import (
 	"flag"
+	"io/ioutil"
 	"log"
 	"os"
+	"os/user"
+	"path/filepath"
 
 	"github.com/hashicorp/packer/helper/communicator/sshkey"
 )
 
 type options struct {
-	Type string
-	Bits int
+	Type     string
+	Bits     int
+	Filename string
 }
 
 func (o *options) AddFlagSets(fs *flag.FlagSet) {
 	fs.StringVar(&o.Type, "type", "rsa", `dsa | ecdsa | ed25519 | rsa
-
-	Specifies the type of key to create.  The possible values are 'dsa', 'ecdsa',
-	'ed25519', or 'rsa' ( the default ).
+Specifies the type of key to create. The possible values are 'dsa', 'ecdsa',
+'ed25519', or 'rsa'.
+`)
+	fs.IntVar(&o.Bits, "bits", 0, `Specifies the number of bits in the key to create. By default maximum 
+number will be picked. For RSA keys, the minimum size is 1024 bits and the 
+default is 3072 bits. Generally, 3072 bits is considered sufficient. DSA 
+keys must be exactly 1024 bits as specified by FIPS 186-2. For ECDSA keys, 
+the bits flag determines the key length by selecting from one of three 
+elliptic curve sizes: 256, 384 or 521 bits. Attempting to use bit lengths
+other than these three values for ECDSA keys will fail. Ed25519 keys have a
+fixed length and the bits flag will be ignored.
 `)
-	fs.IntVar(&o.Bits, "bits", 0, `bits
-
-	Specifies the number of bits in the key to create.  For RSA keys, the min-
-	imum size is 1024 bits and the default is 3072 bits.  Generally, 3072 bits
-	is considered sufficient.  DSA keys must be exactly 1024 bits as specified
-	by FIPS 186-2.  For ECDSA keys, the bits flag determines the key length by
-	selecting from one of three elliptic curve sizes: 256, 384 or 521 bits.
-	Attempting to use bit lengths other than these three values for ECDSA keys
-	will fail.  Ed25519 keys have a fixed length and the bits flag will be
-	ignored.	
+
+	defaultPath := ""
+	user, err := user.Current()
+	if err == nil {
+		defaultPath = filepath.Join(user.HomeDir, ".ssh", "tests")
+	}
+
+	fs.StringVar(&o.Filename, "filename", defaultPath, `Specifies the filename of the key file.
 `)
 }
 
@@ -39,8 +49,7 @@ func main() {
 	cla := options{}
 	cla.AddFlagSets(fs)
 	if err := fs.Parse(os.Args[1:]); err != nil {
-		fs.Usage()
-		os.Exit(1)
+		log.Fatal(err)
 	}
 
 	algo, err := sshkey.AlgorithmString(cla.Type)
@@ -48,13 +57,39 @@ func main() {
 		log.Fatal(err)
 	}
 
+	log.Printf("Generating public/private %s key pair.", algo)
+
 	keypair, err := sshkey.GeneratePair(algo, nil, cla.Bits)
 	if err != nil {
 		log.Fatal(err)
 	}
 
-	log.Printf("keypair.Private:")
-	log.Printf("%s", keypair.Private)
-	log.Printf("keypair.Public:")
-	log.Printf("%s", keypair.Public)
+	if isDir(cla.Filename) {
+		cla.Filename = filepath.Join(cla.Filename, "id_"+algo.String())
+	}
+	if fileExists(cla.Filename) {
+		log.Fatalf("%s already exists.", cla.Filename)
+	}
+	log.Printf("Saving private key to %s", cla.Filename)
+	if err := ioutil.WriteFile(cla.Filename, keypair.Private, 0600); err != nil {
+		log.Fatal(err)
+	}
+	publicFilename := cla.Filename + ".pub"
+	log.Printf("Saving public key to %s", publicFilename)
+	if err := ioutil.WriteFile(publicFilename, keypair.Public, 0644); err != nil {
+		log.Fatal(err)
+	}
+}
+
+func isDir(filename string) bool {
+	info, err := os.Stat(filename)
+	if err != nil {
+		log.Fatal(err)
+	}
+	return info.IsDir()
+}
+
+func fileExists(filename string) bool {
+	_, err := os.Stat(filename)
+	return err == nil
 }