banger/internal/daemon/session/session.go

// Package session contains the pure helpers of the guest-session subsystem:
// bash script generators, on-guest state path helpers, state snapshot
// parsing, and small utilities like ShellQuote and FormatStepError.
//
// The orchestrator methods (StartGuestSession, BeginGuestSessionAttach,
// etc.) stay on *daemon.Daemon and compose these helpers.
package session

import (
	"bufio"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"syscall"

	"banger/internal/model"
	"banger/internal/system"

	"golang.org/x/crypto/ssh"
)

// Constants shared between orchestration and helpers.
const (
	BackendSSH             = "ssh"
	AttachBackendNone      = "none"
	AttachBackendSSHBridge = "ssh_rehydratable"
	AttachModeExclusive    = "exclusive"
	TransportUnixSocket    = "unix_socket"
	StateRoot              = "/root/.local/state/banger/sessions"
	LogTailLineDefault     = 200
)

// StateSnapshot is the decoded per-session state as read from the guest.
type StateSnapshot struct {
	Status    string
	GuestPID  int
	ExitCode  *int
	Alive     bool
	LastError string
}

// -- Guest filesystem paths -------------------------------------------------

func StateDir(id string) string {
	return filepath.ToSlash(filepath.Join(StateRoot, id))
}

func RelativeStateDir(id string) string {
	return strings.TrimPrefix(StateDir(id), "/root/")
}

func ScriptPath(id string) string { return filepath.ToSlash(filepath.Join(StateDir(id), "run.sh")) }
func PIDPath(id string) string    { return filepath.ToSlash(filepath.Join(StateDir(id), "pid")) }
func MonitorPIDPath(id string) string {
	return filepath.ToSlash(filepath.Join(StateDir(id), "monitor_pid"))
}
func ExitCodePath(id string) string {
	return filepath.ToSlash(filepath.Join(StateDir(id), "exit_code"))
}
func StdinPipePath(id string) string {
	return filepath.ToSlash(filepath.Join(StateDir(id), "stdin.pipe"))
}
func StdinKeepalivePIDPath(id string) string {
	return filepath.ToSlash(filepath.Join(StateDir(id), "stdin_keepalive.pid"))
}
func StatusPath(id string) string { return filepath.ToSlash(filepath.Join(StateDir(id), "status")) }
func ErrorPath(id string) string  { return filepath.ToSlash(filepath.Join(StateDir(id), "error")) }
func StdoutLogPath(id string) string {
	return filepath.ToSlash(filepath.Join(StateDir(id), "stdout.log"))
}
func StderrLogPath(id string) string {
	return filepath.ToSlash(filepath.Join(StateDir(id), "stderr.log"))
}

// -- Script generators ------------------------------------------------------

