lpra/internal/agent/terminal.go

package agent

import (
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/creack/pty"
)

// PTY is a pseudo-terminal running an interactive login shell.
type PTY struct {
	cmd    *exec.Cmd
	f      *os.File
	exited chan struct{} // closed once the shell has exited
	once   sync.Once
}

// hangupGrace is how long a shell gets after the hangup to save its history
// and exit before it is killed.
const hangupGrace = 3 * time.Second

// openPTYs are the terminals running now, to close on shutdown.
var openPTYs sync.Map // *PTY -> struct{}

// CloseTerminals hangs up every open terminal and waits for them, so the
// shells save their history before the agent exits.
func CloseTerminals() {
	var wg sync.WaitGroup
	openPTYs.Range(func(k, _ any) bool {
		wg.Add(1)
		go func() { defer wg.Done(); k.(*PTY).Close() }()
		return true
	})
	wg.Wait()
}

// StartPTY launches a login shell attached to a new PTY. With user empty (or
// equal to the agent's own user) it is the agent's shell - root by default,
// since the agent runs as a root systemd service. Any other user is entered via
// `su -l`, which only a root agent can do without a password. dir, when it
// is a folder, is where the shell starts instead of home.
func StartPTY(user, dir string) (*PTY, error) {
	if dir != "" && (!filepath.IsAbs(dir) || strings.ContainsAny(dir, "\x00\r\n") || !isDir(dir)) {
		dir = ""
	}
	var cmd *exec.Cmd
	env := append(os.Environ(), "TERM=xterm-256color")
	if user == "" || user == currentUser() {
		cmd = exec.Command(defaultShell(), "-l")
		// A login shell does not change directory by itself: it would start
		// where the service runs, "/". Start it at home, and say where home
		// is, as a root service often has no $HOME.
		if home := homeDir(""); isDir(home) {
			cmd.Dir = home
			env = append(env, "HOME="+home)
		}
		if dir != "" {
			cmd.Dir = dir
		}
	} else {
		if !validUserName(user) {
			return nil, fmt.Errorf("недопустимое имя пользователя %q", user)
		}
		if os.Geteuid() != 0 {
			return nil, fmt.Errorf("агент запущен от %s и не может открыть сессию от %s", currentUser(), user)
		}
		if !userExists(user) {
			return nil, fmt.Errorf("пользователь %s не найден", user)
		}
		su, err := exec.LookPath("su")
		if err != nil {
			return nil, fmt.Errorf("su не найден: %w", err)
		}
		cmd = exec.Command(su, "-l", user)
		// su -l always goes home; to start elsewhere, its shell changes
		// directory first and becomes the login shell (sh and csh alike).
		if dir != "" {
			cmd = exec.Command(su, "-l", user, "-c", "cd "+shQuote(dir)+"; exec \"$SHELL\" -l")
		}
		// su -l goes to the user's home; starting there too keeps the shell
		// out of "/" should it not.
		if home := homeDir(user); isDir(home) {
			cmd.Dir = home
		}
	}
	cmd.Env = env
	f, err := pty.Start(cmd)
	if err != nil {
		return nil, err
	}
	p := &PTY{cmd: cmd, f: f, exited: make(chan struct{})}
	go func() { _ = cmd.Wait(); close(p.exited) }()
	openPTYs.Store(p, struct{}{})
	return p, nil
}

func isDir(p string) bool {
	info, err := os.Stat(p)
	return err == nil && info.IsDir()
}

// Read reads terminal output.
func (p *PTY) Read(b []byte) (int, error) { return p.f.Read(b) }

// Write writes keystrokes to the terminal.
func (p *PTY) Write(b []byte) (int, error) { return p.f.Write(b) }

// Resize changes the terminal window size.
func (p *PTY) Resize(cols, rows uint16) {
	if cols == 0 || rows == 0 {
		return
	}
	_ = pty.Setsize(p.f, &pty.Winsize{Rows: rows, Cols: cols})
}

// Close hangs the terminal up as a closed window would: the shell gets
// SIGHUP, saves its history and exits. One that does not within hangupGrace
// is killed. Safe to call more than once.
func (p *PTY) Close() error {
	var err error
	p.once.Do(func() {
		// Closing the master hangs up the session; the signal to the
		// shell's process group makes sure it hears, through su as well.
		err = p.f.Close()
		pid := p.cmd.Process.Pid
		_ = syscall.Kill(-pid, syscall.SIGHUP)
		select {
		case <-p.exited:
		case <-time.After(hangupGrace):
			_ = syscall.Kill(-pid, syscall.SIGKILL)
			<-p.exited
		}
		openPTYs.Delete(p)
	})
	return err
}

// shQuote puts s in single quotes for sh and csh alike, closing and
// reopening them around each quote inside.
func shQuote(s string) string {
	return "'" + strings.ReplaceAll(s, "'", `'\''`) + "'"
}