package accounts

import (
	"context"
	"errors"
	"fmt"
	"io/fs"
	"os"
	"path/filepath"
	"slices"
	"strings"
	"syscall"
	"time"
)

// The walk for setuid programs, resumable (novox/hq issue 361).
//
// The walk is the engine's own, in Go rather than find, so it can say how far it got: filepath.WalkDir visits
// every filesystem's tree depth first, each directory's entries in lexical order, so the walk's progress is one
// cursor per filesystem — the last path it visited — and everything before the cursor in that order is done.
// The progress is kept every CheckpointEvery beside the engine's state (SetuidCache), so an engine restarted,
// updated or rebooted midway carries on from its cursor rather than starting over: on the control-node the
// whole walk takes far longer than an engine runs between updates.
//
// **Why a resumed walk is sound.** Only root can make a setuid- or setgid-root file: the agent account can
// neither set the bit on root's file nor give its own file to root. A file that appears behind the cursor while
// the walk runs, or while it waits to be resumed, was made by root — and root's acts through the mesh start a
// new search (apply.MovesRoot), whose walk starts from nothing. A walk is dated by its start, and one resumed
// more than FreshFor after its start is thrown away and begun again.

// SetuidProgress is one walk, as far as it got.
type SetuidProgress struct {
	// ID names the walk; Started is when it began, on the cache's clock: the search is dated by it.
	ID      string    `json:"id"`
	Started time.Time `json:"started"`
	// Mounts are the filesystems it walks, root first; Done those walked to their end.
	Mounts []string `json:"mounts"`
	Done   []string `json:"done"`
	// Cursor is, for a filesystem being walked, the last path visited.
	Cursor map[string]string `json:"cursor"`
	// Found is every setuid- or setgid-root regular file found so far, before the package manager is asked.
	Found []string `json:"found"`
}

// CheckpointEvery is how often a walk's progress is kept.
const CheckpointEvery = time.Minute

// prunedPaths are never walked: the kernel's own trees, and the container runtimes' layers.
var prunedPaths = []string{"/proc", "/sys", "/var/lib/docker", "/var/lib/containers"}

// walker walks for setuid programs. rootUID and rootGID are the owner and group a match has: root's (0), but
// for a test.
type walker struct {
	rootUID, rootGID uint32
	every            time.Duration
}

// walk carries p on to its end, calling checkpoint every w.every and after each filesystem. It ends early only
// with an error — a directory it could not read, or ctx ended — which is an incomplete walk, never "none".
// A file or directory that vanished while it walked is not an error: what still exists was walked.
func (w walker) walk(ctx context.Context, p *SetuidProgress, checkpoint func()) error {
	if p.Cursor == nil {
		p.Cursor = map[string]string{}
	}
	last := time.Now()
	for _, m := range p.Mounts {
		if slices.Contains(p.Done, m) {
			continue
		}
		if err := w.walkMount(ctx, p, m, checkpoint, &last); err != nil {
			return err
		}
		p.Done = append(p.Done, m)
		delete(p.Cursor, m)
		checkpoint()
		last = time.Now()
	}
	slices.Sort(p.Found)
	p.Found = slices.Compact(p.Found)
	return nil
}

func (w walker) walkMount(ctx context.Context, p *SetuidProgress, mount string, checkpoint func(), last *time.Time) error {
	top, err := os.Lstat(mount)
	if err != nil {
		return fmt.Errorf("the filesystem at %s could not be read: %w", mount, err)
	}
	dev, ok := deviceOf(top)
	if !ok || !top.IsDir() {
		return fmt.Errorf("the filesystem at %s is not a directory whose device could be read", mount)
	}
	var cursor []string
	if c := p.Cursor[mount]; c != "" {
		cursor = components(c)
	}
	return filepath.WalkDir(mount, func(path string, d fs.DirEntry, err error) error {
		if ctxErr := ctx.Err(); ctxErr != nil {
			return ctxErr
		}
		if err != nil {
			if errors.Is(err, fs.ErrNotExist) {
				return nil // vanished while it walked
			}
			return fmt.Errorf("the walk for setuid programs could not read %s: %w", path, err)
		}
		at := components(path)
		if cursor != nil && walkOrder(at, cursor) < 0 && !isPrefix(at, cursor) {
			// Before the cursor and not on its way: visited by the walk this resumes.
			if d.IsDir() {
				return filepath.SkipDir
			}
			return nil
		}
		switch {
		case d.IsDir():
			if path != mount && slices.Contains(prunedPaths, path) {
				return filepath.SkipDir
			}
			info, err := os.Lstat(path)
			if errors.Is(err, fs.ErrNotExist) {
				return filepath.SkipDir
			}
			if err != nil {
				return fmt.Errorf("the walk for setuid programs could not read %s: %w", path, err)
			}
			if d, ok := deviceOf(info); !ok || d != dev {
				return filepath.SkipDir // another filesystem, walked as its own or not at all (find's -xdev)
			}
		case d.Type().IsRegular():
			info, err := d.Info() // the entry's own, never through a link
			if errors.Is(err, fs.ErrNotExist) {
				return nil
			}
			if err != nil {
				return fmt.Errorf("the walk for setuid programs could not read %s: %w", path, err)
			}
			if w.setuidRoot(info) {
				p.Found = append(p.Found, path)
			}
		}
		// The cursor only moves on: the directories on the way back to it are visited again, and are before it.
		if cursor == nil || walkOrder(at, cursor) > 0 {
			p.Cursor[mount], cursor = path, at
		}
		if time.Since(*last) >= w.every {
			checkpoint()
			*last = time.Now()
		}
		return nil
	})
}

// setuidRoot is a regular file setuid and owned by root, or setgid to root's group.
func (w walker) setuidRoot(info fs.FileInfo) bool {
	st, ok := info.Sys().(*syscall.Stat_t)
	if !ok || !info.Mode().IsRegular() {
		return false
	}
	return uint32(st.Uid) == w.rootUID && info.Mode()&fs.ModeSetuid != 0 ||
		uint32(st.Gid) == w.rootGID && info.Mode()&fs.ModeSetgid != 0
}

func deviceOf(info fs.FileInfo) (uint64, bool) {
	st, ok := info.Sys().(*syscall.Stat_t)
	if !ok {
		return 0, false
	}
	return uint64(st.Dev), true
}

// components are a path's names, in order; "/" has none.
func components(path string) []string {
	path = strings.Trim(filepath.Clean(path), "/")
	if path == "" {
		return nil
	}
	return strings.Split(path, "/")
}

// walkOrder compares two paths in the order WalkDir visits them: name by name, each in lexical order, a
// directory before what is in it.
func walkOrder(a, b []string) int {
	for i := 0; i < len(a) && i < len(b); i++ {
		if c := strings.Compare(a[i], b[i]); c != 0 {
			return c
		}
	}
	return len(a) - len(b)
}

// isPrefix says a is b or a directory b is in.
func isPrefix(a, b []string) bool {
	return len(a) <= len(b) && slices.Equal(a, b[:len(a)])
}