// Script returns the bash runner installed into the guest for session. It
// sets up state/log paths, optional stdin fifo, and wait-loop around the
// user command.
func Script(sess model.GuestSession) string {
	var script strings.Builder
	script.WriteString("set -euo pipefail\n")
	fmt.Fprintf(&script, "STATE_DIR=%s\n", ShellQuote(sess.GuestStateDir))
	fmt.Fprintf(&script, "STDOUT_LOG=%s\n", ShellQuote(sess.StdoutLogPath))
	fmt.Fprintf(&script, "STDERR_LOG=%s\n", ShellQuote(sess.StderrLogPath))
	fmt.Fprintf(&script, "PID_FILE=%s\n", ShellQuote(PIDPath(sess.ID)))
	fmt.Fprintf(&script, "MONITOR_PID_FILE=%s\n", ShellQuote(MonitorPIDPath(sess.ID)))
	fmt.Fprintf(&script, "EXIT_FILE=%s\n", ShellQuote(ExitCodePath(sess.ID)))
	fmt.Fprintf(&script, "STATUS_FILE=%s\n", ShellQuote(StatusPath(sess.ID)))
	fmt.Fprintf(&script, "ERROR_FILE=%s\n", ShellQuote(ErrorPath(sess.ID)))
	fmt.Fprintf(&script, "STDIN_PIPE=%s\n", ShellQuote(StdinPipePath(sess.ID)))
	fmt.Fprintf(&script, "STDIN_KEEPALIVE_PID_FILE=%s\n", ShellQuote(StdinKeepalivePIDPath(sess.ID)))
	fmt.Fprintf(&script, "SESSION_CWD=%s\n", ShellQuote(DefaultCWD(sess.CWD)))
	script.WriteString("mkdir -p \"$STATE_DIR\"\n")
	script.WriteString(": >\"$STDOUT_LOG\"\n")
	script.WriteString(": >\"$STDERR_LOG\"\n")
	script.WriteString("rm -f \"$EXIT_FILE\" \"$ERROR_FILE\" \"$STDIN_KEEPALIVE_PID_FILE\"\n")
	if sess.StdinMode == model.GuestSessionStdinPipe {
		script.WriteString("rm -f \"$STDIN_PIPE\"\n")
		script.WriteString("mkfifo -m 600 \"$STDIN_PIPE\"\n")
	}
	script.WriteString("printf '%s\\n' \"${BASHPID:-$$}\" >\"$MONITOR_PID_FILE\"\n")
	script.WriteString("printf 'starting\\n' >\"$STATUS_FILE\"\n")
	script.WriteString("cd \"$SESSION_CWD\"\n")
	script.WriteString("exec > >(tee -a \"$STDOUT_LOG\") 2> >(tee -a \"$STDERR_LOG\" >&2)\n")
	for _, line := range EnvLines(sess.Env) {
		script.WriteString(line)
		script.WriteByte('\n')
	}
	script.WriteString("COMMAND=(")
	for _, value := range append([]string{sess.Command}, sess.Args...) {
		script.WriteByte(' ')
		script.WriteString(ShellQuote(value))
	}
	script.WriteString(" )\n")
	if sess.StdinMode == model.GuestSessionStdinPipe {
		script.WriteString("( while :; do sleep 3600; done ) >\"$STDIN_PIPE\" &\n")
		script.WriteString("keepalive=$!\n")
		script.WriteString("printf '%s\\n' \"$keepalive\" >\"$STDIN_KEEPALIVE_PID_FILE\"\n")
		script.WriteString("\"${COMMAND[@]}\" <\"$STDIN_PIPE\" &\n")
	} else {
		script.WriteString("\"${COMMAND[@]}\" &\n")
	}
	script.WriteString("child=$!\n")
	script.WriteString("printf '%s\\n' \"$child\" >\"$PID_FILE\"\n")
	script.WriteString("printf 'running\\n' >\"$STATUS_FILE\"\n")
	script.WriteString("wait \"$child\"\n")
	script.WriteString("rc=$?\n")
	if sess.StdinMode == model.GuestSessionStdinPipe {
		script.WriteString("if [ -f \"$STDIN_KEEPALIVE_PID_FILE\" ]; then kill \"$(cat \"$STDIN_KEEPALIVE_PID_FILE\")\" 2>/dev/null || true; fi\n")
	}
	script.WriteString("printf '%s\\n' \"$rc\" >\"$EXIT_FILE\"\n")
	script.WriteString("if [ \"$rc\" -eq 0 ]; then printf 'exited\\n' >\"$STATUS_FILE\"; else printf 'failed\\n' >\"$STATUS_FILE\"; fi\n")
	script.WriteString("exit \"$rc\"\n")
	return script.String()
}

// InspectScript reads the on-guest state files for sessionID and prints a
// key=value block parseable by ParseState.
func InspectScript(sessionID string) string {
	var script strings.Builder
	script.WriteString("set -euo pipefail\n")
	fmt.Fprintf(&script, "DIR=%s\n", ShellQuote(StateDir(sessionID)))
	script.WriteString("status=''\n")
	script.WriteString("pid=''\n")
	script.WriteString("exit_code=''\n")
	script.WriteString("last_error=''\n")
	script.WriteString("alive=false\n")
	script.WriteString("[ -f \"$DIR/status\" ] && status=\"$(cat \"$DIR/status\")\"\n")
	script.WriteString("[ -f \"$DIR/pid\" ] && pid=\"$(cat \"$DIR/pid\")\"\n")
	script.WriteString("[ -f \"$DIR/exit_code\" ] && exit_code=\"$(cat \"$DIR/exit_code\")\"\n")
	script.WriteString("[ -f \"$DIR/error\" ] && last_error=\"$(cat \"$DIR/error\")\"\n")
	script.WriteString("if [ -n \"$pid\" ] && kill -0 \"$pid\" 2>/dev/null; then alive=true; fi\n")
	script.WriteString("printf 'status=%s\\n' \"$status\"\n")
	script.WriteString("printf 'pid=%s\\n' \"$pid\"\n")
	script.WriteString("printf 'exit=%s\\n' \"$exit_code\"\n")
	script.WriteString("printf 'alive=%s\\n' \"$alive\"\n")
	script.WriteString("printf 'error=%s\\n' \"$last_error\"\n")
	return script.String()
}

