github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/cmd/roachprod/vm/aws/keys.go (about)

     1  // Copyright 2018 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package aws
    12  
    13  import (
    14  	"crypto/sha1"
    15  	"encoding/base64"
    16  	"fmt"
    17  	"io/ioutil"
    18  	"os"
    19  	"strings"
    20  
    21  	"github.com/cockroachdb/errors"
    22  )
    23  
    24  const sshPublicKeyFile = "${HOME}/.ssh/id_rsa.pub"
    25  
    26  // sshKeyExists checks to see if there is a an SSH key with the given name in the given region.
    27  func (p *Provider) sshKeyExists(keyName string, region string) (bool, error) {
    28  	var data struct {
    29  		KeyPairs []struct {
    30  			KeyName string
    31  		}
    32  	}
    33  	args := []string{
    34  		"ec2", "describe-key-pairs",
    35  		"--region", region,
    36  	}
    37  	err := p.runJSONCommand(args, &data)
    38  	if err != nil {
    39  		return false, err
    40  	}
    41  	for _, keyPair := range data.KeyPairs {
    42  		if keyPair.KeyName == keyName {
    43  			return true, nil
    44  		}
    45  	}
    46  	return false, nil
    47  }
    48  
    49  // sshKeyImport takes the user's local, public SSH key and imports it into the ec2 region so that
    50  // we can create new hosts with it.
    51  func (p *Provider) sshKeyImport(keyName string, region string) error {
    52  	keyBytes, err := ioutil.ReadFile(os.ExpandEnv(sshPublicKeyFile))
    53  	if err != nil {
    54  		if os.IsNotExist(err) {
    55  			return errors.Wrapf(err, "please run ssh-keygen externally to create your %s file", sshPublicKeyFile)
    56  		}
    57  		return err
    58  	}
    59  
    60  	var data struct {
    61  		KeyName string
    62  	}
    63  	_ = data.KeyName // silence unused warning
    64  	args := []string{
    65  		"ec2", "import-key-pair",
    66  		"--region", region,
    67  		"--key-name", keyName,
    68  		"--public-key-material", string(keyBytes),
    69  	}
    70  	err = p.runJSONCommand(args, &data)
    71  	// If two roachprod instances run at the same time with the same key, they may
    72  	// race to upload the key pair.
    73  	if err == nil || strings.Contains(err.Error(), "InvalidKeyPair.Duplicate") {
    74  		return nil
    75  	}
    76  	return err
    77  }
    78  
    79  // sshKeyName computes the name of the ec2 ssh key that we'll store the local user's public key in
    80  func (p *Provider) sshKeyName() (string, error) {
    81  	user, err := p.FindActiveAccount()
    82  	if err != nil {
    83  		return "", err
    84  	}
    85  
    86  	keyBytes, err := ioutil.ReadFile(os.ExpandEnv(sshPublicKeyFile))
    87  	if err != nil {
    88  		if os.IsNotExist(err) {
    89  			return "", errors.Wrapf(err, "please run ssh-keygen externally to create your %s file", sshPublicKeyFile)
    90  		}
    91  		return "", err
    92  	}
    93  
    94  	hash := sha1.New()
    95  	if _, err := hash.Write(keyBytes); err != nil {
    96  		return "", err
    97  	}
    98  	hashBytes := hash.Sum(nil)
    99  	hashText := base64.URLEncoding.EncodeToString(hashBytes)
   100  
   101  	return fmt.Sprintf("%s-%s", user, hashText), nil
   102  }