nadir-agent/internal/oscmd/oscmd.go

// Package oscmd holds small helpers shared by modules that shell out to system
// tools (hostnamectl, timedatectl, systemctl, …): a command runner that
// surfaces stderr, and the common "status: ok" HTTP response.
package oscmd

import (
	"bufio"
	"bytes"
	"context"
	"errors"
	"fmt"
	"log"
	"os"
	"os/exec"
	"strings"
	"sync"
	"syscall"
	"time"
)

// cmdTimeout caps how long a synchronous system command may run before it is
// killed, so a wedged tool can't tie up a request indefinitely.
//
// It bounds hangs for the plain Run path (which uses context.Background()).
// RunContext additionally propagates client cancellation. Long ops (package
// install/upgrade) use the streaming runners below, which take a ctx and are
// uncapped.
const cmdTimeout = 60 * time.Second

// CommandRunner allows overriding the command execution function for testing.
// In tests, register a handler with SetMock; the runner translates the
// MockCommand into a /bin/sh script that reproduces the stdout/stderr/exit
// behavior — no helper process, no temp file.
var CommandRunner = func(ctx context.Context, name string, args ...string) *exec.Cmd {
	mockMu.Lock()
	handler, ok := mockCmds[name]
	mockMu.Unlock()
	if !ok {
		return exec.CommandContext(ctx, name, args...)
	}
	return exec.CommandContext(ctx, "/bin/sh", "-c", mockScript(handler(args)))
}

// mockScript builds the shell script that emits a MockCommand's behavior.
// Uses printf so backslashes and percent signs in output pass through verbatim.
func mockScript(m MockCommand) string {
	var b strings.Builder
	if len(m.Lines) > 0 {
		for _, line := range m.Lines {
			b.WriteString("printf '%s\\n' ")
			b.WriteString(shellQuote(line))
			b.WriteByte('\n')
			if m.DelayMs > 0 {
				fmt.Fprintf(&b, "sleep %g\n", float64(m.DelayMs)/1000)
			}
		}
	} else {
		if m.Stdout != "" {
			b.WriteString("printf '%s' ")
			b.WriteString(shellQuote(m.Stdout))
			b.WriteByte('\n')
		}
		if m.Stderr != "" {
			b.WriteString("printf '%s' ")
			b.WriteString(shellQuote(m.Stderr))
			b.WriteString(" 1>&2\n")
		}
	}
	fmt.Fprintf(&b, "exit %d\n", m.ExitCode)
	return b.String()
}

// shellQuote returns s wrapped in single quotes, with embedded single quotes escaped.
func shellQuote(s string) string {
	return "'" + strings.ReplaceAll(s, "'", `'\''`) + "'"
}

// Run executes name with args and returns trimmed stdout. On failure it wraps
// the command's stderr (falling back to the exec error) so handlers can surface
// a meaningful message instead of a bare "exit status 1".
//
// Run uses context.Background(): it is bounded by cmdTimeout but not tied to any
// request. Callers running a slow command on behalf of a request (so a client
// disconnect should kill it) should use RunContext instead.
func Run(name string, args ...string) (string, error) {
	return RunContext(context.Background(), name, args...)
}

// RunContext is Run with a caller-supplied context. The command is killed when
// ctx is cancelled (e.g. the client disconnected) or after cmdTimeout, whichever
// comes first.
func RunContext(ctx context.Context, name string, args ...string) (string, error) {
	ctx, cancel := context.WithTimeout(ctx, cmdTimeout)
	defer cancel()
	cmd := CommandRunner(ctx, name, args...)
	var stderr strings.Builder
	cmd.Stderr = &stderr
	out, err := cmd.Output()
	if err != nil {
		if msg := strings.TrimSpace(stderr.String()); msg != "" {
			return "", errors.New(msg)
		}
		return "", err
	}
	return strings.TrimSpace(string(out)), nil
}

