Os login (#50)

Add OS Login support + change key storage format

cmd/run.go

@@ -1,7 +1,8 @@
 package cmd
 
 import (
-	gcp "speedrun/cloud"
+	"speedrun/cloud"
+	"speedrun/key"
 	"speedrun/marathon"
 	"strings"
 	"time"
@@ -30,7 +31,9 @@ func init() {
 	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")
+	runCmd.Flags().Bool("use-oslogin", false, "Authenticate via OS Login")
 	viper.BindPFlag("gcp.projectid", runCmd.Flags().Lookup("projectid"))
+	viper.BindPFlag("gcp.use-oslogin", runCmd.Flags().Lookup("use-oslogin"))
 	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"))
@@ -41,48 +44,55 @@ func init() {
 
 func run(cmd *cobra.Command, args []string) error {
 	command := strings.Join(args, " ")
-	projectid := viper.GetString("gcp.projectid")
+	project := 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")
+	useOSlogin := viper.GetBool("gcp.use-oslogin")
+
 	filter, err := cmd.Flags().GetString("filter")
 	if err != nil {
 		return err
 	}
 
-	client, err := gcp.NewComputeClient(projectid)
+	gcpClient, err := cloud.NewGCPClient(project)
 	if err != nil {
-		log.Fatal(err.Error())
+		return err
 	}
 
-	privateKeyPath, err := determineKeyFilePath()
+	path, err := determineKeyFilePath()
 	if err != nil {
-		log.Fatal(err.Error())
+		return err
+	}
+
+	k, err := key.Read(path)
+	if err != nil {
+		return err
 	}
 
 	log.Info("Fetching list of GCE instances")
-	instances, err := client.GetInstances(filter)
+	instances, err := gcpClient.GetInstances(filter, usePrivateIP)
 	if err != nil {
-		log.Fatal(err.Error())
+		return err
 	}
+
 	if len(instances) == 0 {
 		log.Warn("No instances found")
 		return nil
 	}
 
-	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
+	if useOSlogin {
+		user, err := gcpClient.GetSAUsername()
+		if err != nil {
+			return err
 		}
+		k.User = user
 	}
 
-	err = m.Run(instanceDict, privateKeyPath, ignoreFingerprint)
+	m := marathon.New(command, timeout, concurrency)
+	err = m.Run(instances, k, ignoreFingerprint)
 	if err != nil {
 		log.Fatal(err.Error())
 	}