github.com/turbot/go-exec-communicator@v0.0.0-20230412124734-9374347749b6/ssh/ssh_test.go (about)

     1  package ssh
     2  
     3  import (
     4  	"bytes"
     5  	"crypto/rand"
     6  	"crypto/rsa"
     7  	"crypto/x509"
     8  	"encoding/pem"
     9  	"io/ioutil"
    10  	"os"
    11  	"path/filepath"
    12  	"testing"
    13  
    14  	"golang.org/x/crypto/ssh"
    15  )
    16  
    17  // verify that we can locate public key data
    18  func TestFindKeyData(t *testing.T) {
    19  	// set up a test directory
    20  	td := t.TempDir()
    21  	cwd, err := os.Getwd()
    22  	if err != nil {
    23  		t.Fatal(err)
    24  	}
    25  	if err := os.Chdir(td); err != nil {
    26  		t.Fatal(err)
    27  	}
    28  	defer os.Chdir(cwd)
    29  
    30  	id := "provisioner_id"
    31  
    32  	pub := generateSSHKey(t, id)
    33  	pubData := pub.Marshal()
    34  
    35  	// backup the pub file, and replace it with a broken file to ensure we
    36  	// extract the public key from the private key.
    37  	if err := os.Rename(id+".pub", "saved.pub"); err != nil {
    38  		t.Fatal(err)
    39  	}
    40  	if err := ioutil.WriteFile(id+".pub", []byte("not a public key"), 0600); err != nil {
    41  		t.Fatal(err)
    42  	}
    43  
    44  	foundData := findIDPublicKey(id)
    45  	if !bytes.Equal(foundData, pubData) {
    46  		t.Fatalf("public key %q does not match", foundData)
    47  	}
    48  
    49  	// move the pub file back, and break the private key file to simulate an
    50  	// encrypted private key
    51  	if err := os.Rename("saved.pub", id+".pub"); err != nil {
    52  		t.Fatal(err)
    53  	}
    54  
    55  	if err := ioutil.WriteFile(id, []byte("encrypted private key"), 0600); err != nil {
    56  		t.Fatal(err)
    57  	}
    58  
    59  	foundData = findIDPublicKey(id)
    60  	if !bytes.Equal(foundData, pubData) {
    61  		t.Fatalf("public key %q does not match", foundData)
    62  	}
    63  
    64  	// check the file by path too
    65  	foundData = findIDPublicKey(filepath.Join(".", id))
    66  	if !bytes.Equal(foundData, pubData) {
    67  		t.Fatalf("public key %q does not match", foundData)
    68  	}
    69  }
    70  
    71  func generateSSHKey(t *testing.T, idFile string) ssh.PublicKey {
    72  	t.Helper()
    73  
    74  	priv, err := rsa.GenerateKey(rand.Reader, 2048)
    75  	if err != nil {
    76  		t.Fatal(err)
    77  	}
    78  
    79  	privFile, err := os.OpenFile(idFile, os.O_RDWR|os.O_CREATE, 0600)
    80  	if err != nil {
    81  		t.Fatal(err)
    82  	}
    83  	defer privFile.Close()
    84  	privPEM := &pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(priv)}
    85  	if err := pem.Encode(privFile, privPEM); err != nil {
    86  		t.Fatal(err)
    87  	}
    88  
    89  	// generate and write public key
    90  	pub, err := ssh.NewPublicKey(&priv.PublicKey)
    91  	if err != nil {
    92  		t.Fatal(err)
    93  	}
    94  
    95  	err = ioutil.WriteFile(idFile+".pub", ssh.MarshalAuthorizedKey(pub), 0600)
    96  	if err != nil {
    97  		t.Fatal(err)
    98  	}
    99  
   100  	return pub
   101  }