main.go

// cshell is a small CLI for managing AWS CloudShell environments and opening an
// interactive session in your terminal. It signs the unofficial CloudShell API
// with SigV4 using your AWS named profile (SSO supported) and hands the session
// off to the AWS session-manager-plugin.
package main

import (
	"context"
	"flag"
	"fmt"
	"os"
	"text/tabwriter"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/config"
)

// envName is the name used for environments this tool creates with a VPC.
const envName = "cshell"

// version is overridden at build time via -ldflags "-X main.version=…".
var version = "dev"

func main() {
	if len(os.Args) < 2 {
		usage()
		os.Exit(2)
	}

	cmd := os.Args[1]
	if cmd == "-h" || cmd == "--help" || cmd == "help" {
		usage()
		return
	}
	if cmd == "version" || cmd == "-v" || cmd == "--version" {
		fmt.Println("cshell", version)
		return
	}

	fs := flag.NewFlagSet(cmd, flag.ExitOnError)
	profile := fs.String("profile", os.Getenv("AWS_PROFILE"), "AWS named profile to use")
	region := fs.String("region", "", "AWS region (defaults to the profile's region)")

	var id, vpcID, subnetID, sgID string
	var yes, inject bool
	switch cmd {
	case "connect", "status", "delete":
		fs.StringVar(&id, "id", "", "environment id (defaults to the existing environment)")
	}
	switch cmd {
	case "connect", "create":
		fs.StringVar(&vpcID, "vpc-id", "", "VPC id to attach (requires -subnet-id and -sg-id)")
		fs.StringVar(&subnetID, "subnet-id", "", "subnet id")
		fs.StringVar(&sgID, "sg-id", "", "security group id")
	case "vpcs":
		fs.StringVar(&vpcID, "vpc-id", "", "list subnets and security groups in this VPC")
	}
	if cmd == "connect" {
		fs.BoolVar(&inject, "inject", false, "inject your credentials into the shell as env vars (non-VPC only)")
	}
	if cmd == "delete" {
		fs.BoolVar(&yes, "yes", false, "skip the confirmation prompt")
	}
	_ = fs.Parse(os.Args[2:])

	ctx := context.Background()
	cfg, err := loadConfig(ctx, *profile, *region)
	if err != nil {
		fatal(err)
	}
	client := NewClient(cfg.Region, cfg.Credentials)

	switch cmd {
	case "list":
		err = cmdList(ctx, client)
	case "status":
		err = cmdStatus(ctx, client, id)
	case "create":
		err = cmdCreate(ctx, client, vpcConfigFrom(vpcID, subnetID, sgID))
	case "connect":
		err = cmdConnect(ctx, client, cfg.Region, id, vpcConfigFrom(vpcID, subnetID, sgID), inject)
	case "delete":
		err = cmdDelete(ctx, client, id, yes)
	case "vpcs":
		err = cmdVpcs(ctx, cfg, vpcID)
	default:
		fmt.Fprintf(os.Stderr, "unknown command %q\n\n", cmd)
		usage()
		os.Exit(2)
	}
	if err != nil {
		fatal(err)
	}
}

func loadConfig(ctx context.Context, profile, region string) (aws.Config, error) {
	opts := []func(*config.LoadOptions) error{}
	if profile != "" {
		opts = append(opts, config.WithSharedConfigProfile(profile))
	}
	if region != "" {
		opts = append(opts, config.WithRegion(region))
	}
	cfg, err := config.LoadDefaultConfig(ctx, opts...)
	if err != nil {
		return cfg, err
	}
	if cfg.Region == "" {
		return cfg, fmt.Errorf("no region configured; pass -region or set one in your profile")
	}
	// Cache credentials so we don't re-resolve (and re-prompt SSO) per call.
	cfg.Credentials = aws.NewCredentialsCache(cfg.Credentials)
	return cfg, nil
}

func vpcConfigFrom(vpc, subnet, sg string) *VpcConfig {
	if vpc != "" && subnet != "" && sg != "" {
		return &VpcConfig{VpcId: vpc, SubnetIds: []string{subnet}, SecurityGroupIds: []string{sg}}
	}
	return nil
}

// resolveTarget finds the environment to act on: by id if given, otherwise the
// single live environment. Returns nil (no error) when none exist so callers can
// decide whether to create one.
func resolveTarget(ctx context.Context, c *Client, id string) (*Environment, error) {
	envs, err := c.DescribeEnvironmentsWithStatus(ctx)
	if err != nil {
		return nil, err
	}

	if id != "" {
		for i := range envs {
			if envs[i].EnvironmentId == id {
				return &envs[i], nil
			}
		}
		return nil, fmt.Errorf("environment %q not found", id)
	}

	var live []Environment
	for _, e := range envs {
		if e.Status != "DELETING" && e.Status != "DELETED" {
			live = append(live, e)
		}
	}
	switch len(live) {
	case 0:
		return nil, nil
	case 1:
		return &live[0], nil
	default:
		ids := ""
		for _, e := range live {
			ids += "\n  " + e.EnvironmentId
		}
		return nil, fmt.Errorf("multiple environments found; specify -id (one of:%s)", ids)
	}
}

