Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Qol touch up #44

Merged
merged 16 commits into from
Mar 15, 2024
12 changes: 9 additions & 3 deletions .gitignore
Original file line number Diff line number Diff line change
@@ -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
22 changes: 12 additions & 10 deletions cmd/cloudexec/cancel.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
bohendo marked this conversation as resolved.
Show resolved Hide resolved
return nil
}

Expand All @@ -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)
bohendo marked this conversation as resolved.
Show resolved Hide resolved
}
}
return nil
Expand Down
29 changes: 15 additions & 14 deletions cmd/cloudexec/clean.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
bohendo marked this conversation as resolved.
Show resolved Hide resolved
Expand All @@ -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
}
}
bohendo marked this conversation as resolved.
Show resolved Hide resolved
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)
}
bohendo marked this conversation as resolved.
Show resolved Hide resolved
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
Expand Down
11 changes: 6 additions & 5 deletions cmd/cloudexec/configure.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
bohendo marked this conversation as resolved.
Show resolved Hide resolved
}
username, err := promptUserInput("Username", user.Username)
if err != nil {
Expand Down Expand Up @@ -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)
Expand All @@ -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 {
Expand Down
7 changes: 4 additions & 3 deletions cmd/cloudexec/init.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import (
"fmt"

"github.com/crytic/cloudexec/pkg/config"
"github.com/crytic/cloudexec/pkg/log"
"github.com/crytic/cloudexec/pkg/s3"
)

Expand All @@ -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
Expand Down Expand Up @@ -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)
Expand All @@ -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)
Expand Down
65 changes: 29 additions & 36 deletions cmd/cloudexec/launch.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
)
Expand Down Expand Up @@ -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.
bohendo marked this conversation as resolved.
Show resolved Hide resolved
[input]
directory = ""
timeout = "48h"
Expand Down Expand Up @@ -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
bohendo marked this conversation as resolved.
Show resolved Hide resolved
existingState, err := state.GetState(config)
if err != nil {
return fmt.Errorf("Failed to get S3 state: %w", err)
Expand All @@ -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")
bohendo marked this conversation as resolved.
Show resolved Hide resolved

return nil
}
Loading
Loading