Merge pull request #44 from crytic/qol-touch-up

Qol touch up

Author: bohendo

.gitignore

@@ -1,9 +1,15 @@
-*.backup
+# Environment
 .direnv/
 .env
+
+# Editors
 .vscode
 .vscode/
-bin/
+
+# Output
 dist/
-example/input.zip
 result
+cloudexec/
+example/input/crytic-export/
+example/input/output/
+example/input.zip

cmd/cloudexec/cancel.go

@@ -6,33 +6,35 @@ import (
 
 	"github.com/crytic/cloudexec/pkg/config"
 	do "github.com/crytic/cloudexec/pkg/digitalocean"
+	"github.com/crytic/cloudexec/pkg/log"
 	"github.com/crytic/cloudexec/pkg/state"
 )
 
 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)
+		log.Info("Job %v is not running, it is %s", job.ID, job.Status)
+    return nil
 	}
-	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)
+	log.Warn("Destroying droplet %s associated with job %v: IP=%v | CreatedAt=%s", 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)")
+		log.Warn("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)
+			log.Info("Droplet %s was not destroyed", job.Droplet.Name)
 			return nil
 		}
 	}
-	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)
+	log.Good("Droplet %v destroyed", job.Droplet.Name)
 	err = existingState.CancelRunningJob(config, job.ID)
 	if err != nil {
-		return fmt.Errorf("Failed to mark job as cancelled: %w", err)
+		return fmt.Errorf("Failed to change job status to cancelled: %w", err)
 	}
+	log.Good("Job %v status changed to cancelled", job.ID)
 	return nil
 }
 
@@ -42,17 +44,17 @@ func CancelAll(config config.Config, existingState *state.State, force bool) err
 		return fmt.Errorf("Failed to get all running servers: %w", err)
 	}
 	if len(droplets) == 0 {
-		fmt.Println("No running servers found")
+		log.Info("No running servers found")
 		return nil
 	}
-	fmt.Printf("Found %v running server(s):\n", len(droplets))
+	log.Info("Found %v running server(s):", 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(config, existingState, &job, force)
 		if err != nil {
-			fmt.Printf("Failed to cancel job %v", job.ID)
+			log.Error("Failed to cancel job %v", job.ID)
 		}
 	}
 	return nil

cmd/cloudexec/clean.go

@@ -5,12 +5,13 @@ import (
 	"strings"
 
 	"github.com/crytic/cloudexec/pkg/config"
+	"github.com/crytic/cloudexec/pkg/log"
 	"github.com/crytic/cloudexec/pkg/s3"
 	"github.com/crytic/cloudexec/pkg/ssh"
 	"github.com/crytic/cloudexec/pkg/state"
 )
 