func cmdList(ctx context.Context, c *Client) error {
	envs, err := c.DescribeEnvironmentsWithStatus(ctx)
	if err != nil {
		return err
	}
	if len(envs) == 0 {
		fmt.Println("No CloudShell environments.")
		return nil
	}
	w := tabwriter.NewWriter(os.Stdout, 0, 2, 2, ' ', 0)
	fmt.Fprintln(w, "ENVIRONMENT ID\tSTATUS\tVPC")
	for _, e := range envs {
		vpc := "-"
		if e.VpcConfig != nil {
			vpc = e.VpcConfig.VpcId
		}
		fmt.Fprintf(w, "%s\t%s\t%s\n", e.EnvironmentId, e.Status, vpc)
	}
	return w.Flush()
}

func cmdStatus(ctx context.Context, c *Client, id string) error {
	target, err := resolveTarget(ctx, c, id)
	if err != nil {
		return err
	}
	if target == nil {
		fmt.Println("No CloudShell environment.")
		return nil
	}
	fmt.Printf("%s  %s\n", target.EnvironmentId, target.Status)
	if target.VpcConfig != nil {
		fmt.Printf("VPC: %s  subnets: %v  security groups: %v\n",
			target.VpcConfig.VpcId, target.VpcConfig.SubnetIds, target.VpcConfig.SecurityGroupIds)
	}
	return nil
}

func cmdCreate(ctx context.Context, c *Client, vpc *VpcConfig) error {
	e, err := c.CreateEnvironment(ctx, vpc)
	if err != nil {
		return err
	}
	fmt.Printf("Created %s (%s)\n", e.EnvironmentId, e.Status)
	return nil
}

func cmdConnect(ctx context.Context, c *Client, region, id string, vpc *VpcConfig, inject bool) error {
	target, err := resolveTarget(ctx, c, id)
	if err != nil {
		return err
	}
	if target == nil {
		fmt.Println("No environment found; creating one…")
		if target, err = c.CreateEnvironment(ctx, vpc); err != nil {
			return err
		}
	}

	if target.VpcConfig != nil {
		fmt.Fprintf(os.Stderr,
			"warning: environment is attached to VPC %s; AWS API calls will time out unless the subnet has internet egress (NAT/IGW) or VPC endpoints\n",
			target.VpcConfig.VpcId)
		if inject {
			fmt.Fprintln(os.Stderr, "warning: skipping credential injection for a VPC-attached environment")
			inject = false
		}
	}

	fmt.Printf("Waiting for %s to be ready…\n", target.EnvironmentId)
	if err := c.WaitForRunning(ctx, target.EnvironmentId, target.Status, 3*time.Minute); err != nil {
		return err
	}

	raw, err := c.CreateSession(ctx, target.EnvironmentId)
	if err != nil {
		return err
	}

	fmt.Printf("Connecting to %s…\n", target.EnvironmentId)
	if inject {
		cr, err := c.Credentials(ctx)
		if err != nil {
			return fmt.Errorf("resolve credentials for injection: %w", err)
		}
		return launchSessionInject(raw, region, buildInjectCommand(region, cr))
	}
	return launchSession(raw, region)
}

func cmdDelete(ctx context.Context, c *Client, id string, yes bool) error {
	target, err := resolveTarget(ctx, c, id)
	if err != nil {
		return err
	}
	if target == nil {
		return fmt.Errorf("no environment to delete")
	}

	if !yes {
		fmt.Printf("Permanently delete %s and its persistent storage? [y/N]: ", target.EnvironmentId)
		var ans string
		fmt.Scanln(&ans)
		if ans != "y" && ans != "Y" {
			fmt.Println("Aborted.")
			return nil
		}
	}

	if err := c.DeleteEnvironment(ctx, target.EnvironmentId); err != nil {
		return err
	}
	fmt.Printf("Deleted %s\n", target.EnvironmentId)
	return nil
}

func usage() {
	fmt.Fprint(os.Stderr, `cshell — manage AWS CloudShell environments from your terminal

Usage:
  cshell <command> [flags]

Commands:
  connect   Connect to a CloudShell environment (creating one if needed)
  list      List CloudShell environments and their status
  status    Show the status of an environment
  create    Create a CloudShell environment
  delete    Delete a CloudShell environment
  vpcs      List VPCs (or subnets + security groups with -vpc-id)
  version   Print the cshell version

Common flags:
  -profile <name>   AWS named profile (default: $AWS_PROFILE)
  -region <region>  AWS region (default: the profile's region)
  -id <id>          Target environment id (connect/status/delete)

VPC flags (connect/create):
  -vpc-id, -subnet-id, -sg-id   Attach the environment to a VPC

Connect flags:
  -inject   Inject your credentials into the shell as env vars (non-VPC only)

Examples:
  cshell connect
  cshell connect -profile dev -region eu-west-1
  cshell connect -inject
  cshell list
  cshell delete -id abcdefgh-... -yes
`)
}

func fatal(err error) {
	fmt.Fprintln(os.Stderr, "error:", err)
	os.Exit(1)
}