cloudshell.go

package main

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	v4 "github.com/aws/aws-sdk-go-v2/aws/signer/v4"
)

// service is the SigV4 service name for the (unofficial) CloudShell API.
const service = "cloudshell"

// VpcConfig optionally attaches the environment to a VPC.
type VpcConfig struct {
	VpcId            string   `json:"VpcId"`
	SecurityGroupIds []string `json:"SecurityGroupIds"`
	SubnetIds        []string `json:"SubnetIds"`
}

// Environment is a CloudShell environment. Note: describeEnvironments returns
// only EnvironmentId; Status/VpcConfig come from getEnvironmentStatus.
type Environment struct {
	EnvironmentId   string     `json:"EnvironmentId"`
	Status          string     `json:"Status,omitempty"`
	EnvironmentName string     `json:"EnvironmentName,omitempty"`
	StatusReason    string     `json:"StatusReason,omitempty"`
	VpcConfig       *VpcConfig `json:"VpcConfig,omitempty"`
}

// Client talks to the CloudShell JSON API, signing each request with SigV4.
type Client struct {
	region       string
	creds        aws.CredentialsProvider
	signer       *v4.Signer
	http         *http.Client
	pollInterval time.Duration
}

// Credentials resolves the current AWS credentials (e.g. for shell injection).
func (c *Client) Credentials(ctx context.Context) (aws.Credentials, error) {
	return c.creds.Retrieve(ctx)
}

// NewClient builds a Client for the given region and credential provider.
func NewClient(region string, creds aws.CredentialsProvider) *Client {
	return &Client{
		region:       region,
		creds:        creds,
		signer:       v4.NewSigner(),
		http:         http.DefaultClient,
		pollInterval: 2 * time.Second,
	}
}

// call POSTs a SigV4-signed JSON request to the given API action and decodes the
// response into out (which may be nil, or a *json.RawMessage for the raw body).
func (c *Client) call(ctx context.Context, action string, body, out any) error {
	payload, err := json.Marshal(body)
	if err != nil {
		return err
	}

	url := fmt.Sprintf("https://%s.%s.amazonaws.com/%s", service, c.region, action)
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(payload))
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/json")

	creds, err := c.creds.Retrieve(ctx)
	if err != nil {
		return fmt.Errorf("resolve credentials: %w", err)
	}
	sum := sha256.Sum256(payload)
	if err := c.signer.SignHTTP(ctx, creds, req, hex.EncodeToString(sum[:]), service, c.region, time.Now()); err != nil {
		return fmt.Errorf("sign request: %w", err)
	}

	resp, err := c.http.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	data, _ := io.ReadAll(resp.Body)
	if resp.StatusCode/100 != 2 {
		return fmt.Errorf("%s failed (%d): %s", action, resp.StatusCode, string(data))
	}
	if out != nil && len(data) > 0 {
		return json.Unmarshal(data, out)
	}
	return nil
}

// DescribeEnvironments lists environments. The response shape varies, so it is
// coerced defensively (wrapped object, bare array, or single object).
func (c *Client) DescribeEnvironments(ctx context.Context) ([]Environment, error) {
	var raw json.RawMessage
	if err := c.call(ctx, "describeEnvironments", map[string]any{}, &raw); err != nil {
		return nil, err
	}

	var wrapped struct {
		Environments []Environment `json:"Environments"`
	}
	if json.Unmarshal(raw, &wrapped) == nil && wrapped.Environments != nil {
		return wrapped.Environments, nil
	}
	var arr []Environment
	if json.Unmarshal(raw, &arr) == nil {
		return arr, nil
	}
	var single Environment
	if json.Unmarshal(raw, &single) == nil && single.EnvironmentId != "" {
		return []Environment{single}, nil
	}
	return nil, nil
}

// DescribeEnvironmentsWithStatus enriches each environment with its real status
// (describeEnvironments returns only the id). Environments that can't be resolved
// — e.g. lingering after deletion — are marked DELETED.
func (c *Client) DescribeEnvironmentsWithStatus(ctx context.Context) ([]Environment, error) {
	envs, err := c.DescribeEnvironments(ctx)
	if err != nil {
		return nil, err
	}
	for i := range envs {
		s, err := c.GetEnvironmentStatus(ctx, envs[i].EnvironmentId)
		if err != nil {
			envs[i].Status = "DELETED"
			continue
		}
		envs[i].Status = s.Status
		envs[i].VpcConfig = s.VpcConfig
		if s.EnvironmentName != "" {
			envs[i].EnvironmentName = s.EnvironmentName
		}
	}
	return envs, nil
}

// GetEnvironmentStatus returns the current status (and VpcConfig) of an environment.
func (c *Client) GetEnvironmentStatus(ctx context.Context, id string) (*Environment, error) {
	var e Environment
	err := c.call(ctx, "getEnvironmentStatus", map[string]string{"EnvironmentId": id}, &e)
	return &e, err
}

// CreateEnvironment creates an environment, optionally attached to a VPC. When an
// environment already exists, CloudShell returns the existing one (idempotent).
func (c *Client) CreateEnvironment(ctx context.Context, vpc *VpcConfig) (*Environment, error) {
	body := map[string]any{}
	if vpc != nil {
		body["EnvironmentName"] = envName
		body["VpcConfig"] = vpc
	}
	var e Environment
	err := c.call(ctx, "createEnvironment", body, &e)
	return &e, err
}

// StartEnvironment resumes a suspended environment.
func (c *Client) StartEnvironment(ctx context.Context, id string) error {
	return c.call(ctx, "startEnvironment", map[string]string{"EnvironmentId": id}, nil)
}

// CreateSession returns the raw session payload to hand to session-manager-plugin.
func (c *Client) CreateSession(ctx context.Context, id string) (json.RawMessage, error) {
	var raw json.RawMessage
	err := c.call(ctx, "createSession", map[string]string{"EnvironmentId": id}, &raw)
	return raw, err
}

// DeleteEnvironment permanently deletes an environment and its persistent storage.
func (c *Client) DeleteEnvironment(ctx context.Context, id string) error {
	return c.call(ctx, "deleteEnvironment", map[string]string{"EnvironmentId": id}, nil)
}

// WaitForRunning drives an environment to RUNNING: it resumes a suspended one
// (once), waits through CREATING/RESUMING, and fails on deletion or timeout.
func (c *Client) WaitForRunning(ctx context.Context, id, initial string, timeout time.Duration) error {
	deadline := time.Now().Add(timeout)
	status := initial
	startIssued := false

	for {
		switch status {
		case "RUNNING":
			return nil
		case "DELETING", "DELETED":
			return fmt.Errorf("environment is %s; cannot connect", status)
		case "SUSPENDED", "SUSPENDING":
			if !startIssued {
				_ = c.StartEnvironment(ctx, id)
				startIssued = true
			}
		default:
			startIssued = false
		}

		if time.Now().After(deadline) {
			return fmt.Errorf("timed out waiting for environment to become available")
		}
		time.Sleep(c.pollInterval)

		if e, err := c.GetEnvironmentStatus(ctx, id); err == nil {
			status = e.Status
		}
	}
}