From 48199c5aa831d56ec65396f3f02a4e6cc8c22836 Mon Sep 17 00:00:00 2001
From: Adrien Delorme <adrien.delorme@icloud.com>
Date: Tue, 13 Oct 2020 16:25:24 +0200
Subject: [PATCH] add basic test to see if generated files are parseable

---
 helper/communicator/sshkey/generate.go      | 31 +++++++++++----
 helper/communicator/sshkey/generate_test.go | 43 +++++++++++++++++++++
 2 files changed, 67 insertions(+), 7 deletions(-)
 create mode 100644 helper/communicator/sshkey/generate_test.go

diff --git a/helper/communicator/sshkey/generate.go b/helper/communicator/sshkey/generate.go
index 51e17627a..a977c7f46 100644
--- a/helper/communicator/sshkey/generate.go
+++ b/helper/communicator/sshkey/generate.go
@@ -47,17 +47,34 @@ func NewPair(public, private interface{}) (*Pair, error) {
 	}
 
 	privBlk := &pem.Block{
-		Type:    "OPENSSH PRIVATE KEY",
+		Type:    "PRIVATE KEY",
 		Headers: nil,
 		Bytes:   kb,
 	}
 
-	switch private.(type) {
-	case *rsa.PrivateKey:
-		privBlk.Type = "RSA PRIVATE KEY"
+	publicKey, err := ssh.NewPublicKey(public)
+	if err != nil {
+		return nil, err
 	}
+	return &Pair{
+		Private: pem.EncodeToMemory(privBlk),
+		Public:  ssh.MarshalAuthorizedKey(publicKey),
+	}, nil
+}
 
-	publicKey, err := ssh.NewPublicKey(public)
+func PairFromEC(key *ecdsa.PrivateKey) (*Pair, error) {
+	kb, err := x509.MarshalECPrivateKey(key)
+	if err != nil {
+		return nil, err
+	}
+
+	privBlk := &pem.Block{
+		Type:    "EC PRIVATE KEY",
+		Headers: nil,
+		Bytes:   kb,
+	}
+
+	publicKey, err := ssh.NewPublicKey(&key.PublicKey)
 	if err != nil {
 		return nil, err
 	}
@@ -120,7 +137,7 @@ func GeneratePair(t Algorithm, rand io.Reader, bits int) (*Pair, error) {
 	switch t {
 	case DSA:
 		if bits == 0 {
-			bits = 3072
+			bits = 1024
 		}
 		var sizes dsa.ParameterSizes
 		switch bits {
@@ -167,7 +184,7 @@ func GeneratePair(t Algorithm, rand io.Reader, bits int) (*Pair, error) {
 		if err != nil {
 			return nil, err
 		}
-		return NewPair(&ecdsakey.PublicKey, ecdsakey)
+		return PairFromEC(ecdsakey)
 	case ED25519:
 		publicKey, privateKey, err := ed25519.GenerateKey(rand)
 		if err != nil {
diff --git a/helper/communicator/sshkey/generate_test.go b/helper/communicator/sshkey/generate_test.go
new file mode 100644
index 000000000..aee35599f
--- /dev/null
+++ b/helper/communicator/sshkey/generate_test.go
@@ -0,0 +1,43 @@
+package sshkey
+
+import (
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	"golang.org/x/crypto/ssh"
+)
+
+func TestGeneratePair_parseable(t *testing.T) {
+	type args struct {
+		t Algorithm
+	}
+	tests := []struct {
+		t Algorithm
+	}{
+		{DSA},
+		{RSA},
+		{ECDSA},
+		{ED25519},
+	}
+	for _, tt := range tests {
+		t.Run(tt.t.String(), func(t *testing.T) {
+			got, err := GeneratePair(tt.t, nil, 0)
+			if err != nil {
+				t.Errorf("GeneratePair() error = %v", err)
+				return
+			}
+
+			privateKey, err := ssh.ParsePrivateKey(got.Private)
+			if err != nil {
+				t.Fatal(err)
+			}
+			publicKey, _, _, _, err := ssh.ParseAuthorizedKey(got.Public)
+			if err != nil {
+				t.Fatalf("%v: %s", err, got.Public)
+			}
+			if diff := cmp.Diff(privateKey.PublicKey().Marshal(), publicKey.Marshal()); diff != "" {
+				t.Fatalf("wrong public key: %s", diff)
+			}
+		})
+	}
+}