Merge pull request #43 from crytic/pull-and-clean

Pull and clean
crytic · Mar 15, 2024 · 1f11b45 · 1f11b45
2 parents 20251d4 + a0783b9
commit 1f11b45
Show file tree

Hide file tree

Showing 11 changed files with 366 additions and 241 deletions.
diff --git a/cmd/cloudexec/cancel.go b/cmd/cloudexec/cancel.go
@@ -6,38 +6,54 @@ import (
 
 	"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 {
+func CancelJob(config config.Config, existingState *state.State, job *state.Job, force bool) 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("Destroying droplet %s associated with job %v: IP=%v | CreatedAt=%s\n", job.Droplet.Name, job.ID, job.Droplet.IP, job.Droplet.Created)
+	if !force { // Ask for confirmation before cleaning this job if no force flag
+		fmt.Println("Confirm? (y/n)")
+		var response string
+		fmt.Scanln(&response)
+		if strings.ToLower(response) != "y" {
+			fmt.Printf("Droplet %s was not destroyed\n", job.Droplet.Name)
+			return nil
 		}
-		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("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("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)
+	}
+	return nil
+}
+
+func CancelAll(config config.Config, existingState *state.State, force bool) error {
+	droplets, err := do.GetAllDroplets(config)
+	if err != nil {
+		return fmt.Errorf("Failed to get all running servers: %w", err)
+	}
+	if len(droplets) == 0 {
+		fmt.Println("No running servers found")
+		return nil
+	}
+	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
 		}
-		fmt.Printf("Marking job %v as cancelled...\n", job.Droplet.ID)
-		err = existingState.CancelRunningJob(config, job.ID)
+		err = CancelJob(config, existingState, &job, force)
 		if err != nil {
-			return fmt.Errorf("Failed to mark job as cancelled: %w", err)
+			fmt.Printf("Failed to cancel job %v", job.ID)
 		}
-		fmt.Println("Done")
-	} else {
-		fmt.Printf("Job %v was not cancelled\n", job.ID)
 	}
 	return nil
 }
diff --git a/cmd/cloudexec/clean.go b/cmd/cloudexec/clean.go
@@ -5,60 +5,70 @@ import (
 	"strings"
 
 	"github.com/crytic/cloudexec/pkg/config"
-	do "github.com/crytic/cloudexec/pkg/digitalocean"
 	"github.com/crytic/cloudexec/pkg/s3"
+	"github.com/crytic/cloudexec/pkg/ssh"
 	"github.com/crytic/cloudexec/pkg/state"
 )
 
-func ConfirmCancelAll(config config.Config, existingState *state.State) error {
-	droplets, err := do.GetAllDroplets(config)
+func CleanBucketJob(config config.Config, existingState *state.State, jobID int64, force bool) error {
+	prefix := fmt.Sprintf("job-%v", jobID)
+	objects, err := s3.ListObjects(config, prefix)
 	if err != nil {
-		return fmt.Errorf("Failed to get all running servers: %w", err)
+		return fmt.Errorf("Failed to list objects in bucket with prefix %s: %w", prefix, err)
 	}
-	if len(droplets) == 0 {
-		fmt.Printf("Zero servers found\n")
+	// Confirm job data deletion
+	var numToRm int = len(objects)
+	if numToRm == 0 {
+		fmt.Printf("Bucket is already empty.\n")
 		return nil
 	}
-	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
+	fmt.Printf("Removing ALL input, output, and logs associated with %s...\n", prefix)
+	if !force { // Ask for confirmation before cleaning this job if no force flag
+		fmt.Println("Confirm? (y/n)")
+		var response string
+		fmt.Scanln(&response)
+		if strings.ToLower(response) != "y" {
+			fmt.Printf("Job %v was not cleaned\n", jobID)
+			return nil
 		}
-		err = CancelJob(&job, existingState, config)
+	}
+	fmt.Printf("Deleting bucket contents...\n")
+	// Delete all objects in the bucket
+	for _, object := range objects {
+		fmt.Println("Deleting object: ", object)
+		err = s3.DeleteObject(config, object)
 		if err != nil {
-			fmt.Printf("Failed to cancel job %v", job.ID)
+			return err
 		}
 	}
-	return nil
-}
-
-func ResetBucket(config config.Config) error {
-	objects, err := s3.ListObjects(config, "")
+	fmt.Printf("Deleted %d objects from bucket, removing job %v from state file..\n", numToRm, jobID)
+	newState := &state.State{}
+	deleteJob := state.Job{
+		ID:     jobID,
+		Delete: true,
+	}
+	newState.CreateJob(deleteJob)
+	err = state.MergeAndSave(config, newState)
 	if err != nil {
-		return fmt.Errorf("Failed to list objects in bucket: %w", err)
+		return fmt.Errorf("Error removing %s from state file: %w\n", prefix, err)
 	}
+	fmt.Printf("Removing ssh config for job %v...\n", jobID)
+	err = ssh.DeleteSSHConfig(jobID)
+	if err != nil {
+		return fmt.Errorf("Failed to delete ssh config: %w", err)
+	}
+	return nil
+}
 
-	// Confirm bucket deletion
-	var numToRm int = len(objects)
-	if numToRm == 0 {
-		fmt.Printf("Bucket is already empty.\n")
+func CleanBucketAll(config config.Config, existingState *state.State, force bool) error {
+	if len(existingState.Jobs) == 0 {
+		fmt.Println("No jobs are available")
 		return nil
-	} else {
-		fmt.Printf("Removing the first %d items from bucket...\n", numToRm)
-		fmt.Println("Confirm? (y/n)")
-		var response string
-		fmt.Scanln(&response)
-		if strings.ToLower(response) == "y" {
-			fmt.Printf("Deleting bucket contents...\n")
-			// Delete all objects in the bucket
-			for _, object := range objects {
-				fmt.Println("Deleting object: ", object)
-				err = s3.DeleteObject(config, object)
-				if err != nil {
-					return err
-				}
-			}
-			fmt.Printf("Deleted %d objects from bucket...\n", numToRm)
+	}
+	for _, job := range existingState.Jobs {
+		err := CleanBucketJob(config, existingState, job.ID, force)
+		if err != nil {
+			return err
 		}
 	}
 	return nil

diff --git a/cmd/cloudexec/launch.go b/cmd/cloudexec/launch.go
@@ -115,7 +115,7 @@ func Launch(config config.Config, dropletSize string, dropletRegion string, lc L
 	newState.CreateJob(newJob)
 	// sync state to bucket
 	fmt.Printf("Adding new job to the state...\n")
-	err = state.UpdateState(config, newState)
+	err = state.MergeAndSave(config, newState)
 	if err != nil {
 		return fmt.Errorf("Failed to update S3 state: %w", err)
 	}
@@ -161,7 +161,7 @@ func Launch(config config.Config, dropletSize string, dropletRegion string, lc L
 		}
 	}
 	fmt.Printf("Uploading new state to %s\n", bucketName)
-	err = state.UpdateState(config, newState)
+	err = state.MergeAndSave(config, newState)
 	if err != nil {
 		return fmt.Errorf("Failed to update S3 state: %w", err)
 	}