package main

import (
	"bufio"
	"bytes"
	"context"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"syscall"
	"time"
)

// CommandTimeout bounds every command a tool runs: a bus that hangs must cost a tool call twenty
// seconds, never the runtime's thirty.
const CommandTimeout = 20 * time.Second

// ReadCap is the most of one command's output the module reads. An introspection document of the
// service manager is about 100 KiB; nothing the module asks is near a mebibyte.
const ReadCap = 1024 * 1024

// AnswerCap is the most entries a tool answers in one list (names, messages, denials).
const AnswerCap = 500

// Runner runs one command and answers its standard output. Injected, so every tool is tested against
// recorded answers rather than this machine's bus.
type Runner func(ctx context.Context, name string, args ...string) (string, error)

// Streamer runs one command for at most the context's time and hands each line of its output to
// line, which says whether it wants more. The end of the time is the normal end, not an error.
// Injected, so dbus_monitor is tested without a bus.
type Streamer func(ctx context.Context, line func(string) bool, name string, args ...string) error

// ExecRunner runs a command, bounded by CommandTimeout. A failure carries what it said on stderr.
func ExecRunner(ctx context.Context, name string, args ...string) (string, error) {
	ctx, cancel := context.WithTimeout(ctx, CommandTimeout)
	defer cancel()
	cmd := exec.CommandContext(ctx, name, args...)
	cmd.Env = append(os.Environ(), "LC_ALL=C", "SYSTEMD_PAGER=", "SYSTEMD_COLORS=0")
	var stdout, stderr bytes.Buffer
	cmd.Stdout, cmd.Stderr = &stdout, &stderr
	err := cmd.Run()
	out := capped(stdout.String(), ReadCap)
	if ctx.Err() == context.DeadlineExceeded {
		return out, fmt.Errorf("%s did not answer within %s", name, CommandTimeout)
	}
	if err != nil {
		said := strings.TrimSpace(stderr.String())
		if said == "" {
			said = strings.TrimSpace(stdout.String())
		}
		return out, fmt.Errorf("%s %s: %w: %s", name, strings.Join(args, " "), err, capped(said, 2048))
	}
	return out, nil
}

// ExecStreamer runs a command until the context ends, line by line. A line longer than ReadCap is
// skipped, never held: a monitored message's line carries its body, which the module never keeps.
func ExecStreamer(ctx context.Context, line func(string) bool, name string, args ...string) error {
	cmd := exec.CommandContext(ctx, name, args...)
	cmd.Env = append(os.Environ(), "LC_ALL=C", "SYSTEMD_PAGER=", "SYSTEMD_COLORS=0")
	// A terminate, which sudo passes on to what it runs; a kill would leave a root monitor behind.
	cmd.Cancel = func() error { return cmd.Process.Signal(syscall.SIGTERM) }
	cmd.WaitDelay = 2 * time.Second
	var stderr bytes.Buffer
	cmd.Stderr = &stderr
	out, err := cmd.StdoutPipe()
	if err != nil {
		return err
	}
	if err := cmd.Start(); err != nil {
		return err
	}
	r := bufio.NewReaderSize(out, 64*1024)
	var cur []byte
	tooLong := false
	for {
		chunk, isPrefix, err := r.ReadLine()
		if err != nil {
			break
		}
		if !tooLong {
			cur = append(cur, chunk...)
			if len(cur) > ReadCap {
				cur, tooLong = cur[:0], true
			}
		}
		if isPrefix {
			continue
		}
		if !tooLong && !line(string(cur)) {
			break
		}
		cur, tooLong = cur[:0], false
	}
	_ = cmd.Process.Signal(syscall.SIGTERM)
	werr := cmd.Wait()
	if ctx.Err() != nil {
		return nil // the time ran out: the normal end of a bounded watch
	}
	if werr != nil && stderr.Len() > 0 {
		return fmt.Errorf("%s: %w: %s", name, werr, capped(strings.TrimSpace(stderr.String()), 2048))
	}
	return nil
}

func capped(s string, n int) string {
	if len(s) <= n {
		return s
	}
	return s[:n] + "\n… (cut)"
}

// Machine is what the module reads and acts on: a filesystem root (the real one, or a test's tree),
// a way to run commands, a way to stream one, and the environment the runtime gave it.
type Machine struct {
	Root   string
	Run    Runner
	Stream Streamer
	Env    func(string) string
	UID    int
	Now    func() time.Time
}

