dawid/speedrun
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 8 Wiki Activity

Initial commit

Browse Source
This commit is contained in:
Dawid Pogorzelski
2021-03-21 23:10:46 +01:00
commit 3144420a47
22 changed files with 2364 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
cmd/init.go Normal file
View File

@@ -0,0 +1,15 @@
package cmd
import (
"speedrun/config"
"github.com/spf13/cobra"
)
var initCmd = &cobra.Command{
Use: "init",
Short: "Initialize speedrun",
RunE: func(cmd *cobra.Command, args []string) error {
return config.Create()
},
}

241
cmd/key.go Normal file
View File

@@ -0,0 +1,241 @@
package cmd
import (
"crypto/ed25519"
"crypto/rand"
"crypto/x509"
"encoding/pem"
"fmt"
"io/ioutil"
"path/filepath"
gcp "speedrun/cloud"
"github.com/alitto/pond"
"github.com/apex/log"
homedir "github.com/mitchellh/go-homedir"
"github.com/spf13/cobra"
"github.com/spf13/viper"
"golang.org/x/crypto/ssh"
)
var keyCmd = &cobra.Command{
Use: "key",
Short: "Manage ssh keys",
TraverseChildren: true,
}
var newKeyCmd = &cobra.Command{
Use: "new",
Short: "Create a new ssh key",
PreRun: func(cmd *cobra.Command, args []string) {
initConfig()
},
RunE: newKey,
}
var setKeyCmd = &cobra.Command{
Use: "set",
Short: "Set key in the project or instance metadata",
Example: " speedrun key set \n speedrun key set --filter \"labels.foo = bar AND labels.environment = staging\"",
PreRun: func(cmd *cobra.Command, args []string) {
initConfig()
},
RunE: setKey,
}
var removeKeyCmd = &cobra.Command{
Use: "remove",
Short: "Remove key from the project metadata or instance metadata",
Example: " speedrun key remove \n speedrun key remove --filter \"labels.foo = bar AND labels.environment = staging\"",
PreRun: func(cmd *cobra.Command, args []string) {
initConfig()
},
RunE: removeKey,
}
func init() {
setKeyCmd.Flags().String("filter", "", "Set the key only on matching instances")
removeKeyCmd.Flags().String("filter", "", "Set the key only on matching instances")
keyCmd.AddCommand(newKeyCmd)
keyCmd.AddCommand(setKeyCmd)
keyCmd.AddCommand(removeKeyCmd)
}
func determineKeyFilePath() (string, error) {
log.Debug("Determining private key path")
home, err := homedir.Dir()
if err != nil {
return "", err
}
path := filepath.Join(home, ".speedrun/privatekey")
return path, nil
}
func newKey(cmd *cobra.Command, args []string) error {
log.Debug("Generating new private key")
_, privKey, err := ed25519.GenerateKey(rand.Reader)
if err != nil {
return err
}
log.Debug("Converting private key to PKCS8 format")
pemBlock := &pem.Block{}
pemBlock.Type = "PRIVATE KEY"
pemBlock.Bytes, err = x509.MarshalPKCS8PrivateKey(privKey)
if err != nil {
return err
}
log.Debug("Encoding the key to PEM format")
privateKey := pem.EncodeToMemory(pemBlock)
err = writeKeyFile(privateKey)
if err != nil {
return err
}
return nil
}
func writeKeyFile(key []byte) error {
privateKeyPath, err := determineKeyFilePath()
if err != nil {
return err
}
log.Debugf("Writing priviate key to %s", privateKeyPath)
err = ioutil.WriteFile(privateKeyPath, key, 0600)
if err != nil {
return err
}
return nil
}
func loadKeyPair() (ssh.PublicKey, ssh.Signer, error) {
privateKeyPath, err := determineKeyFilePath()
if err != nil {
return nil, nil, err
}
file, err := readKeyFile(privateKeyPath)
if err != nil {
return nil, nil, fmt.Errorf("couldn't find private key. Use 'speedrun key new' to generate a new one")
}
signer, err := ssh.ParsePrivateKey(file)
if err != nil {
return nil, nil, err
}
pubKey := signer.PublicKey()
return pubKey, signer, nil
}
func readKeyFile(path string) ([]byte, error) {
cleanPath := filepath.Clean(path)
absPath, err := filepath.Abs(cleanPath)
if err != nil {
return nil, err
}
file, err := ioutil.ReadFile(absPath)
if err != nil {
return nil, err
}
return file, nil
}
func setKey(cmd *cobra.Command, args []string) error {
client, err := gcp.NewComputeClient(viper.GetString("gcp.projectid"))
if err != nil {
return err
}
pubKey, _, err := loadKeyPair()
if err != nil {
return err
}
filter, err := cmd.Flags().GetString("filter")
if err != nil {
return err
}
if filter != "" {
log.Info("Setting public key in the instance metadata")
instances, err := client.GetInstances(filter)
if err != nil {
return err
}
if len(instances) == 0 {
log.Warn("no instances found")
}
pool := pond.New(10, 0, pond.MinWorkers(10))
for i := 0; i < len(instances); i++ {
n := i
pool.Submit(func() {
client.AddKeyToMetadata(instances[n], pubKey)
})
}
pool.StopAndWait()
} else {
log.Info("Setting public key in the project metadata")
err = client.AddKeyToMetadataP(pubKey)
if err != nil {
return err
}
}
return nil
}
func removeKey(cmd *cobra.Command, args []string) error {
client, err := gcp.NewComputeClient(viper.GetString("gcp.projectid"))
if err != nil {
return err
}
pubKey, _, err := loadKeyPair()
if err != nil {
return err
}
filter, err := cmd.Flags().GetString("filter")
if err != nil {
return err
}
if filter != "" {
log.Info("Removing public from the instance metadata")
instances, err := client.GetInstances(filter)
if err != nil {
return err
}
if len(instances) == 0 {
log.Warn("no instances found")
}
pool := pond.New(10, 0, pond.MinWorkers(10))
for i := 0; i < len(instances); i++ {
n := i
pool.Submit(func() {
client.RemoveKeyFromMetadata(instances[n], pubKey)
})
}
pool.StopAndWait()
} else {
log.Info("Removing public key from the project metadata")
err = client.RemoveKeyFromMetadataP(pubKey)
if err != nil {
return err
}
}
return nil
}