// SignalScript sends signal to sessionID's runner and monitor processes.
func SignalScript(sessionID, signal string) string {
	var script strings.Builder
	script.WriteString("set -euo pipefail\n")
	fmt.Fprintf(&script, "DIR=%s\n", ShellQuote(StateDir(sessionID)))
	fmt.Fprintf(&script, "SIGNAL=%s\n", ShellQuote(signal))
	script.WriteString("pid=''\n")
	script.WriteString("monitor=''\n")
	script.WriteString("keepalive=''\n")
	script.WriteString("[ -f \"$DIR/pid\" ] && pid=\"$(cat \"$DIR/pid\")\"\n")
	script.WriteString("[ -f \"$DIR/monitor_pid\" ] && monitor=\"$(cat \"$DIR/monitor_pid\")\"\n")
	script.WriteString("[ -f \"$DIR/stdin_keepalive.pid\" ] && keepalive=\"$(cat \"$DIR/stdin_keepalive.pid\")\"\n")
	script.WriteString("printf 'stopping\\n' >\"$DIR/status\"\n")
	script.WriteString("if [ -n \"$pid\" ]; then kill -${SIGNAL} \"$pid\" 2>/dev/null || true; fi\n")
	script.WriteString("if [ -n \"$monitor\" ]; then kill -${SIGNAL} \"$monitor\" 2>/dev/null || true; fi\n")
	script.WriteString("if [ -n \"$keepalive\" ]; then kill -${SIGNAL} \"$keepalive\" 2>/dev/null || true; fi\n")
	return script.String()
}

// CWDPreflightScript verifies cwd exists on the guest.
func CWDPreflightScript(cwd string) string {
	var script strings.Builder
	script.WriteString("set -euo pipefail\n")
	fmt.Fprintf(&script, "DIR=%s\n", ShellQuote(DefaultCWD(cwd)))
	script.WriteString("if [ ! -d \"$DIR\" ]; then echo \"missing cwd: $DIR\"; exit 1; fi\n")
	return script.String()
}

// CommandPreflightScript verifies each command is resolvable on the guest.
func CommandPreflightScript(commands []string) string {
	var script strings.Builder
	script.WriteString("set -euo pipefail\n")
	script.WriteString("check_command() {\n")
	script.WriteString("  cmd=\"$1\"\n")
	script.WriteString("  case \"$cmd\" in\n")
	script.WriteString("    */*) [ -x \"$cmd\" ] || { echo \"missing command: $cmd\"; exit 1; } ;;\n")
	script.WriteString("    *) command -v \"$cmd\" >/dev/null 2>&1 || { echo \"missing command: $cmd\"; exit 1; } ;;\n")
	script.WriteString("  esac\n")
	script.WriteString("}\n")
	for _, command := range commands {
		fmt.Fprintf(&script, "check_command %s\n", ShellQuote(command))
	}
	return script.String()
}

// AttachInputCommand returns the guest command that creates/opens the stdin
// fifo for sessionID and cats attach-side bytes into it.
func AttachInputCommand(sessionID string) string {
	path := StdinPipePath(sessionID)
	return "bash -lc " + ShellQuote(fmt.Sprintf("set -euo pipefail\n[ -p %s ] || mkfifo -m 600 %s\nexec cat > %s\n", ShellQuote(path), ShellQuote(path), ShellQuote(path)))
}

// AttachTailCommand returns the guest command that tails a log file and
// streams new content back to the attach bridge.
func AttachTailCommand(path string) string {
	return "bash -lc " + ShellQuote(fmt.Sprintf("set -euo pipefail\ntouch %s\nexec tail -n 0 -F %s 2>/dev/null\n", ShellQuote(path), ShellQuote(path)))
}

// EnvLines returns deterministic `export KEY=value` lines for the session
// launcher, ordered by key.
func EnvLines(values map[string]string) []string {
	if len(values) == 0 {
		return nil
	}
	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 {
		lines = append(lines, "export "+key+"="+ShellQuote(values[key]))
	}
	return lines
}

// -- State snapshot helpers -------------------------------------------------

// ParseState decodes the key=value output produced by InspectScript.
func ParseState(raw string) (StateSnapshot, error) {
	var snapshot StateSnapshot
	scanner := bufio.NewScanner(strings.NewReader(raw))
	for scanner.Scan() {
		line := scanner.Text()
		key, value, ok := strings.Cut(line, "=")
		if !ok {
			continue
		}
		switch strings.TrimSpace(key) {
		case "status":
			snapshot.Status = strings.TrimSpace(value)
		case "pid":
			if pid, err := strconv.Atoi(strings.TrimSpace(value)); err == nil {
				snapshot.GuestPID = pid
			}
		case "exit":
			if exitCode, err := strconv.Atoi(strings.TrimSpace(value)); err == nil {
				snapshot.ExitCode = &exitCode
			}
		case "alive":
			snapshot.Alive = strings.TrimSpace(value) == "true"
		case "error":
			snapshot.LastError = strings.TrimSpace(value)
		}
	}
	return snapshot, scanner.Err()
}

