Merge pull request #42 from crytic/attach-to-jobid

Attach to specific job id

cmd/cloudexec/cancel.go

@@ -0,0 +1,43 @@
+package main
+
+import (
+	"fmt"
+	"strings"
+
+	"github.com/crytic/cloudexec/pkg/config"
+	do "github.com/crytic/cloudexec/pkg/digitalocean"
+	"github.com/crytic/cloudexec/pkg/ssh"
+	"github.com/crytic/cloudexec/pkg/state"
+)
+
+func CancelJob(job *state.Job, existingState *state.State, config config.Config) error {
+	if job.Status != state.Provisioning && job.Status != state.Running {
+		return fmt.Errorf("Job %v is not running, it is %s", job.ID, job.Status)
+	}
+
+	fmt.Printf("Droplet %s associated with job %v: IP=%v | CreatedAt=%s\n", job.Droplet.Name, job.ID, job.Droplet.IP, job.Droplet.Created)
+	fmt.Println("Destroy this droplet? (y/n)")
+	var response string
+	fmt.Scanln(&response)
+	if strings.ToLower(response) == "y" {
+		fmt.Printf("Destroying droplet %v...\n", job.Droplet.ID)
+		err := do.DeleteDroplet(config, job.Droplet.ID)
+		if err != nil {
+			return fmt.Errorf("Failed to destroy droplet: %w", err)
+		}
+		fmt.Printf("Removing ssh config for droplet %v...\n", job.Droplet.ID)
+		err = ssh.DeleteSSHConfig(job.ID)
+		if err != nil {
+			return fmt.Errorf("Failed to delete ssh config: %w", err)
+		}
+		fmt.Printf("Marking job %v as cancelled...\n", job.Droplet.ID)
+		err = existingState.CancelRunningJob(config, job.ID)
+		if err != nil {
+			return fmt.Errorf("Failed to mark job as cancelled: %w", err)
+		}
+		fmt.Println("Done")
+	} else {
+		fmt.Printf("Job %v was not cancelled\n", job.ID)
+	}
+	return nil
+}

cmd/cloudexec/clean.go

@@ -7,65 +7,44 @@ import (
 	"github.com/crytic/cloudexec/pkg/config"
 	do "github.com/crytic/cloudexec/pkg/digitalocean"
 	"github.com/crytic/cloudexec/pkg/s3"
+	"github.com/crytic/cloudexec/pkg/state"
 )
 