// Here is the machine this process runs on.
func Here() *Machine {
	return &Machine{Root: "/", Run: ExecRunner, Stream: ExecStreamer, Env: os.Getenv, UID: os.Getuid(), Now: time.Now}
}

func (m *Machine) path(p string) string { return filepath.Join(m.Root, p) }

func (m *Machine) read(p string) string {
	b, err := os.ReadFile(m.path(p))
	if err != nil {
		return ""
	}
	return strings.TrimSpace(string(b))
}

func (m *Machine) exists(p string) bool {
	_, err := os.Stat(m.path(p))
	return err == nil
}

func (m *Machine) glob(pattern string) []string {
	got, _ := filepath.Glob(m.path(pattern))
	out := make([]string, 0, len(got))
	for _, g := range got {
		rel, err := filepath.Rel(m.Root, g)
		if err != nil {
			continue
		}
		out = append(out, "/"+rel)
	}
	return out
}

// privileged runs a command as root without asking for a password (sudo -n), as the other modules'
// tools do: the runtime runs as the operator's account, and watching the system bus is root's.
func (m *Machine) privileged(ctx context.Context, name string, args ...string) (string, error) {
	return m.Run(ctx, "sudo", append([]string{"-n", name}, args...)...)
}

// Buses are the two buses a tool can be pointed at.
var Buses = []string{"system", "session"}

// ErrNoSession is the answer on a machine where this account has no session bus: the servers.
var ErrNoSession = errors.New("no session bus for this account on this machine")

// SessionAddress is the account's session bus: the runtime's DBUS_SESSION_BUS_ADDRESS, else the
// user manager's socket under XDG_RUNTIME_DIR or /run/user/<uid>. Only a socket that exists counts.
func (m *Machine) SessionAddress() (string, error) {
	if a := m.Env("DBUS_SESSION_BUS_ADDRESS"); strings.HasPrefix(a, "unix:path=") {
		p := strings.TrimPrefix(a, "unix:path=")
		if i := strings.IndexByte(p, ','); i >= 0 {
			p = p[:i]
		}
		if m.exists(p) {
			return a, nil
		}
	} else if a != "" {
		return a, nil // an abstract or other address: taken as given
	}
	dir := m.Env("XDG_RUNTIME_DIR")
	if dir == "" {
		dir = "/run/user/" + strconv.Itoa(m.UID)
	}
	if p := dir + "/bus"; m.exists(p) {
		return "unix:path=" + p, nil
	}
	return "", fmt.Errorf("%w (no socket at %s/bus)", ErrNoSession, dir)
}

// busArgs are busctl's words for one bus.
func (m *Machine) busArgs(bus string) ([]string, error) {
	switch bus {
	case "", "system":
		return []string{"--system"}, nil
	case "session":
		a, err := m.SessionAddress()
		if err != nil {
			return nil, err
		}
		return []string{"--address=" + a}, nil
	}
	return nil, fmt.Errorf("bus is system or session, not %q", bus)
}

// UnitOf is the systemd unit a process runs in, from its cgroup: readable for any process.
func (m *Machine) UnitOf(pid uint32) string {
	if pid == 0 {
		return ""
	}
	return UnitFromCgroup(m.read(fmt.Sprintf("/proc/%d/cgroup", pid)))
}

// ProcessName is a process's command name.
func (m *Machine) ProcessName(pid uint32) string {
	if pid == 0 {
		return ""
	}
	return m.read(fmt.Sprintf("/proc/%d/comm", pid))
}

// UnitFromCgroup is the innermost service, socket or scope in a cgroup v2 path.
func UnitFromCgroup(cgroup string) string {
	for _, line := range strings.Split(cgroup, "\n") {
		if !strings.HasPrefix(line, "0::") {
			continue
		}
		parts := strings.Split(strings.TrimPrefix(line, "0::"), "/")
		for i := len(parts) - 1; i >= 0; i-- {
			p := parts[i]
			if strings.HasSuffix(p, ".service") || strings.HasSuffix(p, ".scope") {
				return p
			}
		}
	}
	return ""
}

// BootID is this boot's id: a bus that changed across a boot did not restart, the machine did.
func (m *Machine) BootID() string { return m.read("/proc/sys/kernel/random/boot_id") }