|
| 1 | +package main |
| 2 | + |
| 3 | +import ( |
| 4 | + "bytes" |
| 5 | + "context" |
| 6 | + "crypto/sha256" |
| 7 | + "encoding/hex" |
| 8 | + "encoding/json" |
| 9 | + "fmt" |
| 10 | + "io" |
| 11 | + "net/http" |
| 12 | + "time" |
| 13 | + |
| 14 | + "github.com/aws/aws-sdk-go-v2/aws" |
| 15 | + v4 "github.com/aws/aws-sdk-go-v2/aws/signer/v4" |
| 16 | +) |
| 17 | + |
| 18 | +// service is the SigV4 service name for the (unofficial) CloudShell API. |
| 19 | +const service = "cloudshell" |
| 20 | + |
| 21 | +// VpcConfig optionally attaches the environment to a VPC. |
| 22 | +type VpcConfig struct { |
| 23 | + VpcId string `json:"VpcId"` |
| 24 | + SecurityGroupIds []string `json:"SecurityGroupIds"` |
| 25 | + SubnetIds []string `json:"SubnetIds"` |
| 26 | +} |
| 27 | + |
| 28 | +// Environment is a CloudShell environment. Note: describeEnvironments returns |
| 29 | +// only EnvironmentId; Status/VpcConfig come from getEnvironmentStatus. |
| 30 | +type Environment struct { |
| 31 | + EnvironmentId string `json:"EnvironmentId"` |
| 32 | + Status string `json:"Status,omitempty"` |
| 33 | + EnvironmentName string `json:"EnvironmentName,omitempty"` |
| 34 | + StatusReason string `json:"StatusReason,omitempty"` |
| 35 | + VpcConfig *VpcConfig `json:"VpcConfig,omitempty"` |
| 36 | +} |
| 37 | + |
| 38 | +// Client talks to the CloudShell JSON API, signing each request with SigV4. |
| 39 | +type Client struct { |
| 40 | + region string |
| 41 | + creds aws.CredentialsProvider |
| 42 | + signer *v4.Signer |
| 43 | + http *http.Client |
| 44 | + pollInterval time.Duration |
| 45 | +} |
| 46 | + |
| 47 | +// Credentials resolves the current AWS credentials (e.g. for shell injection). |
| 48 | +func (c *Client) Credentials(ctx context.Context) (aws.Credentials, error) { |
| 49 | + return c.creds.Retrieve(ctx) |
| 50 | +} |
| 51 | + |
| 52 | +// NewClient builds a Client for the given region and credential provider. |
| 53 | +func NewClient(region string, creds aws.CredentialsProvider) *Client { |
| 54 | + return &Client{ |
| 55 | + region: region, |
| 56 | + creds: creds, |
| 57 | + signer: v4.NewSigner(), |
| 58 | + http: http.DefaultClient, |
| 59 | + pollInterval: 2 * time.Second, |
| 60 | + } |
| 61 | +} |
| 62 | + |
| 63 | +// call POSTs a SigV4-signed JSON request to the given API action and decodes the |
| 64 | +// response into out (which may be nil, or a *json.RawMessage for the raw body). |
| 65 | +func (c *Client) call(ctx context.Context, action string, body, out any) error { |
| 66 | + payload, err := json.Marshal(body) |
| 67 | + if err != nil { |
| 68 | + return err |
| 69 | + } |
| 70 | + |
| 71 | + url := fmt.Sprintf("https://%s.%s.amazonaws.com/%s", service, c.region, action) |
| 72 | + req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(payload)) |
| 73 | + if err != nil { |
| 74 | + return err |
| 75 | + } |
| 76 | + req.Header.Set("Content-Type", "application/json") |
| 77 | + |
| 78 | + creds, err := c.creds.Retrieve(ctx) |
| 79 | + if err != nil { |
| 80 | + return fmt.Errorf("resolve credentials: %w", err) |
| 81 | + } |
| 82 | + sum := sha256.Sum256(payload) |
| 83 | + if err := c.signer.SignHTTP(ctx, creds, req, hex.EncodeToString(sum[:]), service, c.region, time.Now()); err != nil { |
| 84 | + return fmt.Errorf("sign request: %w", err) |
| 85 | + } |
| 86 | + |
| 87 | + resp, err := c.http.Do(req) |
| 88 | + if err != nil { |
| 89 | + return err |
| 90 | + } |
| 91 | + defer resp.Body.Close() |
| 92 | + |
| 93 | + data, _ := io.ReadAll(resp.Body) |
| 94 | + if resp.StatusCode/100 != 2 { |
| 95 | + return fmt.Errorf("%s failed (%d): %s", action, resp.StatusCode, string(data)) |
| 96 | + } |
| 97 | + if out != nil && len(data) > 0 { |
| 98 | + return json.Unmarshal(data, out) |
| 99 | + } |
| 100 | + return nil |
| 101 | +} |
| 102 | + |
| 103 | +// DescribeEnvironments lists environments. The response shape varies, so it is |
| 104 | +// coerced defensively (wrapped object, bare array, or single object). |
| 105 | +func (c *Client) DescribeEnvironments(ctx context.Context) ([]Environment, error) { |
| 106 | + var raw json.RawMessage |
| 107 | + if err := c.call(ctx, "describeEnvironments", map[string]any{}, &raw); err != nil { |
| 108 | + return nil, err |
| 109 | + } |
| 110 | + |
| 111 | + var wrapped struct { |
| 112 | + Environments []Environment `json:"Environments"` |
| 113 | + } |
| 114 | + if json.Unmarshal(raw, &wrapped) == nil && wrapped.Environments != nil { |
| 115 | + return wrapped.Environments, nil |
| 116 | + } |
| 117 | + var arr []Environment |
| 118 | + if json.Unmarshal(raw, &arr) == nil { |
| 119 | + return arr, nil |
| 120 | + } |
| 121 | + var single Environment |
| 122 | + if json.Unmarshal(raw, &single) == nil && single.EnvironmentId != "" { |
| 123 | + return []Environment{single}, nil |
| 124 | + } |
| 125 | + return nil, nil |
| 126 | +} |
| 127 | + |
| 128 | +// DescribeEnvironmentsWithStatus enriches each environment with its real status |
| 129 | +// (describeEnvironments returns only the id). Environments that can't be resolved |
| 130 | +// — e.g. lingering after deletion — are marked DELETED. |
| 131 | +func (c *Client) DescribeEnvironmentsWithStatus(ctx context.Context) ([]Environment, error) { |
| 132 | + envs, err := c.DescribeEnvironments(ctx) |
| 133 | + if err != nil { |
| 134 | + return nil, err |
| 135 | + } |
| 136 | + for i := range envs { |
| 137 | + s, err := c.GetEnvironmentStatus(ctx, envs[i].EnvironmentId) |
| 138 | + if err != nil { |
| 139 | + envs[i].Status = "DELETED" |
| 140 | + continue |
| 141 | + } |
| 142 | + envs[i].Status = s.Status |
| 143 | + envs[i].VpcConfig = s.VpcConfig |
| 144 | + if s.EnvironmentName != "" { |
| 145 | + envs[i].EnvironmentName = s.EnvironmentName |
| 146 | + } |
| 147 | + } |
| 148 | + return envs, nil |
| 149 | +} |
| 150 | + |
| 151 | +// GetEnvironmentStatus returns the current status (and VpcConfig) of an environment. |
| 152 | +func (c *Client) GetEnvironmentStatus(ctx context.Context, id string) (*Environment, error) { |
| 153 | + var e Environment |
| 154 | + err := c.call(ctx, "getEnvironmentStatus", map[string]string{"EnvironmentId": id}, &e) |
| 155 | + return &e, err |
| 156 | +} |
| 157 | + |
| 158 | +// CreateEnvironment creates an environment, optionally attached to a VPC. When an |
| 159 | +// environment already exists, CloudShell returns the existing one (idempotent). |
| 160 | +func (c *Client) CreateEnvironment(ctx context.Context, vpc *VpcConfig) (*Environment, error) { |
| 161 | + body := map[string]any{} |
| 162 | + if vpc != nil { |
| 163 | + body["EnvironmentName"] = envName |
| 164 | + body["VpcConfig"] = vpc |
| 165 | + } |
| 166 | + var e Environment |
| 167 | + err := c.call(ctx, "createEnvironment", body, &e) |
| 168 | + return &e, err |
| 169 | +} |
| 170 | + |
| 171 | +// StartEnvironment resumes a suspended environment. |
| 172 | +func (c *Client) StartEnvironment(ctx context.Context, id string) error { |
| 173 | + return c.call(ctx, "startEnvironment", map[string]string{"EnvironmentId": id}, nil) |
| 174 | +} |
| 175 | + |
| 176 | +// CreateSession returns the raw session payload to hand to session-manager-plugin. |
| 177 | +func (c *Client) CreateSession(ctx context.Context, id string) (json.RawMessage, error) { |
| 178 | + var raw json.RawMessage |
| 179 | + err := c.call(ctx, "createSession", map[string]string{"EnvironmentId": id}, &raw) |
| 180 | + return raw, err |
| 181 | +} |
| 182 | + |
| 183 | +// DeleteEnvironment permanently deletes an environment and its persistent storage. |
| 184 | +func (c *Client) DeleteEnvironment(ctx context.Context, id string) error { |
| 185 | + return c.call(ctx, "deleteEnvironment", map[string]string{"EnvironmentId": id}, nil) |
| 186 | +} |
| 187 | + |
| 188 | +// WaitForRunning drives an environment to RUNNING: it resumes a suspended one |
| 189 | +// (once), waits through CREATING/RESUMING, and fails on deletion or timeout. |
| 190 | +func (c *Client) WaitForRunning(ctx context.Context, id, initial string, timeout time.Duration) error { |
| 191 | + deadline := time.Now().Add(timeout) |
| 192 | + status := initial |
| 193 | + startIssued := false |
| 194 | + |
| 195 | + for { |
| 196 | + switch status { |
| 197 | + case "RUNNING": |
| 198 | + return nil |
| 199 | + case "DELETING", "DELETED": |
| 200 | + return fmt.Errorf("environment is %s; cannot connect", status) |
| 201 | + case "SUSPENDED", "SUSPENDING": |
| 202 | + if !startIssued { |
| 203 | + _ = c.StartEnvironment(ctx, id) |
| 204 | + startIssued = true |
| 205 | + } |
| 206 | + default: |
| 207 | + startIssued = false |
| 208 | + } |
| 209 | + |
| 210 | + if time.Now().After(deadline) { |
| 211 | + return fmt.Errorf("timed out waiting for environment to become available") |
| 212 | + } |
| 213 | + time.Sleep(c.pollInterval) |
| 214 | + |
| 215 | + if e, err := c.GetEnvironmentStatus(ctx, id); err == nil { |
| 216 | + status = e.Status |
| 217 | + } |
| 218 | + } |
| 219 | +} |
0 commit comments