// InspectStateFromDir reads the state files directly from stateDir (used
// when the guest is offline and we can mount the work disk from the host).
func InspectStateFromDir(stateDir string) (StateSnapshot, error) {
	var snapshot StateSnapshot
	statusData, _ := os.ReadFile(filepath.Join(stateDir, "status"))
	snapshot.Status = strings.TrimSpace(string(statusData))
	pidData, _ := os.ReadFile(filepath.Join(stateDir, "pid"))
	if pidValue, err := strconv.Atoi(strings.TrimSpace(string(pidData))); err == nil {
		snapshot.GuestPID = pidValue
	}
	exitData, _ := os.ReadFile(filepath.Join(stateDir, "exit_code"))
	if exitValue, err := strconv.Atoi(strings.TrimSpace(string(exitData))); err == nil {
		snapshot.ExitCode = &exitValue
	}
	errorData, _ := os.ReadFile(filepath.Join(stateDir, "error"))
	snapshot.LastError = strings.TrimSpace(string(errorData))
	if snapshot.GuestPID != 0 {
		snapshot.Alive = ProcessAlive(snapshot.GuestPID)
	}
	return snapshot, nil
}

// ApplyStateSnapshot mutates sess in place to reflect snapshot. vmRunning
// captures whether the VM is currently up so stale in-flight sessions can be
// failed when the VM is gone.
func ApplyStateSnapshot(sess *model.GuestSession, snapshot StateSnapshot, vmRunning bool) {
	if sess == nil {
		return
	}
	if snapshot.GuestPID != 0 {
		sess.GuestPID = snapshot.GuestPID
	}
	if snapshot.LastError != "" {
		sess.LastError = snapshot.LastError
	}
	if snapshot.ExitCode != nil {
		sess.ExitCode = snapshot.ExitCode
		sess.Attachable = false
		sess.Reattachable = false
		if sess.StartedAt.IsZero() {
			sess.StartedAt = model.Now()
		}
		if sess.EndedAt.IsZero() {
			sess.EndedAt = model.Now()
		}
		if *snapshot.ExitCode == 0 {
			sess.Status = model.GuestSessionStatusExited
		} else {
			sess.Status = model.GuestSessionStatusFailed
		}
		return
	}
	if snapshot.Alive {
		if sess.StartedAt.IsZero() {
			sess.StartedAt = model.Now()
		}
		sess.Status = model.GuestSessionStatusRunning
		return
	}
	if !vmRunning && (sess.Status == model.GuestSessionStatusStarting || sess.Status == model.GuestSessionStatusRunning || sess.Status == model.GuestSessionStatusStopping) {
		sess.Status = model.GuestSessionStatusFailed
		sess.Attachable = false
		sess.Reattachable = false
		if sess.LastError == "" {
			sess.LastError = "vm is not running"
		}
		if sess.EndedAt.IsZero() {
			sess.EndedAt = model.Now()
		}
		return
	}
	if snapshot.Status == string(model.GuestSessionStatusRunning) {
		if sess.StartedAt.IsZero() {
			sess.StartedAt = model.Now()
		}
		sess.Status = model.GuestSessionStatusRunning
	}
	if sess.Status == model.GuestSessionStatusRunning && sess.StdinMode == model.GuestSessionStdinPipe {
		sess.Attachable = true
		sess.Reattachable = true
		if sess.AttachBackend == "" {
			sess.AttachBackend = AttachBackendSSHBridge
		}
		if sess.AttachMode == "" {
			sess.AttachMode = AttachModeExclusive
		}
	}
}

// StateChanged reports whether any materially observable field differs
// between before and after, guiding whether to persist an update.
func StateChanged(before, after model.GuestSession) bool {
	if before.Status != after.Status || before.GuestPID != after.GuestPID || before.LastError != after.LastError || before.Attachable != after.Attachable || before.Reattachable != after.Reattachable || before.AttachBackend != after.AttachBackend || before.AttachMode != after.AttachMode || before.LaunchStage != after.LaunchStage || before.LaunchMessage != after.LaunchMessage || before.LaunchRawLog != after.LaunchRawLog {
		return true
	}
	if before.StartedAt != after.StartedAt || before.EndedAt != after.EndedAt {
		return true
	}
	switch {
	case before.ExitCode == nil && after.ExitCode == nil:
		return false
	case before.ExitCode == nil || after.ExitCode == nil:
		return true
	default:
		return *before.ExitCode != *after.ExitCode
	}
}