-func ConfirmDeleteDroplets(config config.Config, dropletName string, instanceToJobs map[int64][]int64) ([]int64, error) {
-	var confirmedToDelete []int64
-	instances, err := do.GetDropletsByName(config, dropletName)
+func ConfirmCancelAll(config config.Config, existingState *state.State) error {
+	droplets, err := do.GetAllDroplets(config)
 	if err != nil {
-		return confirmedToDelete, fmt.Errorf("Failed to get droplets by name: %w", err)
+		return fmt.Errorf("Failed to get all running servers: %w", err)
 	}
-	if instanceToJobs == nil {
-		return confirmedToDelete, fmt.Errorf("Given instanceToJobs argument must not be nil")
+	if len(droplets) == 0 {
+		fmt.Printf("Zero servers found\n")
+		return nil
 	}
-	if len(instances) > 0 {
-		fmt.Printf("Existing %s instance(s) found:\n", dropletName)
-		for _, instance := range instances {
-			// get a pretty string describing the jobs associated with this instance
-			jobs := instanceToJobs[int64(instance.ID)]
-			var prettyJobs string
-			if len(jobs) == 0 {
-				prettyJobs = "none"
-			} else {
-				jobStrings := make([]string, len(jobs))
-				for i, job := range jobs {
-					jobStrings[i] = fmt.Sprint(job)
-				}
-				prettyJobs = strings.Join(jobStrings, ", ")
-			}
-
-			fmt.Printf("  - %v (IP: %v) (Jobs: %s) created at %v\n", instance.Name, instance.IP, prettyJobs, instance.Created)
-			fmt.Println("destroy this droplet? (y/n)")
-			var response string
-			fmt.Scanln(&response)
-			if strings.ToLower(response) == "y" {
-				fmt.Printf("Destroying droplet %v...\n", instance.ID)
-				err = do.DeleteDroplet(config, instance.ID)
-				if err != nil {
-					return confirmedToDelete, fmt.Errorf("Failed to destroy droplet: %w", err)
-				}
-				confirmedToDelete = append(confirmedToDelete, instance.ID)
-			}
+	fmt.Printf("Found %v running server(s):\n", len(droplets))
+	for _, job := range existingState.Jobs {
+		if job.Status != state.Provisioning && job.Status != state.Running {
+			continue // skip jobs that aren't running
+		}
+		err = CancelJob(&job, existingState, config)
+		if err != nil {
+			fmt.Printf("Failed to cancel job %v", job.ID)
 		}
-	} else {
-		fmt.Printf("Zero %s instances found\n", dropletName)
 	}
-	return confirmedToDelete, nil
+	return nil
 }
 
-func ResetBucket(config config.Config, bucketName string, spacesAccessKey string, spacesSecretKey string, spacesRegion string) error {
-	objects, err := s3.ListObjects(config, bucketName, "")
+func ResetBucket(config config.Config) error {
+	objects, err := s3.ListObjects(config, "")
 	if err != nil {
-		return fmt.Errorf("Failed to list objects in bucket '%s': %w", bucketName, err)
+		return fmt.Errorf("Failed to list objects in bucket: %w", err)
 	}
 
 	// Confirm bucket deletion
 	var numToRm int = len(objects)
 	if numToRm == 0 {
-		fmt.Printf("Bucket '%s' is already empty.\n", bucketName)
+		fmt.Printf("Bucket is already empty.\n")
 		return nil
 	} else {
-		fmt.Printf("Removing the first %d items from bucket %s...\n", numToRm, bucketName)
+		fmt.Printf("Removing the first %d items from bucket...\n", numToRm)
 		fmt.Println("Confirm? (y/n)")
 		var response string
 		fmt.Scanln(&response)
@@ -74,12 +53,12 @@ func ResetBucket(config config.Config, bucketName string, spacesAccessKey string
 			// Delete all objects in the bucket
 			for _, object := range objects {
 				fmt.Println("Deleting object: ", object)
-				err = s3.DeleteObject(config, bucketName, object)
+				err = s3.DeleteObject(config, object)
 				if err != nil {
 					return err
 				}
 			}
-			fmt.Printf("Deleted %d objects in bucket '%s'...\n", numToRm, bucketName)
+			fmt.Printf("Deleted %d objects from bucket...\n", numToRm)
 		}
 	}
 	return nil

cmd/cloudexec/init.go

@@ -7,7 +7,8 @@ import (
 	"github.com/crytic/cloudexec/pkg/s3"
 )
 
-func Init(config config.Config, bucket string) error {
+func Init(config config.Config) error {
+	bucketName := fmt.Sprintf("cloudexec-%s", config.Username)
 	// Get a list of existing buckets
 	listBucketsOutput, err := s3.ListBuckets(config)
 	if err != nil {
@@ -17,63 +18,63 @@ func Init(config config.Config, bucket string) error {
 	// Return if the desired bucket already exists
 	bucketExists := false
 	for _, thisBucket := range listBucketsOutput {
-		if thisBucket == bucket {
+		if thisBucket == bucketName {
 			bucketExists = true
 		}
 	}
 
 	if !bucketExists {
 		// Create a new bucket
-		fmt.Printf("Creating new %s bucket...\n", bucket)
-		err = s3.CreateBucket(config, bucket)
+		fmt.Printf("Creating new %s bucket...\n", bucketName)
+		err = s3.CreateBucket(config)
 		if err != nil {
 			return err
 		}
 	}
 
 	// Ensure versioning is enabled, necessary if bucket creation was interrupted
-	err = s3.SetVersioning(config, bucket)
+	err = s3.SetVersioning(config)
 	if err != nil {
 		return err
 	}
 
 	// Initialize bucket state if not already present
-	err = initState(config, bucket)
+	err = initState(config, bucketName)
 	if err != nil {
-		return fmt.Errorf("Failed to initialize state for bucket %s: %w", bucket, err)
+		return fmt.Errorf("Failed to initialize state for bucket %s: %w", bucketName, err)
 	}
 
 	return nil
 }
 
-func initState(config config.Config, bucket string) error {
+func initState(config config.Config, bucketName string) error {
 	// Check if the state directory already exists
 	stateDir := "state/"
-	stateDirExists, err := s3.ObjectExists(config, bucket, stateDir)
+	stateDirExists, err := s3.ObjectExists(config, stateDir)
 	if err != nil {
 		return fmt.Errorf("Failed to check whether the state directory exists: %w", err)
 	}
 	// Create the state directory if it does not already exist
 	if !stateDirExists {
-		fmt.Printf("Creating new state directory at %s/%s\n", bucket, stateDir)
-		err = s3.PutObject(config, bucket, stateDir, []byte{})
+		fmt.Printf("Creating new state directory at %s/%s\n", bucketName, stateDir)
+		err = s3.PutObject(config, stateDir, []byte{})
 		if err != nil {
-			return fmt.Errorf("Failed to create state directory at %s/%s: %w", bucket, stateDir, err)
+			return fmt.Errorf("Failed to create state directory at %s/%s: %w", bucketName, stateDir, err)
 		}
 	}
 
 	// Check if the state file already exists
 	statePath := "state/state.json"
-	statePathExists, err := s3.ObjectExists(config, bucket, statePath)
+	statePathExists, err := s3.ObjectExists(config, statePath)
 	if err != nil {
 		return fmt.Errorf("Failed to check whether the state file exists: %w", err)
 	}
 	// Create the initial state file if it does not already exist
 	if !statePathExists {
-		fmt.Printf("Creating new state file at %s/%s\n", bucket, statePath)
-		err = s3.PutObject(config, bucket, statePath, []byte("{}"))
+		fmt.Printf("Creating new state file at %s/%s\n", bucketName, statePath)
+		err = s3.PutObject(config, statePath, []byte("{}"))
 		if err != nil {
-			return fmt.Errorf("Failed to create state file in bucket %s: %w", bucket, err)
+			return fmt.Errorf("Failed to create state file in bucket %s: %w", bucketName, err)
 		}
 	}
 

cmd/cloudexec/launch.go

@@ -88,8 +88,7 @@ func Launch(config config.Config, dropletSize string, dropletRegion string, lc L
 	bucketName := fmt.Sprintf("cloudexec-%s", username)
 
 	// get existing state from bucket
-	fmt.Printf("Getting existing state from bucket %s...\n", bucketName)
-	existingState, err := state.GetState(config, bucketName)
+	existingState, err := state.GetState(config)
 	if err != nil {
 		return fmt.Errorf("Failed to get S3 state: %w", err)
 	}
@@ -107,16 +106,16 @@ func Launch(config config.Config, dropletSize string, dropletRegion string, lc L
 	newState := &state.State{}
 	startedAt := time.Now().Unix()
 
-	newJob := state.JobInfo{
+	newJob := state.Job{
 		Name:      lc.Input.JobName,
 		ID:        thisJobId,
 		Status:    state.Provisioning,
 		StartedAt: startedAt,
 	}
 	newState.CreateJob(newJob)
 	// sync state to bucket
-	fmt.Printf("Updating state in bucket %s...\n", bucketName)
-	err = state.UpdateState(config, bucketName, newState)
+	fmt.Printf("Adding new job to the state...\n")
+	err = state.UpdateState(config, newState)
 	if err != nil {
 		return fmt.Errorf("Failed to update S3 state: %w", err)
 	}
@@ -125,7 +124,7 @@ func Launch(config config.Config, dropletSize string, dropletRegion string, lc L
 	sourcePath := lc.Input.Directory // TODO: verify that this path exists & throw informative error if not
 	destPath := fmt.Sprintf("job-%v", thisJobId)
 	fmt.Printf("Compressing and uploading contents of directory %s to bucket %s/%s...\n", sourcePath, bucketName, destPath)
-	err = UploadDirectoryToSpaces(config, bucketName, sourcePath, destPath)
+	err = UploadDirectoryToSpaces(config, sourcePath, destPath)
 	if err != nil {
 		return fmt.Errorf("Failed to upload files: %w", err)
 	}
@@ -145,7 +144,7 @@ func Launch(config config.Config, dropletSize string, dropletRegion string, lc L
 	}
 
 	fmt.Printf("Creating new %s droplet in %s for job %d...\n", dropletSize, config.DigitalOcean.SpacesRegion, thisJobId)
-	droplet, err := do.CreateDroplet(config, username, config.DigitalOcean.SpacesRegion, dropletSize, userData, thisJobId, publicKey)
+	droplet, err := do.CreateDroplet(config, config.DigitalOcean.SpacesRegion, dropletSize, userData, thisJobId, publicKey)
 	if err != nil {
 		return fmt.Errorf("Failed to create droplet: %w", err)
 	}
@@ -162,26 +161,21 @@ func Launch(config config.Config, dropletSize string, dropletRegion string, lc L
 		}
 	}
 	fmt.Printf("Uploading new state to %s\n", bucketName)
-	err = state.UpdateState(config, bucketName, newState)
+	err = state.UpdateState(config, newState)
 	if err != nil {
 		return fmt.Errorf("Failed to update S3 state: %w", err)
 	}
 
 	// Add the droplet to the SSH config file
-	fmt.Println("Deleting old cloudexec instance from SSH config file...")
-	err = ssh.DeleteSSHConfig("cloudexec")
-	if err != nil {
-		return fmt.Errorf("Failed to delete old cloudexec entry from SSH config file: %w", err)
-	}
 	fmt.Println("Adding droplet to SSH config file...")
-	err = ssh.AddSSHConfig(droplet.IP)
+	err = ssh.AddSSHConfig(thisJobId, droplet.IP)
 	if err != nil {
 		return fmt.Errorf("Failed to add droplet to SSH config file: %w", err)
 	}
 
 	// Ensure we can SSH into the droplet
 	fmt.Println("Ensuring we can SSH into the droplet...")
-	err = ssh.WaitForSSHConnection()
+	err = ssh.WaitForSSHConnection(thisJobId)
 	if err != nil {
 		return fmt.Errorf("Failed to SSH into the droplet: %w", err)
 	}

cmd/cloudexec/logs.go

@@ -10,10 +10,10 @@ import (
 	"github.com/crytic/cloudexec/pkg/s3"
 )
 
-func GetLogsFromBucket(config config.Config, jobID int, bucketName string) error {
+func GetLogsFromBucket(config config.Config, jobID int64) error {
 	itemKey := fmt.Sprintf("job-%d/logs/cloud-init-output.log", jobID)
 
-	log, err := s3.GetObject(config, bucketName, itemKey)
+	log, err := s3.GetObject(config, itemKey)
 	if err != nil {
 		if err.Error() == "The specified key does not exist." {
 			return fmt.Errorf("The specified job logs do not exist. Please check the job ID and try again.\n")