// RunStdin is Run with data fed to the command's stdin. Use it for secrets
// (e.g. piping "user:password" to chpasswd) so they never appear in argv/ps.
func RunStdin(stdin, name string, args ...string) (string, error) {
	ctx, cancel := context.WithTimeout(context.Background(), cmdTimeout)
	defer cancel()
	cmd := CommandRunner(ctx, name, args...)
	cmd.Stdin = strings.NewReader(stdin)
	var stderr strings.Builder
	cmd.Stderr = &stderr
	out, err := cmd.Output()
	if err != nil {
		if msg := strings.TrimSpace(stderr.String()); msg != "" {
			return "", errors.New(msg)
		}
		return "", err
	}
	return strings.TrimSpace(string(out)), nil
}

// RunLines runs the command and splits stdout into non-empty lines.
func RunLines(name string, args ...string) ([]string, error) {
	out, err := Run(name, args...)
	if err != nil {
		return nil, err
	}
	if out == "" {
		return []string{}, nil
	}
	return strings.Split(out, "\n"), nil
}

// RunStatus runs name and returns trimmed stdout plus the process exit code. It
// does NOT treat a non-zero exit as an error, because some tools signal state
// that way (dnf `check-update` exits 100 when updates exist; pacman `-Qu` exits
// 1 when there are none). err is set only when the command could not be run.
func RunStatus(name string, args ...string) (string, int, error) {
	ctx, cancel := context.WithTimeout(context.Background(), cmdTimeout)
	defer cancel()
	cmd := CommandRunner(ctx, name, args...)
	out, err := cmd.Output()
	if err != nil {
		var ee *exec.ExitError
		if errors.As(err, &ee) {
			return strings.TrimSpace(string(out)), ee.ExitCode(), nil
		}
		return "", -1, err
	}
	return strings.TrimSpace(string(out)), 0, nil
}

// RunStream runs a long-lived command (e.g. `journalctl -f`) and pushes its
// stdout lines onto the returned channel until the process exits or ctx is
// cancelled. Cancelling ctx kills the process (exec.CommandContext) and closes
// the channel, so an SSE handler can stop simply by cancelling its request
// context (client disconnect).
func RunStream(ctx context.Context, name string, args ...string) (<-chan string, error) {
	cmd := CommandRunner(ctx, name, args...)
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, err
	}
	if err := cmd.Start(); err != nil {
		return nil, err
	}
	ch := make(chan string)
	go func() {
		defer close(ch)
		defer cmd.Wait() // reap; process is already killed on ctx cancel
		sc := bufio.NewScanner(stdout)
		sc.Buffer(make([]byte, 64*1024), 1024*1024) // tolerate long log lines
		for sc.Scan() {
			select {
			case ch <- sc.Text():
			case <-ctx.Done():
				return
			}
		}
		// A scan error (e.g. a line over the buffer cap) ends the loop silently
		// otherwise; surface it so a truncated stream is diagnosable. Context
		// cancellation returns above, so this only fires on a real read error.
		if err := sc.Err(); err != nil && ctx.Err() == nil {
			log.Printf("oscmd: stream %s: %v", name, err)
		}
	}()
	return ch, nil
}