// -- Launch helpers ---------------------------------------------------------

// DefaultName returns a friendly session name: caller-provided if non-empty,
// otherwise `<command-base>-<short-id>`.
func DefaultName(id, command, explicit string) string {
	if trimmed := strings.TrimSpace(explicit); trimmed != "" {
		return trimmed
	}
	base := filepath.Base(strings.TrimSpace(command))
	if base == "." || base == string(filepath.Separator) || base == "" {
		base = "session"
	}
	return base + "-" + system.ShortID(id)
}

// DefaultCWD returns value if non-empty, else /root.
func DefaultCWD(value string) string {
	if trimmed := strings.TrimSpace(value); trimmed != "" {
		return trimmed
	}
	return "/root"
}

// FailLaunch annotates sess as launch-failed with stage/message/raw log and
// returns it for persistence.
func FailLaunch(sess model.GuestSession, stage, message, rawLog string) model.GuestSession {
	now := model.Now()
	sess.Status = model.GuestSessionStatusFailed
	sess.LastError = strings.TrimSpace(message)
	sess.Attachable = false
	sess.Reattachable = false
	sess.LaunchStage = strings.TrimSpace(stage)
	sess.LaunchMessage = strings.TrimSpace(message)
	sess.LaunchRawLog = strings.TrimSpace(rawLog)
	sess.UpdatedAt = now
	sess.EndedAt = now
	return sess
}

// NormalizeRequiredCommands returns a de-duplicated, order-preserving list
// of required commands, with the session command first.
func NormalizeRequiredCommands(command string, extras []string) []string {
	ordered := make([]string, 0, len(extras)+1)
	seen := map[string]struct{}{}
	appendValue := func(value string) {
		trimmed := strings.TrimSpace(value)
		if trimmed == "" {
			return
		}
		if _, ok := seen[trimmed]; ok {
			return
		}
		seen[trimmed] = struct{}{}
		ordered = append(ordered, trimmed)
	}
	appendValue(command)
	for _, extra := range extras {
		appendValue(extra)
	}
	return ordered
}

// -- Small utilities --------------------------------------------------------

// ShellQuote returns value single-quoted for bash, escaping embedded quotes.
func ShellQuote(value string) string {
	return "'" + strings.ReplaceAll(value, "'", `'"'"'`) + "'"
}

// ExitCode extracts the exit status from an ssh.ExitError, returning
// (0, true) for nil errors.
func ExitCode(err error) (int, bool) {
	if err == nil {
		return 0, true
	}
	var exitErr *ssh.ExitError
	if errors.As(err, &exitErr) {
		return exitErr.ExitStatus(), true
	}
	return 0, false
}

// CloneStringMap returns a shallow copy of values, or nil if empty.
func CloneStringMap(values map[string]string) map[string]string {
	if len(values) == 0 {
		return nil
	}
	cloned := make(map[string]string, len(values))
	for key, value := range values {
		cloned[key] = value
	}
	return cloned
}

// TailFileContent returns the last N lines of a file, or "" if the file is
// missing.
func TailFileContent(path string, lines int) (string, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		if os.IsNotExist(err) {
			return "", nil
		}
		return "", err
	}
	if lines <= 0 {
		return string(data), nil
	}
	parts := strings.Split(string(data), "\n")
	if len(parts) <= lines {
		return string(data), nil
	}
	return strings.Join(parts[len(parts)-lines-1:], "\n"), nil
}

// ProcessAlive returns true if the process with pid exists. The syscallKill
// override is exposed for tests that need to simulate alive/dead processes.
func ProcessAlive(pid int) bool {
	if pid <= 0 {
		return false
	}
	return syscallKill(pid, syscall.Signal(0)) == nil
}

// syscallKill is a test seam: tests replace it to stub process-alive checks.
var syscallKill = func(pid int, signal os.Signal) error {
	proc, err := os.FindProcess(pid)
	if err != nil {
		return err
	}
	return proc.Signal(signal)
}

// FormatStepError wraps err with an action label and trimmed on-guest log.
func FormatStepError(action string, err error, log string) error {
	log = strings.TrimSpace(log)
	if log == "" {
		return fmt.Errorf("%s: %w", action, err)
	}
	return fmt.Errorf("%s: %w: %s", action, err, log)
}