gosh/prompt.go

package main

// Custom prompts with $PS1 and $PS2, using the backslash escapes of bash:
//
//	\u user name          \h host name up to the first '.'   \H host name
//	\w working directory  \W its last component (~ for home) \$ # for root, else $
//	\t time 24h HH:MM:SS  \T 12h HH:MM:SS  \@ 12h am/pm  \A 24h HH:MM  \d date
//	\j number of jobs     \l terminal      \s shell name  \v \V version
//	\n newline  \r carriage return  \a bell  \e escape  \nnn octal  \\ backslash
//	\[ ... \]   begin/end of characters that take no space (escape sequences)
//
// After the escapes, $VAR, ${...}, $(...) and $((...)) are expanded (in the
// same pass, so the values of escapes are not expanded again).

import (
	"fmt"
	"os"
	"os/user"
	"path/filepath"
	"strconv"
	"strings"
	"time"
)

// promptInfo holds the values for the escapes of a prompt string.
type promptInfo struct {
	user, host, wd, home string
	root                 bool
	jobs                 int
	tty                  string
	now                  time.Time
}

// currentPromptInfo collects the promptInfo of the running shell.
func currentPromptInfo() promptInfo {
	info := promptInfo{user: "?", host: "unknown", wd: "?", home: getHomeDir(),
		root: os.Geteuid() == 0, jobs: len(jobTable), now: time.Now()}
	if wd, err := os.Getwd(); err == nil {
		info.wd = wd
	}
	if usr, err := user.Current(); err == nil && usr.Username != "" {
		info.user = usr.Username
	}
	if host, err := os.Hostname(); err == nil {
		info.host = host
	}
	if tty, err := os.Readlink("/proc/self/fd/0"); err == nil && strings.HasPrefix(tty, "/dev/") {
		info.tty = filepath.Base(tty)
	}
	return info
}

// expandPrompt expands the escapes and $ expressions of a PS1 or PS2 value.
// The values of escapes like \w and of expansions are sanitized, so a
// directory or git branch name cannot write escape sequences to the
// terminal; colors have to be written as \e[...m in the prompt itself.
func expandPrompt(ps string, info promptInfo) (string, error) {
	// $(...) in the prompt must not change $? or the status of the next
	// command line
	defer func(status int) { substStatus = status }(substStatus)

	var b strings.Builder
	for i := 0; i < len(ps); i++ {
		ch := ps[i]
		if ch == '$' {
			value, n, err := expandDollar(ps[i:])
			if err != nil {
				return "", err
			}
			b.WriteString(sanitize(value))
			i += n - 1
			continue
		}
		if ch != '\\' || i+1 == len(ps) {
			b.WriteByte(ch)
			continue
		}
		i++
		switch esc := ps[i]; esc {
		case 'u':
			b.WriteString(sanitize(info.user))
		case 'h':
			host, _, _ := strings.Cut(info.host, ".")
			b.WriteString(sanitize(host))
		case 'H':
			b.WriteString(sanitize(info.host))
		case 'w':
			b.WriteString(sanitize(tildePath(info.wd, info.home)))
		case 'W':
			b.WriteString(sanitize(baseDir(info.wd, info.home)))
		case '$':
			if info.root {
				b.WriteByte('#')
			} else {
				b.WriteByte('$')
			}
		case 't':
			b.WriteString(info.now.Format("15:04:05"))
		case 'T':
			b.WriteString(info.now.Format("03:04:05"))
		case '@':
			b.WriteString(info.now.Format("03:04 PM"))
		case 'A':
			b.WriteString(info.now.Format("15:04"))
		case 'd':
			b.WriteString(info.now.Format("Mon Jan 02"))
		case 'j':
			b.WriteString(strconv.Itoa(info.jobs))
		case 'l':
			b.WriteString(sanitize(info.tty))
		case 's':
			b.WriteString("gosh")
		case 'v':
			b.WriteString(shortVersion(version))
		case 'V':
			b.WriteString(sanitize(version))
		case 'n':
			b.WriteByte('\n')
		case 'r':
			b.WriteByte('\r')
		case 'a':
			b.WriteByte('\a')
		case 'e':
			b.WriteByte('\x1b')
		case '\\':
			b.WriteByte('\\')
		case '[', ']':
			// Markers for bash's readline. Our readline finds the width of
			// the prompt itself (see splitPrompt).
		default:
			if n := octalLen(ps[i:]); n > 0 {
				code, _ := strconv.ParseUint(ps[i:i+n], 8, 8)
				b.WriteByte(byte(code))
				i += n - 1
			} else {
				// Unknown escapes are kept, like in bash
				b.WriteByte('\\')
				b.WriteByte(esc)
			}
		}
	}
	return b.String(), nil
}

