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