// RunStreamCombined runs a command and streams its merged stdout+stderr line by
// line, the way the command's own terminal output looks. Lines are split on \r
// as well as \n so progress redraws (apt/dnf) stream instead of buffering until
// a newline. extraEnv is appended to the process environment.
//
// It returns a lines channel and a one-shot error channel that delivers the
// process's exit status after the lines channel closes (nil = success).
// Cancelling ctx kills the process.
func RunStreamCombined(ctx context.Context, extraEnv []string, name string, args ...string) (<-chan string, <-chan error, error) {
	cmd := CommandRunner(ctx, name, args...)
	if len(extraEnv) > 0 {
		if cmd.Env == nil {
			cmd.Env = append(os.Environ(), extraEnv...)
		} else {
			cmd.Env = append(cmd.Env, extraEnv...)
		}
	}
	// One OS pipe for both streams gives the natural interleaving; passing an
	// *os.File hands the fd straight to the child (no copy goroutine).
	pr, pw, err := os.Pipe()
	if err != nil {
		return nil, nil, err
	}
	cmd.Stdout, cmd.Stderr = pw, pw
	if err := cmd.Start(); err != nil {
		pr.Close()
		pw.Close()
		return nil, nil, err
	}
	pw.Close() // the child holds its own dup; the reader sees EOF when it exits

	lines := make(chan string)
	errc := make(chan error, 1)
	go func() {
		defer close(lines)
		defer func() { errc <- cmd.Wait() }()
		defer pr.Close()
		sc := bufio.NewScanner(pr)
		sc.Buffer(make([]byte, 64*1024), 1024*1024)
		sc.Split(scanLinesCR)
		for sc.Scan() {
			if line := sc.Text(); line != "" {
				select {
				case lines <- line:
				case <-ctx.Done():
					return
				}
			}
		}
		// Surface a read error (e.g. a line over the buffer cap) so a truncated
		// stream is diagnosable. ctx cancellation returns above, so this only
		// fires on a real error; the exit status still travels via errc.
		if err := sc.Err(); err != nil && ctx.Err() == nil {
			log.Printf("oscmd: stream %s: %v", name, err)
		}
	}()
	return lines, errc, nil
}

// scanLinesCR is a bufio.SplitFunc that breaks on either \r or \n, so terminal
// progress redraws (carriage returns) surface as they happen.
func scanLinesCR(data []byte, atEOF bool) (advance int, token []byte, err error) {
	if atEOF && len(data) == 0 {
		return 0, nil, nil
	}
	if i := bytes.IndexAny(data, "\r\n"); i >= 0 {
		return i + 1, data[:i], nil
	}
	if atEOF {
		return len(data), data, nil
	}
	return 0, nil, nil // request more data
}

// ParseKV parses key=value lines (e.g. from timedatectl/systemctl show) into a map.
func ParseKV(lines []string) map[string]string {
	m := make(map[string]string, len(lines))
	for _, line := range lines {
		if k, v, ok := strings.Cut(line, "="); ok {
			m[k] = v
		}
	}
	return m
}

// RunDetached starts name with args in a new process group (Setsid) and returns
// immediately once the child has started. The child is reaped in the background
// so it does not become a zombie.
//
// Use this for operations where the synchronous path would kill the caller
// (e.g. "systemctl restart nadir" would SIGTERM the process serving the
// request before the response is written).
func RunDetached(name string, args ...string) (*StatusOutput, error) {
	cmd := exec.Command(name, args...)
	cmd.SysProcAttr = &syscall.SysProcAttr{Setsid: true}
	if err := cmd.Start(); err != nil {
		return nil, err
	}
	go cmd.Wait()
	return OK(), nil
}

// StatusOutput is the shared response for write operations that just report
// success. Reusing one type means all such endpoints share a single OpenAPI
// schema.
type StatusOutput struct {
	Body struct {
		Status string `json:"status" example:"ok" doc:"Always \"ok\" on success"`
	}
}

// OK returns a populated StatusOutput.
func OK() *StatusOutput {
	out := &StatusOutput{}
	out.Body.Status = "ok"
	return out
}

// --- Mocking helpers for testing ---------------------------------------------

// MockCommand holds the behavior for a mocked command.
type MockCommand struct {
	Stdout   string
	Stderr   string
	ExitCode int
	Lines    []string
	DelayMs  int
}

var (
	mockMu   sync.Mutex
	mockCmds = make(map[string]func(args []string) MockCommand)
)

// SetMock registers a mock handler function for the given command name.
func SetMock(name string, handler func(args []string) MockCommand) {
	mockMu.Lock()
	defer mockMu.Unlock()
	mockCmds[name] = handler
}

// ClearMocks removes all registered mock command handlers.
func ClearMocks() {
	mockMu.Lock()
	defer mockMu.Unlock()
	clear(mockCmds)
}