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  }