-func CleanBucketJob(config config.Config, existingState *state.State, jobID int64, force bool) error {
+func CleanJob(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 {
@@ -19,56 +20,56 @@ func CleanBucketJob(config config.Config, existingState *state.State, jobID int6
 	// Confirm job data deletion
 	var numToRm int = len(objects)
 	if numToRm == 0 {
-		fmt.Printf("Bucket is already empty.\n")
+		log.Info("Bucket is already empty.")
 		return nil
 	}
-	fmt.Printf("Removing ALL input, output, and logs associated with %s...\n", prefix)
+	log.Warn("Removing all input, output, logs, and configuration associated with %s", prefix)
 	if !force { // Ask for confirmation before cleaning this job if no force flag
-		fmt.Println("Confirm? (y/n)")
+		log.Warn("Confirm? (y/n)")
 		var response string
 		fmt.Scanln(&response)
 		if strings.ToLower(response) != "y" {
-			fmt.Printf("Job %v was not cleaned\n", jobID)
+			log.Info("Job %v was not cleaned", jobID)
 			return nil
 		}
 	}
-	fmt.Printf("Deleting bucket contents...\n")
+	log.Wait("Deleting bucket contents")
 	// Delete all objects in the bucket
 	for _, object := range objects {
-		fmt.Println("Deleting object: ", object)
+		log.Info("Deleting object: %s", object)
 		err = s3.DeleteObject(config, object)
 		if err != nil {
 			return err
 		}
 	}
-	fmt.Printf("Deleted %d objects from bucket, removing job %v from state file..\n", numToRm, jobID)
+	log.Good("Bucket is clean")
 	newState := &state.State{}
 	deleteJob := state.Job{
 		ID:     jobID,
 		Delete: true,
 	}
 	newState.CreateJob(deleteJob)
 	err = state.MergeAndSave(config, newState)
+	log.Good("Removed job %v from state file", jobID)
 	if err != nil {
-		return fmt.Errorf("Error removing %s from state file: %w\n", prefix, err)
+		return fmt.Errorf("Error removing %s from state file: %w", 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
 }
 
-func CleanBucketAll(config config.Config, existingState *state.State, force bool) error {
+func CleanAll(config config.Config, existingState *state.State, force bool) error {
 	if len(existingState.Jobs) == 0 {
-		fmt.Println("No jobs are available")
+		log.Info("No jobs are available")
 		return nil
 	}
 	for _, job := range existingState.Jobs {
-		err := CleanBucketJob(config, existingState, job.ID, force)
+		err := CleanJob(config, existingState, job.ID, force)
 		if err != nil {
-			return err
+			log.Error("Failed to clean job %v", job.ID)
 		}
 	}
 	return nil

cmd/cloudexec/configure.go

@@ -9,15 +9,16 @@ import (
 	"os/user"
 	"strings"
 
-	"github.com/crytic/cloudexec/pkg/config"
 	"golang.org/x/term"
+
+	"github.com/crytic/cloudexec/pkg/config"
+	"github.com/crytic/cloudexec/pkg/log"
 )
 
 func Configure() error {
 	user, err := user.Current()
 	if err != nil {
-		fmt.Printf("Failed to get current user: %v", err)
-		os.Exit(1)
+		return fmt.Errorf("Failed to get current user: %v", err)
 	}
 	username, err := promptUserInput("Username", user.Username)
 	if err != nil {
@@ -63,7 +64,7 @@ func Configure() error {
 }
 
 func promptSecretInput(prompt, defaultValue string) (string, error) {
-	fmt.Printf("%s [%s]: ", prompt, defaultValue)
+	log.Info("%s [%s]: ", prompt, defaultValue)
 	rawInput, err := term.ReadPassword(int(os.Stdin.Fd()))
 	if err != nil {
 		return "", fmt.Errorf("Failed to read input: %w", err)
@@ -81,7 +82,7 @@ func promptSecretInput(prompt, defaultValue string) (string, error) {
 
 func promptUserInput(prompt, defaultValue string) (string, error) {
 	reader := bufio.NewReader(os.Stdin)
-	fmt.Printf("%s [%s]: ", prompt, defaultValue)
+	log.Info("%s [%s]: ", prompt, defaultValue)
 
 	input, err := reader.ReadString('\n')
 	if err != nil {

cmd/cloudexec/init.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 
 	"github.com/crytic/cloudexec/pkg/config"
+	"github.com/crytic/cloudexec/pkg/log"
 	"github.com/crytic/cloudexec/pkg/s3"
 )
 
@@ -25,7 +26,7 @@ func Init(config config.Config) error {
 
 	if !bucketExists {
 		// Create a new bucket
-		fmt.Printf("Creating new %s bucket...\n", bucketName)
+		log.Wait("Creating new %s bucket", bucketName)
 		err = s3.CreateBucket(config)
 		if err != nil {
 			return err
@@ -56,7 +57,7 @@ func initState(config config.Config, bucketName string) error {
 	}
 	// Create the state directory if it does not already exist
 	if !stateDirExists {
-		fmt.Printf("Creating new state directory at %s/%s\n", bucketName, stateDir)
+		log.Wait("Creating new state directory at %s/%s", bucketName, stateDir)
 		err = s3.PutObject(config, stateDir, []byte{})
 		if err != nil {
 			return fmt.Errorf("Failed to create state directory at %s/%s: %w", bucketName, stateDir, err)
@@ -71,7 +72,7 @@ func initState(config config.Config, bucketName string) error {
 	}
 	// Create the initial state file if it does not already exist
 	if !statePathExists {
-		fmt.Printf("Creating new state file at %s/%s\n", bucketName, statePath)
+		log.Wait("Creating new state file at %s/%s", bucketName, statePath)
 		err = s3.PutObject(config, statePath, []byte("{}"))
 		if err != nil {
 			return fmt.Errorf("Failed to create state file in bucket %s: %w", bucketName, err)

cmd/cloudexec/launch.go

@@ -9,6 +9,7 @@ import (
 	"github.com/BurntSushi/toml"
 	"github.com/crytic/cloudexec/pkg/config"
 	do "github.com/crytic/cloudexec/pkg/digitalocean"
+	"github.com/crytic/cloudexec/pkg/log"
 	"github.com/crytic/cloudexec/pkg/ssh"
 	"github.com/crytic/cloudexec/pkg/state"
 )
@@ -43,7 +44,7 @@ func InitLaunchConfig() error {
 
 	// Write the default launch config to the file
 	_, err = launchConfigFile.WriteString(`
-# Set the directory to upload to the droplet.
+# Set the directory to upload to the server.
 [input]
 directory = ""
 timeout = "48h"
@@ -83,11 +84,8 @@ func LoadLaunchConfig(launchConfigPath string) (LaunchConfig, error) {
 	return lc, nil
 }
 
-func Launch(config config.Config, dropletSize string, dropletRegion string, lc LaunchConfig) error {
-	username := config.Username
-	bucketName := fmt.Sprintf("cloudexec-%s", username)
-
-	// get existing state from bucket
+func Launch(config config.Config, serverSize string, serverRegion string, lc LaunchConfig) error {
+	// get existing state
 	existingState, err := state.GetState(config)
 	if err != nil {
 		return fmt.Errorf("Failed to get S3 state: %w", err)
@@ -100,90 +98,85 @@ func Launch(config config.Config, dropletSize string, dropletRegion string, lc L
 	} else {
 		latestJobId = latestJob.ID
 	}
-	thisJobId := latestJobId + 1
+	jobID := latestJobId + 1
 
 	// update state struct with a new job
 	newState := &state.State{}
 	startedAt := time.Now().Unix()
 
 	newJob := state.Job{
 		Name:      lc.Input.JobName,
-		ID:        thisJobId,
+		ID:        jobID,
 		Status:    state.Provisioning,
 		StartedAt: startedAt,
 	}
 	newState.CreateJob(newJob)
 	// sync state to bucket
-	fmt.Printf("Adding new job to the state...\n")
 	err = state.MergeAndSave(config, newState)
+	log.Info("Registered new job with id %v", jobID)
 	if err != nil {
 		return fmt.Errorf("Failed to update S3 state: %w", err)
 	}
 
 	// upload local files to the bucket
 	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)
+	destPath := fmt.Sprintf("job-%v", jobID)
 	err = UploadDirectoryToSpaces(config, sourcePath, destPath)
 	if err != nil {
 		return fmt.Errorf("Failed to upload files: %w", err)
 	}
 
 	// Get or create an SSH key
-	fmt.Println("Getting or creating SSH key pair...")
 	publicKey, err := ssh.GetOrCreateSSHKeyPair()
 	if err != nil {
 		return fmt.Errorf("Failed to get or creating SSH key pair: %w", err)
 	}
 
 	// Prepare user data
-	fmt.Println("Generating user data...")
 	userData, err := GenerateUserData(config, lc)
 	if err != nil {
 		return fmt.Errorf("Failed to generate user data: %w", err)
 	}
 
-	fmt.Printf("Creating new %s droplet in %s for job %d...\n", dropletSize, config.DigitalOcean.SpacesRegion, thisJobId)
-	droplet, err := do.CreateDroplet(config, config.DigitalOcean.SpacesRegion, dropletSize, userData, thisJobId, publicKey)
+	log.Wait("Creating new %s server in %s for job %d", serverSize, config.DigitalOcean.SpacesRegion, jobID)
+	server, err := do.CreateDroplet(config, config.DigitalOcean.SpacesRegion, serverSize, userData, jobID, publicKey)
 	if err != nil {
-		return fmt.Errorf("Failed to create droplet: %w", err)
+		return fmt.Errorf("Failed to create server: %w", err)
 	}
+	log.Good("Server created with IP: %v", server.IP)
 
-	fmt.Printf("Droplet created with IP: %v\n", droplet.IP)
-
-	// Add the droplet info to state
-	fmt.Println("Adding new droplet info to state...")
+	// Add the server info to state
 	updatedAt := time.Now().Unix()
 	for i, job := range newState.Jobs {
-		if job.ID == thisJobId {
-			newState.Jobs[i].Droplet = droplet
+		if job.ID == jobID {
+			newState.Jobs[i].Droplet = server
 			newState.Jobs[i].UpdatedAt = updatedAt
 		}
 	}
-	fmt.Printf("Uploading new state to %s\n", bucketName)
 	err = state.MergeAndSave(config, newState)
 	if err != nil {
 		return fmt.Errorf("Failed to update S3 state: %w", err)
 	}
+	log.Info("Saved new server info to state")
 
-	// Add the droplet to the SSH config file
-	fmt.Println("Adding droplet to SSH config file...")
-	err = ssh.AddSSHConfig(thisJobId, droplet.IP)
+	// Add the server to the SSH config file
+	err = ssh.AddSSHConfig(jobID, server.IP)
 	if err != nil {
-		return fmt.Errorf("Failed to add droplet to SSH config file: %w", err)
+		return fmt.Errorf("Failed to add server to SSH config file: %w", err)
 	}
+	log.Info("Added cloudexec-%v to SSH config", jobID)
 
-	// Ensure we can SSH into the droplet
-	fmt.Println("Ensuring we can SSH into the droplet...")
-	err = ssh.WaitForSSHConnection(thisJobId)
+	// Ensure we can SSH into the server
+	log.Wait("Waiting for our new server to wake up")
+	err = ssh.WaitForSSHConnection(jobID)
 	if err != nil {
-		return fmt.Errorf("Failed to SSH into the droplet: %w", err)
+		return fmt.Errorf("Failed to SSH into the server: %w", err)
 	}
-	fmt.Println("SSH connection established!")
-	fmt.Println("Launch complete")
-	fmt.Println("You can now attach to the running job with: cloudexec attach")
-	fmt.Println("Stream logs from the droplet with: cloudexec logs")
-	fmt.Println("SSH to your droplet with: ssh cloudexec")
+	log.Good("Good Morning!")
+	fmt.Println()
+	log.Info("Stream logs from the server with: cloudexec logs")
+	log.Info("SSH to your server with: ssh cloudexec-%v", jobID)
+	log.Info("Once setup is complete, you can attach to the running job with: cloudexec attach")
 
 	return nil
 }