52
cmd/root.go Normal file
View File

@@ -0,0 +1,52 @@
package cmd
import (
"path"
"github.com/apex/log"
"github.com/spf13/cobra"
"github.com/spf13/viper"
)
var cfgFile string
//Execute runs the root command
func Execute() {
// cobra.OnInitialize(initConfig)
var rootCmd = &cobra.Command{
Use: "speedrun",
Short: "Cloud first command execution",
Version: "0.1.0",
}
rootCmd.AddCommand(initCmd)
rootCmd.AddCommand(keyCmd)
rootCmd.AddCommand(runCmd)
rootCmd.PersistentFlags().StringVarP(&cfgFile, "config", "c", "$HOME/.speedrun/config.toml", "config file")
rootCmd.PersistentFlags().StringP("loglevel", "l", "info", "Log level")
viper.BindPFlag("loglevel", rootCmd.PersistentFlags().Lookup("loglevel"))
if err := rootCmd.Execute(); err != nil {
log.Fatal(err.Error())
}
}
func initConfig() {
dir, file := path.Split(cfgFile)
viper.SetConfigName(file)
viper.SetConfigType("toml")
viper.AddConfigPath(dir)
if err := viper.ReadInConfig(); err != nil {
if _, ok := err.(viper.ConfigFileNotFoundError); ok {
if err != nil {
log.Error(err.Error())
log.Fatal("Run `speedrun init` first")
}
} else {
log.Fatal(err.Error())
}
}
log.SetLevelFromString(viper.GetString("loglevel"))
}

94
cmd/run.go Normal file
View File

@@ -0,0 +1,94 @@
package cmd
import (
"fmt"
gcp "speedrun/cloud"
"speedrun/colors"
"speedrun/marathon"
"strings"
"time"
"github.com/apex/log"
"github.com/spf13/cobra"
"github.com/spf13/viper"
)
var runCmd = &cobra.Command{
Use: "run <command to run>",
Short: "Run command on remote servers",
Example: " speedrun run whoami\n speedrun run whoami --only-failures --filter \"labels.foo = bar AND labels.environment = staging\"",
Args: cobra.MinimumNArgs(1),
PreRun: func(cmd *cobra.Command, args []string) {
initConfig()
},
RunE: run,
}
func init() {
runCmd.Flags().String("filter", "", "Fetch instances that match the filter")
runCmd.Flags().String("projectid", "", "Override GCP project id")
runCmd.Flags().Bool("only-failures", false, "Print only failures and errors")
runCmd.Flags().Bool("ignore-fingerprint", false, "Ignore host's fingerprint mismatch")
runCmd.Flags().Duration("timeout", time.Duration(10*time.Second), "SSH connection timeout")
runCmd.Flags().Int("concurrency", 100, "Number of maximum concurrent SSH workers")
runCmd.Flags().Bool("use-private-ip", false, "Connect to private IPs instead of public ones")
viper.BindPFlag("gcp.projectid", runCmd.Flags().Lookup("projectid"))
viper.BindPFlag("ssh.timeout", runCmd.Flags().Lookup("timeout"))
viper.BindPFlag("ssh.ignore-fingerprint", runCmd.Flags().Lookup("ignore-fingerprint"))
viper.BindPFlag("ssh.only-failures", runCmd.Flags().Lookup("only-failures"))
viper.BindPFlag("ssh.concurrency", runCmd.Flags().Lookup("concurrency"))
viper.BindPFlag("ssh.use-private-ip", runCmd.Flags().Lookup("use-private-ip"))
}
func run(cmd *cobra.Command, args []string) error {
command := strings.Join(args, " ")
projectid := viper.GetString("gcp.projectid")
timeout := viper.GetDuration("ssh.timeout")
ignoreFingerprint := viper.GetBool("ssh.ignore-fingerprint")
onlyFailures := viper.GetBool("ssh.only-failures")
concurrency := viper.GetInt("ssh.concurrency")
usePrivateIP := viper.GetBool("ssh.use-private-ip")
filter, err := cmd.Flags().GetString("filter")
if err != nil {
return err
}
client, err := gcp.NewComputeClient(projectid)
if err != nil {
log.Fatal(err.Error())
}
privateKeyPath, err := determineKeyFilePath()
if err != nil {
log.Fatal(err.Error())
}
log.Info("Fetching list of GCE instances")
instances, err := client.GetInstances(filter)
if err != nil {
log.Fatal(err.Error())
}
if len(instances) == 0 {
log.Warn("No instances found")
return nil
}
log.Info(fmt.Sprintf("Running [%s]", colors.Blue(command)))
m := marathon.New(command, timeout, concurrency)
instanceDict := map[string]string{}
for _, instance := range instances {
if usePrivateIP {
instanceDict[instance.NetworkInterfaces[0].NetworkIP] = instance.Name
} else {
instanceDict[instance.NetworkInterfaces[0].AccessConfigs[0].NatIP] = instance.Name
}
}
err = m.Run(instanceDict, privateKeyPath, ignoreFingerprint)
if err != nil {
log.Fatal(err.Error())
}
m.PrintResult(onlyFailures)
return nil
}