Merge pull request #11 from crytic/refactor-image-management

Refactor image management

.github/workflows/release.yml

@@ -53,4 +53,4 @@ jobs:
           version: latest
           args: release --clean --skip-validate
         env:
-          GITHUB_TOKEN: ${{ secrets.OLDSJ_GITHUB_TOKEN }}
+          GITHUB_TOKEN: ${{ secrets.RELEASE_TOKEN }}

.goreleaser.yml

@@ -45,7 +45,7 @@ signs:
 
 release:
   github:
-    owner: trailofbits
+    owner: crytic
     name: cloudexec
 
 brews:

README.md

@@ -111,12 +111,6 @@ Confirm `cloudexec` has access to DigitalOcean.
 cloudexec check
 ```
 
-### Initialize your CloudExec environment
-
-```bash
-cloudexec init
-```
-
 ### Launch a new remote job
 
 Generate a cloudexec.toml configuration file in the current directory.
@@ -195,3 +189,17 @@ cloudexec clean
 ```
 
 Note that there is often a delay while deleting files from Digital Ocean Spaces buckets.
+
+## Optional: Create a CloudExec DigitalOcean image
+
+Building and uploading a dedicated DigitalOcean image for `cloudexec` will simplify your launch configuration and improve startup times.
+
+To do so, install `packer` with `brew install packer`. If you're using `nix` and `direnv`, it's added to your PATH via the flake's dev shell.
+
+To build and upload a docker image, run the following command. Make sure your DigitalOcean API key is either in your env vars or replace it with the actual token.
+
+`packer build -var do_api_token=$DIGITALOCEAN_API_KEY cloudexec.pkr.hcl`
+
+This will take care of everything and if you visit the [DigitalOcean snapshots page](https://cloud.digitalocean.com/images/snapshots/droplets), you'll see a snapshot called `cloudexec-20230920164605` or similar. `cloudexec` will search for snapshots starts with a `cloudexec-` prefix and it will use the one with the most recent timestamp string.
+
+Now, you can remove everything from the setup command in the example launch config or replace it to install additional tools.

cmd/cloudexec/init.go

@@ -7,25 +7,74 @@ import (
 	"github.com/crytic/cloudexec/pkg/s3"
 )
 
-func Init(username string, config config.Config) error {
-	bucketName := fmt.Sprintf("cloudexec-%s-trailofbits", username)
-	// Create a new bucket (or get an existing one)
-	err := s3.GetOrCreateBucket(config, username)
+func Init(config config.Config, bucket string) error {
+	// Get a list of existing buckets
+	listBucketsOutput, err := s3.ListBuckets(config)
 	if err != nil {
-		return fmt.Errorf("Failed to get bucket for %s: %w", username, err)
+		return fmt.Errorf("Failed to list buckets: %w", err)
 	}
-	fmt.Printf("Using bucket: %v\n", bucketName)
 
-	// Create the state directory
-	err = s3.PutObject(config, bucketName, "state/", []byte{})
+	// Return if the desired bucket already exists
+	bucketExists := false
+	for _, thisBucket := range listBucketsOutput {
+		if thisBucket == bucket {
+			bucketExists = true
+		}
+	}
+
+	if !bucketExists {
+		// Create a new bucket
+		fmt.Printf("Creating new %s bucket...\n", bucket)
+		err = s3.CreateBucket(config, bucket)
+		if err != nil {
+			return err
+		}
+	}
+
+	// Ensure versioning is enabled, necessary if bucket creation was interrupted
+	err = s3.SetVersioning(config, bucket)
 	if err != nil {
-		return fmt.Errorf("Failed to create state directory in bucket %s: %w", bucketName, err)
+		return err
 	}
 
-	// Create the initial state file
-	err = s3.PutObject(config, bucketName, "state/state.json", []byte("{}"))
+	// Initialize bucket state if not already present
+	err = initState(config, bucket)
 	if err != nil {
-		return fmt.Errorf("Failed to create state file in bucket %s: %w", bucketName, err)
+		return fmt.Errorf("Failed to initialize state for bucket %s: %w", bucket, err)
+	}
+
+	return nil
+}
+
+func initState(config config.Config, bucket string) error {
+	// Check if the state directory already exists
+	stateDir := "state/"
+	stateDirExists, err := s3.ObjectExists(config, bucket, 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{})
+		if err != nil {
+			return fmt.Errorf("Failed to create state directory at %s/%s: %w", bucket, stateDir, err)
+		}
+	}
+
+	// Check if the state file already exists
+	statePath := "state/state.json"
+	statePathExists, err := s3.ObjectExists(config, bucket, 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("{}"))
+		if err != nil {
+			return fmt.Errorf("Failed to create state file in bucket %s: %w", bucket, err)
+		}
 	}
 
 	return nil

cmd/cloudexec/launch.go

@@ -11,7 +11,6 @@ import (
 	"github.com/BurntSushi/toml"
 	"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"
 )
@@ -87,13 +86,7 @@ func LoadLaunchConfig(launchConfigPath string) (LaunchConfig, error) {
 
 func Launch(user *user.User, config config.Config, dropletSize string, dropletRegion string, lc LaunchConfig) error {
 	username := user.Username
-	bucketName := fmt.Sprintf("cloudexec-%s-trailofbits", username)
-
-	fmt.Printf("Getting or creating new bucket for %s...\n", username)
-	err := s3.GetOrCreateBucket(config, username)
-	if err != nil {
-		return fmt.Errorf("Failed to get bucket for %s: %w", username, err)
-	}
+	bucketName := fmt.Sprintf("cloudexec-%s", username)
 
 	// get existing state from bucket
 	fmt.Printf("Getting existing state from bucket %s...\n", bucketName)

cmd/cloudexec/main.go

@@ -31,7 +31,8 @@ func main() {
 		os.Exit(1)
 	}
 	userName := user.Username
-	bucketName := fmt.Sprintf("cloudexec-%s-trailofbits", userName)
+	// TODO: sanitize username usage in bucketname
+	bucketName := fmt.Sprintf("cloudexec-%s", userName)
 
 	// Attempt to load the configuration
 	config, configErr := LoadConfig(configFilePath)
@@ -63,23 +64,6 @@ func main() {
 					return nil
 				},
 			},
-			{
-				Name:    "init",
-				Usage:   "Initialize the cloud environment",
-				Aliases: []string{"i"},
-				Action: func(*cli.Context) error {
-					// Abort on configuration error
-					if configErr != nil {
-						return configErr
-					}
-
-					err = Init(userName, config)
-					if err != nil {
-						return err
-					}
-					return nil
-				},
-			},
 			{
 				Name:  "configure",
 				Usage: "Configure credentials",
@@ -148,6 +132,12 @@ func main() {
 					dropletSize := c.String("size")
 					dropletRegion := c.String("region")
 
+					// Initialize the s3 state
+					err = Init(config, bucketName)
+					if err != nil {
+						return err
+					}
+
 					fmt.Printf("Launching a %s droplet in the %s region\n", dropletSize, dropletRegion)
 					err = Launch(user, config, dropletSize, dropletRegion, lc)
 					if err != nil {
@@ -172,6 +162,12 @@ func main() {
 						return configErr
 					}
 
+					// Initialize the s3 state
+					err = Init(config, bucketName)
+					if err != nil {
+						return err
+					}
+
 					existingState, err := state.GetState(config, bucketName)
 					if err != nil {
 						return err
@@ -206,6 +202,12 @@ func main() {
 						return configErr
 					}
 
+					// Initialize the s3 state
+					err = Init(config, bucketName)
+					if err != nil {
+						return err
+					}
+
 					instanceToJobs, err := state.GetJobIdsByInstance(config, bucketName)
 					if err != nil {
 						return err
@@ -244,6 +246,12 @@ func main() {
 						return configErr
 					}
 
+					// Initialize the s3 state
+					err = Init(config, bucketName)
+					if err != nil {
+						return err
+					}
+
 					instanceToJobs, err := state.GetJobIdsByInstance(config, bucketName)
 					if err != nil {
 						return err
@@ -287,6 +295,12 @@ func main() {
 						return configErr
 					}
 
+					// Initialize the s3 state
+					err = Init(config, bucketName)
+					if err != nil {
+						return err
+					}
+
 					existingState, err := state.GetState(config, bucketName)
 					if err != nil {
 						return err
@@ -329,6 +343,12 @@ func main() {
 						return configErr
 					}
 
+					// Initialize the s3 state
+					err = Init(config, bucketName)
+					if err != nil {
+						return err
+					}
+
 					existingState, err := state.GetState(config, bucketName)
 					if err != nil {
 						return err
@@ -392,6 +412,12 @@ func main() {
 								return configErr
 							}
 
+							// Initialize the s3 state
+							err = Init(config, bucketName)
+							if err != nil {
+								return err
+							}
+
 							// Retrieve existing state
 							existingState, err := state.GetState(config, bucketName)
 							if err != nil {
@@ -414,6 +440,12 @@ func main() {
 								return configErr
 							}
 
+							// Initialize the s3 state
+							err = Init(config, bucketName)
+							if err != nil {
+								return err
+							}
+
 							jobID := c.Args().First() // Get the job ID from the arguments
 							if jobID == "" {
 								fmt.Println("Please provide a job ID to remove")
@@ -449,6 +481,12 @@ func main() {
 								return configErr
 							}
 
+							// Initialize the s3 state
+							err = Init(config, bucketName)
+							if err != nil {
+								return err
+							}
+
 							// Retrieve existing state
 							existingState, err := state.GetState(config, bucketName)
 							if err != nil {
@@ -474,6 +512,13 @@ func main() {
 					if configErr != nil {
 						return configErr
 					}
+
+					// Initialize the s3 state
+					err = Init(config, bucketName)
+					if err != nil {
+						return err
+					}
+
 					// First check if there's a running job
 					existingState, err := state.GetState(config, bucketName)
 					if err != nil {

cmd/cloudexec/user_data.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	_ "embed"
 	"fmt"
+	"strings"
 	"text/template"
 	"time"
 
@@ -36,13 +37,14 @@ func GenerateUserData(config config.Config, lc LaunchConfig) (string, error) {
 	timeoutStr := fmt.Sprintf("%d", int(timeout.Seconds()))
 
 	// Set the values for the template
+	// double quotes are escaped so the command strings can be safely contained by double quotes in bash
 	data := UserData{
 		SpacesAccessKey:   config.DigitalOcean.SpacesAccessKey,
 		SpacesSecretKey:   config.DigitalOcean.SpacesSecretKey,
 		SpacesRegion:      config.DigitalOcean.SpacesRegion,
 		DigitalOceanToken: config.DigitalOcean.ApiKey,
-		SetupCommands:     lc.Commands.Setup,
-		RunCommand:        lc.Commands.Run,
+		SetupCommands:     strings.ReplaceAll(lc.Commands.Setup, `"`, `\"`),
+		RunCommand:        strings.ReplaceAll(lc.Commands.Run, `"`, `\"`),
 		Timeout:           timeoutStr,
 	}