// tildePath replaces the home directory at the start of path with ~.
func tildePath(path, home string) string {
	if home != "" && home != "/" && (path == home || strings.HasPrefix(path, home+"/")) {
		return "~" + strings.TrimPrefix(path, home)
	}
	return path
}

// baseDir returns the last component of path for \W: ~ for the home
// directory, / for the root directory.
func baseDir(path, home string) string {
	if home != "" && path == home {
		return "~"
	}
	if path == "/" {
		return path
	}
	return filepath.Base(path)
}

// shortVersion returns major.minor of a version like v0.19.0 for \v.
func shortVersion(v string) string {
	v = strings.TrimPrefix(v, "v")
	parts := strings.SplitN(v, ".", 3)
	if len(parts) >= 2 {
		return sanitize(parts[0] + "." + parts[1])
	}
	return sanitize(v)
}

// octalLen returns the length of the octal number (up to 3 digits, at most
// 0377) at the start of s.
func octalLen(s string) int {
	n := 0
	for n < len(s) && n < 3 && s[n] >= '0' && s[n] <= '7' {
		n++
	}
	if n == 3 && s[0] > '3' {
		n = 2
	}
	return n
}

// splitPrompt splits an expanded prompt into what is written to the terminal
// before readline starts and the prompt that readline shows. readline only
// handles a single line and only knows the width of SGR sequences
// (\e[...m), so all lines but the last and every other control sequence of
// the last line (e.g. \e]0;title\a for the window title) go into prefix.
func splitPrompt(prompt string) (prefix, line string) {
	var pre strings.Builder
	if i := strings.LastIndexByte(prompt, '\n'); i >= 0 {
		pre.WriteString(prompt[:i+1])
		prompt = prompt[i+1:]
	}
	var b strings.Builder
	for i := 0; i < len(prompt); {
		ch := prompt[i]
		switch {
		case ch == '\x1b':
			n := escapeLen(prompt[i:])
			if seq := prompt[i : i+n]; isSGRSequence(seq) {
				b.WriteString(seq)
			} else {
				pre.WriteString(seq)
			}
			i += n
		case ch < ' ' || ch == 0x7f:
			pre.WriteByte(ch)
			i++
		default:
			b.WriteByte(ch)
			i++
		}
	}
	return pre.String(), b.String()
}

// escapeLen returns the length of the escape sequence at the start of s
// (which starts with ESC): a CSI sequence ESC [ ... final byte, an OSC
// sequence ESC ] ... BEL or ESC \, or ESC followed by one character.
func escapeLen(s string) int {
	if len(s) < 2 {
		return len(s)
	}
	switch s[1] {
	case '[':
		for i := 2; i < len(s); i++ {
			if s[i] >= 0x40 && s[i] <= 0x7e {
				return i + 1
			}
		}
		return len(s)
	case ']', 'P', '_', '^':
		for i := 2; i < len(s); i++ {
			if s[i] == '\a' {
				return i + 1
			}
			if s[i] == '\x1b' && i+1 < len(s) && s[i+1] == '\\' {
				return i + 2
			}
		}
		return len(s)
	}
	return 2
}

// isSGRSequence reports whether seq is a complete ESC [ ... m sequence.
func isSGRSequence(seq string) bool {
	return len(seq) >= 3 && strings.HasPrefix(seq, "\x1b[") && strings.HasSuffix(seq, "m") && isSGR(seq[2:len(seq)-1])
}

// promptString returns the expanded value of the prompt variable name (PS1
// or PS2), or ok == false if it is not set. An error in the prompt is reported and
// the default prompt is used.
func promptString(name string) (string, bool) {
	ps, ok := os.LookupEnv(name)
	if !ok {
		return "", false
	}
	prompt, err := expandPrompt(ps, currentPromptInfo())
	if err != nil {
		fmt.Fprintf(os.Stderr, "gosh: %s: %s\n", name, sanitize(err.Error()))
		return "", false
	}
	return prompt, true
}