session.go

package main

import (
	"fmt"
	"io"
	"os"
	"os/exec"
	"os/signal"
	"syscall"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/creack/pty"
	"golang.org/x/term"
)

func pluginPath() (string, error) {
	p, err := exec.LookPath("session-manager-plugin")
	if err != nil {
		return "", fmt.Errorf("session-manager-plugin not found in PATH; install it: " +
			"https://docs.aws.amazon.com/systems-manager/latest/userguide/session-manager-working-with-install-plugin.html")
	}
	return p, nil
}

// launchSession execs session-manager-plugin with the session payload, inheriting
// the terminal so the user gets an interactive shell.
func launchSession(raw []byte, region string) error {
	path, err := pluginPath()
	if err != nil {
		return err
	}
	cmd := exec.Command(path, string(raw), region, "StartSession")
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	return cmd.Run()
}

// launchSessionInject runs session-manager-plugin on a PTY and, after a short
// delay, injects export commands so the shell has working AWS credentials. This
// mirrors the VSCode extension's opt-in injection (HISTCONTROL + a leading space
// keep the secrets out of history; `clear` wipes them from the screen).
func launchSessionInject(raw []byte, region, injectCmd string) error {
	path, err := pluginPath()
	if err != nil {
		return err
	}
	if !term.IsTerminal(int(os.Stdin.Fd())) {
		return fmt.Errorf("-inject requires an interactive terminal")
	}

	cmd := exec.Command(path, string(raw), region, "StartSession")
	ptmx, err := pty.Start(cmd)
	if err != nil {
		return err
	}
	defer func() { _ = ptmx.Close() }()

	// Keep the PTY sized to the real terminal.
	ch := make(chan os.Signal, 1)
	signal.Notify(ch, syscall.SIGWINCH)
	go func() {
		for range ch {
			_ = pty.InheritSize(os.Stdin, ptmx)
		}
	}()
	ch <- syscall.SIGWINCH
	defer signal.Stop(ch)

	oldState, err := term.MakeRaw(int(os.Stdin.Fd()))
	if err != nil {
		return err
	}
	defer func() { _ = term.Restore(int(os.Stdin.Fd()), oldState) }()

	go func() { _, _ = io.Copy(ptmx, os.Stdin) }()
	go func() {
		time.Sleep(5 * time.Second)
		_, _ = io.WriteString(ptmx, injectCmd)
	}()

	_, _ = io.Copy(os.Stdout, ptmx)
	return nil
}

// buildInjectCommand builds the shell input that exports the given credentials.
// `\r` acts as Enter inside the raw PTY.
func buildInjectCommand(region string, cr aws.Credentials) string {
	exports := " export AWS_DEFAULT_REGION='" + region + "' AWS_REGION='" + region + "'" +
		" AWS_ACCESS_KEY_ID='" + cr.AccessKeyID + "'" +
		" AWS_SECRET_ACCESS_KEY='" + cr.SecretAccessKey + "'"
	if cr.SessionToken != "" {
		exports += " AWS_SESSION_TOKEN='" + cr.SessionToken + "'"
	}
	return "export HISTCONTROL=ignorespace\r" + exports + "\r clear\r"
}