[codex] add context setup and env commands

README.md

@@ -59,6 +59,18 @@ unic init --force              # Overwrite existing config
 
 # Update to latest version
 unic update                    # Auto-detects install method (brew vs binary)
+
+# Print shell exports for the current context
+eval "$(unic env)"
+
+# Print shell exports for a named context
+eval "$(unic env staging-creds)"
+
+# Interactively select/setup a context, set it current, and copy exports to clipboard
+unic context setup
+
+# Clear the current context and copy cleanup commands to clipboard
+unic context unset
 ```
 
 ## Configuration
@@ -81,9 +93,19 @@ defaults:
   region: us-east-1
 
 contexts:
-  # SSO authentication
+  # SSO base context for one-step setup.
+  # `unic context setup` will log in, let you pick an account/role,
+  # then create or reuse a concrete context automatically.
   - name: dev-sso
     region: ap-northeast-2
+    profile: my-sso-profile
+    auth_type: sso
+    sso_start_url: https://my-sso-portal.awsapps.com/start
+
+  # SSO authentication
+  - name: dev-sso-123456789012-developerrole
+    region: ap-northeast-2
+    profile: my-sso-profile
     auth_type: sso
     sso_start_url: https://my-sso-portal.awsapps.com/start
     sso_account_id: "123456789012"
@@ -114,6 +136,25 @@ contexts:
 
 **Priority**: CLI flags (`--profile`, `--region`) > context settings > config defaults > hardcoded defaults (`us-east-1`)
 
+### One-Step Context Setup
+
+`unic context setup` is designed for interactive setup:
+
+```bash
+unic context setup
+```
+
+Behavior:
+
+- Prompts and progress messages go to `stderr`
+- Shell `export` / `unset` commands are copied to the clipboard
+- Credential contexts export `AWS_PROFILE` and region vars
+- Assume-role contexts export temporary STS credentials
+- SSO base contexts can log in, list accessible accounts and roles, and save a concrete context automatically
+- `~/.aws/credentials` is not modified by this flow
+
+`unic context unset` clears the `current` context from `~/.config/unic/config.yaml` and copies AWS cleanup commands to the clipboard so you can quickly reset your shell environment.
+
 ## Currently Implemented Features
 
 | Service | Feature | Status |

internal/auth/env.go

@@ -0,0 +1,129 @@
+package auth
+
+import (
+	"context"
+	"fmt"
+	"sort"
+	"strings"
+
+	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/aws/aws-sdk-go-v2/service/sts"
+
+	"unic/internal/config"
+	awsservice "unic/internal/services/aws"
+)
+
+var assumeRoleFn = assumeRoleEnv
+var resolveSSORoleFn = resolveSSORoleEnv
+
+// BuildEnvExports renders shell export commands for the given context.
+func BuildEnvExports(ctx context.Context, cfg *config.Config) (string, error) {
+	switch cfg.AuthType {
+	case config.AuthTypeSSO:
+		if cfg.SSOAccountID == "" || cfg.SSORoleName == "" {
+			return "", fmt.Errorf("context %q is an SSO base context; run `unic context setup` first", cfg.ContextName)
+		}
+		values, err := resolveSSORoleFn(ctx, cfg)
+		if err != nil {
+			return "", err
+		}
+		values["AWS_REGION"] = cfg.Region
+		values["AWS_DEFAULT_REGION"] = cfg.Region
+		values["AWS_PROFILE"] = ""
+		return renderExports(values), nil
+
+	case config.AuthTypeAssumeRole:
+		values, err := assumeRoleFn(ctx, cfg)
+		if err != nil {
+			return "", err
+		}
+		values["AWS_REGION"] = cfg.Region
+		values["AWS_DEFAULT_REGION"] = cfg.Region
+		values["AWS_PROFILE"] = ""
+		return renderExports(values), nil
+
+	case config.AuthTypeCredential, config.AuthTypeDefault:
+		profile := cfg.Profile
+		if profile == "" {
+			profile = "default"
+		}
+		return renderExports(map[string]string{
+			"AWS_PROFILE":           profile,
+			"AWS_REGION":            cfg.Region,
+			"AWS_DEFAULT_REGION":    cfg.Region,
+			"AWS_ACCESS_KEY_ID":     "",
+			"AWS_SECRET_ACCESS_KEY": "",
+			"AWS_SESSION_TOKEN":     "",
+		}), nil
+
+	default:
+		return "", fmt.Errorf("unsupported auth type %q", cfg.AuthType)
+	}
+}
+
+func assumeRoleEnv(ctx context.Context, cfg *config.Config) (map[string]string, error) {
+	baseCfg, err := awsservice.LoadBaseConfig(ctx, cfg.Region, cfg.Profile)
+	if err != nil {
+		return nil, fmt.Errorf("failed to load AWS config: %w", err)
+	}
+
+	client := sts.NewFromConfig(baseCfg)
+	input := &sts.AssumeRoleInput{
+		RoleArn:         aws.String(cfg.RoleArn),
+		RoleSessionName: aws.String("unic-env"),
+	}
+	if cfg.ExternalID != "" {
+		input.ExternalId = aws.String(cfg.ExternalID)
+	}
+
+	out, err := client.AssumeRole(ctx, input)
+	if err != nil {
+		return nil, fmt.Errorf("failed to assume role %s: %w", cfg.RoleArn, err)
+	}
+
+	creds := out.Credentials
+	return map[string]string{
+		"AWS_ACCESS_KEY_ID":     aws.ToString(creds.AccessKeyId),
+		"AWS_SECRET_ACCESS_KEY": aws.ToString(creds.SecretAccessKey),
+		"AWS_SESSION_TOKEN":     aws.ToString(creds.SessionToken),
+	}, nil
+}
+
+func resolveSSORoleEnv(ctx context.Context, cfg *config.Config) (map[string]string, error) {
+	awsCfg, err := awsservice.ResolveSSOCredentials(ctx, cfg)
+	if err != nil {
+		return nil, err
+	}
+	creds, err := awsCfg.Credentials.Retrieve(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("failed to retrieve SSO credentials: %w", err)
+	}
+	return map[string]string{
+		"AWS_ACCESS_KEY_ID":     creds.AccessKeyID,
+		"AWS_SECRET_ACCESS_KEY": creds.SecretAccessKey,
+		"AWS_SESSION_TOKEN":     creds.SessionToken,
+	}, nil
+}
+
+func renderExports(values map[string]string) string {
+	keys := make([]string, 0, len(values))
+	for key := range values {
+		keys = append(keys, key)
+	}
+	sort.Strings(keys)
+
+	lines := make([]string, 0, len(keys))
+	for _, key := range keys {
+		value := values[key]
+		if value == "" {
+			lines = append(lines, fmt.Sprintf("unset %s", key))
+			continue
+		}
+		lines = append(lines, fmt.Sprintf("export %s=%s", key, shellQuote(value)))
+	}
+	return strings.Join(lines, "\n")
+}
+
+func shellQuote(value string) string {
+	return "'" + strings.ReplaceAll(value, "'", `'\''`) + "'"
+}

internal/auth/env_test.go

@@ -0,0 +1,83 @@
+package auth
+
+import (
+	"context"
+	"strings"
+	"testing"
+
+	"unic/internal/config"
+)
+
+func TestBuildEnvExportsCredentialContext(t *testing.T) {
+	cfg := &config.Config{
+		ContextName: "dev",
+		AuthType:    config.AuthTypeCredential,
+		Profile:     "dev-profile",
+		Region:      "ap-northeast-2",
+	}
+
+	exports, err := BuildEnvExports(context.Background(), cfg)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	for _, expected := range []string{
+		"export AWS_PROFILE='dev-profile'",
+		"export AWS_REGION='ap-northeast-2'",
+		"export AWS_DEFAULT_REGION='ap-northeast-2'",
+		"unset AWS_ACCESS_KEY_ID",
+		"unset AWS_SECRET_ACCESS_KEY",
+		"unset AWS_SESSION_TOKEN",
+	} {
+		if !strings.Contains(exports, expected) {
+			t.Fatalf("expected exports to contain %q, got:\n%s", expected, exports)
+		}
+	}
+}
+
+func TestBuildEnvExportsAssumeRoleContext(t *testing.T) {
+	orig := assumeRoleFn
+	t.Cleanup(func() { assumeRoleFn = orig })
+
+	assumeRoleFn = func(ctx context.Context, cfg *config.Config) (map[string]string, error) {
+		return map[string]string{
+			"AWS_ACCESS_KEY_ID":     "AKIA123",
+			"AWS_SECRET_ACCESS_KEY": "secret",
+			"AWS_SESSION_TOKEN":     "token",
+		}, nil
+	}
+
+	cfg := &config.Config{
+		ContextName: "prod",
+		AuthType:    config.AuthTypeAssumeRole,
+		Profile:     "base",
+		Region:      "us-east-1",
+		RoleArn:     "arn:aws:iam::111111111111:role/Admin",
+	}
+
+	exports, err := BuildEnvExports(context.Background(), cfg)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if !strings.Contains(exports, "export AWS_ACCESS_KEY_ID='AKIA123'") {
+		t.Fatalf("expected assumed credentials in exports, got:\n%s", exports)
+	}
+	if !strings.Contains(exports, "unset AWS_PROFILE") {
+		t.Fatalf("expected AWS_PROFILE to be unset for temp credentials, got:\n%s", exports)
+	}
+}
+
+func TestBuildEnvExportsRejectsIncompleteSSOContext(t *testing.T) {
+	cfg := &config.Config{
+		ContextName: "base-sso",
+		AuthType:    config.AuthTypeSSO,
+		Region:      "us-east-1",
+		SSOStartURL: "https://example.awsapps.com/start",
+	}
+
+	_, err := BuildEnvExports(context.Background(), cfg)
+	if err == nil || !strings.Contains(err.Error(), "run `unic context setup` first") {
+		t.Fatalf("expected incomplete SSO error, got %v", err)
+	}
+}