bee/audit/internal/app/blackbox.go

package app

import (
	"bytes"
	"context"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io/fs"
	"os"
	"os/exec"
	"path/filepath"
	"sort"
	"strings"
	"sync"
	"time"

	"bee/audit/internal/platform"
)

const (
	blackboxMarkerName        = ".bee-blackbox"
	blackboxDiscoverInterval  = 2 * time.Second
	blackboxMinFlushPeriod    = 1 * time.Second
	blackboxMaxFlushPeriod    = 30 * time.Second
	blackboxRecoveryFastCount = 5
)

var DefaultBlackboxStatePath = DefaultExportDir + "/blackbox-state.json"

var (
	blackboxExecCommand = exec.Command
	blackboxNow         = func() time.Time { return time.Now().UTC() }
)

type BlackboxMarker struct {
	Version      int    `json:"version"`
	EnrollmentID string `json:"enrollment_id"`
	CreatedAtUTC string `json:"created_at_utc"`
	Host         string `json:"host,omitempty"`
}

type BlackboxTargetStatus struct {
	EnrollmentID      string                   `json:"enrollment_id"`
	Device            string                   `json:"device"`
	FS                platform.RemovableTarget `json:"fs"`
	BootFolder        string                   `json:"boot_folder"`
	Status            string                   `json:"status"`
	LastSyncAtUTC     string                   `json:"last_sync_at_utc,omitempty"`
	LastCycleDuration string                   `json:"last_cycle_duration,omitempty"`
	FlushPeriod       string                   `json:"flush_period"`
	LastError         string                   `json:"last_error,omitempty"`
	Mountpoint        string                   `json:"mountpoint,omitempty"`
}

type BlackboxState struct {
	Status           string                 `json:"status"`
	BootStartedAtUTC string                 `json:"boot_started_at_utc"`
	BootFolder       string                 `json:"boot_folder"`
	UpdatedAtUTC     string                 `json:"updated_at_utc"`
	Targets          []BlackboxTargetStatus `json:"targets"`
}

type blackboxRuntime struct {
	exportDir   string
	statePath   string
	system      *platform.System
	bootStarted time.Time
	bootFolder  string

	mu      sync.Mutex
	workers map[string]*blackboxWorker
}

type discoveredBlackboxTarget struct {
	marker       BlackboxMarker
	target       platform.RemovableTarget
	seenMount    string
	mountedByBee bool
}

type blackboxWorker struct {
	runtime      *blackboxRuntime
	enrollmentID string

	mu           sync.Mutex
	target       platform.RemovableTarget
	marker       BlackboxMarker
	mountpoint   string
	mountedByBee bool
	status       string
	lastSyncAt   time.Time
	lastDuration time.Duration
	flushPeriod  time.Duration
	lastError    string
	fastCycles   int
	stopCh       chan struct{}
	stoppedCh    chan struct{}
}

func RunBlackbox(ctx context.Context, exportDir, statePath string, system *platform.System) error {
	exportDir = strings.TrimSpace(exportDir)
	if exportDir == "" {
		exportDir = DefaultExportDir
	}
	statePath = strings.TrimSpace(statePath)
	if statePath == "" {
		statePath = DefaultBlackboxStatePath
	}
	if system == nil {
		system = platform.New()
	}
	bootStarted, err := bootStartedAtUTC()
	if err != nil {
		bootStarted = blackboxNow()
	}
	rt := &blackboxRuntime{
		exportDir:   exportDir,
		statePath:   statePath,
		system:      system,
		bootStarted: bootStarted,
		bootFolder:  SupportBundleBaseName(bootStarted),
		workers:     make(map[string]*blackboxWorker),
	}
	_ = os.MkdirAll(filepath.Dir(statePath), 0755)
	rt.persistState()
	ticker := time.NewTicker(blackboxDiscoverInterval)
	defer ticker.Stop()
	for {
		rt.reconcile()
		select {
		case <-ctx.Done():
			rt.stopAll()
			return ctx.Err()
		case <-ticker.C:
		}
	}
}

func ReadBlackboxState(path string) (BlackboxState, error) {
	path = strings.TrimSpace(path)
	if path == "" {
		path = DefaultBlackboxStatePath
	}
	raw, err := os.ReadFile(path)
	if err != nil {
		return BlackboxState{}, err
	}
	var state BlackboxState
	if err := json.Unmarshal(raw, &state); err != nil {
		return BlackboxState{}, err
	}
	return state, nil
}

func EnableBlackboxTarget(target platform.RemovableTarget) (BlackboxMarker, error) {
	target = sanitizeRemovableTarget(target)
	if target.Device == "" {
		return BlackboxMarker{}, fmt.Errorf("device is required")
	}
	mountpoint, mountedByBee, err := ensureMountedTarget(target, "marker")
	if err != nil {
		return BlackboxMarker{}, err
	}
	defer func() {
		if mountedByBee {
			_ = unmountTarget(mountpoint)
		}
	}()

	marker, _, err := readBlackboxMarker(mountpoint)
	if err != nil && !errors.Is(err, os.ErrNotExist) {
		return BlackboxMarker{}, err
	}
	if marker.EnrollmentID == "" {
		marker = BlackboxMarker{
			Version:      1,
			EnrollmentID: newBlackboxEnrollmentID(),
			CreatedAtUTC: blackboxNow().Format(time.RFC3339),
			Host:         hostnameOr("unknown"),
		}
	}
	if err := writeBlackboxMarker(mountpoint, marker); err != nil {
		return BlackboxMarker{}, err
	}
	return marker, nil
}

func DisableBlackboxTarget(device, enrollmentID string) error {
	device = strings.TrimSpace(device)
	enrollmentID = strings.TrimSpace(enrollmentID)
	if device == "" && enrollmentID == "" {
		return fmt.Errorf("device or enrollment_id is required")
	}
	system := platform.New()
	targets, err := system.ListRemovableTargets()
	if err != nil {
		return err
	}
	for _, target := range targets {
		target = sanitizeRemovableTarget(target)
		mountpoint, mountedByBee, mountErr := ensureMountedTarget(target, "marker")
		if mountErr != nil {
			continue
		}
		remove := false
		marker, _, err := readBlackboxMarker(mountpoint)
		if err == nil {
			if enrollmentID != "" && marker.EnrollmentID == enrollmentID {
				remove = true
			}
			if device != "" && target.Device == device {
				remove = true
			}
		}
		if remove {
			err = os.Remove(filepath.Join(mountpoint, blackboxMarkerName))
		}
		if mountedByBee {
			_ = unmountTarget(mountpoint)
		}
		if remove {
			return err
		}
	}
	return os.ErrNotExist
}

func (rt *blackboxRuntime) reconcile() {
	discovered, _ := rt.discoverMarkedTargets()

	rt.mu.Lock()
	defer rt.mu.Unlock()

	seen := make(map[string]struct{}, len(discovered))
	for _, found := range discovered {
		seen[found.marker.EnrollmentID] = struct{}{}
		worker, ok := rt.workers[found.marker.EnrollmentID]
		if !ok {
			worker = newBlackboxWorker(rt, found)
			rt.workers[found.marker.EnrollmentID] = worker
			go worker.run()
			continue
		}
		worker.update(found)
	}
	for id, worker := range rt.workers {
		if _, ok := seen[id]; ok {
			continue
		}
		worker.stop()
		delete(rt.workers, id)
	}
	rt.persistStateLocked()
}

func (rt *blackboxRuntime) stopAll() {
	rt.mu.Lock()
	workers := make([]*blackboxWorker, 0, len(rt.workers))
	for _, worker := range rt.workers {
		workers = append(workers, worker)
	}
	rt.workers = map[string]*blackboxWorker{}
	rt.persistStateLocked()
	rt.mu.Unlock()
	for _, worker := range workers {
		worker.stop()
	}
}

func (rt *blackboxRuntime) discoverMarkedTargets() ([]discoveredBlackboxTarget, error) {
	targets, err := rt.system.ListRemovableTargets()
	if err != nil {
		return nil, err
	}
	var out []discoveredBlackboxTarget
	for _, rawTarget := range targets {
		target := sanitizeRemovableTarget(rawTarget)
		if target.Device == "" {
			continue
		}
		mountpoint, mountedByBee, err := ensureMountedTarget(target, "probe")
		if err != nil {
			continue
		}
		marker, ok, err := readBlackboxMarker(mountpoint)
		if mountedByBee && !ok {
			_ = unmountTarget(mountpoint)
		}
		if err != nil || !ok || marker.EnrollmentID == "" {
			continue
		}
		if mountedByBee {
			_ = unmountTarget(mountpoint)
		}
		out = append(out, discoveredBlackboxTarget{
			marker:       marker,
			target:       target,
			seenMount:    mountpoint,
			mountedByBee: mountedByBee,
		})
	}
	sort.Slice(out, func(i, j int) bool {
		return out[i].marker.EnrollmentID < out[j].marker.EnrollmentID
	})
	return out, nil
}

func newBlackboxWorker(rt *blackboxRuntime, found discoveredBlackboxTarget) *blackboxWorker {
	return &blackboxWorker{
		runtime:      rt,
		enrollmentID: found.marker.EnrollmentID,
		target:       found.target,
		marker:       found.marker,
		flushPeriod:  blackboxMinFlushPeriod,
		status:       "running",
		stopCh:       make(chan struct{}),
		stoppedCh:    make(chan struct{}),
	}
}

func (w *blackboxWorker) run() {
	defer close(w.stoppedCh)
	for {
		start := time.Now()
		err := w.syncCycle()
		duration := time.Since(start)
		w.finishCycle(duration, err)

		wait := w.currentFlushPeriod()
		timer := time.NewTimer(wait)
		select {
		case <-w.stopCh:
			timer.Stop()
			w.cleanup()
			return
		case <-timer.C:
		}
	}
}

func (w *blackboxWorker) update(found discoveredBlackboxTarget) {
	w.mu.Lock()
	defer w.mu.Unlock()
	w.target = found.target
	w.marker = found.marker
}

func (w *blackboxWorker) stop() {
	select {
	case <-w.stopCh:
	default:
		close(w.stopCh)
	}
	<-w.stoppedCh
}

func (w *blackboxWorker) currentFlushPeriod() time.Duration {
	w.mu.Lock()
	defer w.mu.Unlock()
	return w.flushPeriod
}

func (w *blackboxWorker) finishCycle(duration time.Duration, err error) {
	w.mu.Lock()
	defer w.mu.Unlock()
	w.lastDuration = duration
	if err != nil {
		w.status = "degraded"
		w.lastError = err.Error()
		w.fastCycles = 0
		w.flushPeriod = adjustFlushPeriod(w.flushPeriod, duration, false, 0)
	} else {
		w.status = "running"
		w.lastSyncAt = blackboxNow()
		w.lastError = ""
		if duration <= w.flushPeriod/2 {
			w.fastCycles++
		} else {
			w.fastCycles = 0
		}
		w.flushPeriod = adjustFlushPeriod(w.flushPeriod, duration, true, w.fastCycles)
	}
	w.runtime.persistState()
}

func adjustFlushPeriod(current, duration time.Duration, success bool, fastCycles int) time.Duration {
	if current <= 0 {
		current = blackboxMinFlushPeriod
	}
	if duration <= 0 {
		duration = current
	}
	next := current
	if duration > current {
		growA := time.Duration(float64(current) * 1.25)
		growB := time.Duration(float64(duration) * 1.25)
		if growB > growA {
			next = growB
		} else {
			next = growA
		}
	}
	if success && fastCycles >= blackboxRecoveryFastCount {
		next = time.Duration(float64(current) * 0.9)
	}
	if next < blackboxMinFlushPeriod {
		next = blackboxMinFlushPeriod
	}
	if next > blackboxMaxFlushPeriod {
		next = blackboxMaxFlushPeriod
	}
	return next
}

func (w *blackboxWorker) syncCycle() error {
	target, marker := w.snapshotTarget()
	mountpoint, mountedByBee, err := ensureMountedTarget(target, marker.EnrollmentID)
	if err != nil {
		return err
	}
	w.recordMountpoint(mountpoint, mountedByBee)

	root := filepath.Join(mountpoint, w.runtime.bootFolder)
	if err := os.MkdirAll(filepath.Join(root, "export"), 0755); err != nil {
		return err
	}
	if err := syncDirectoryTree(w.runtime.exportDir, filepath.Join(root, "export")); err != nil {
		return err
	}
	if err := w.captureSnapshots(root); err != nil {
		return err
	}
	return syncFilesystem(root)
}

func (w *blackboxWorker) cleanup() {
	w.mu.Lock()
	mountpoint := w.mountpoint
	mountedByBee := w.mountedByBee
	w.mu.Unlock()
	if mountedByBee && mountpoint != "" {
		_ = unmountTarget(mountpoint)
	}
}

func (w *blackboxWorker) snapshotTarget() (platform.RemovableTarget, BlackboxMarker) {
	w.mu.Lock()
	defer w.mu.Unlock()
	return w.target, w.marker
}

func (w *blackboxWorker) recordMountpoint(mountpoint string, mountedByBee bool) {
	w.mu.Lock()
	defer w.mu.Unlock()
	w.mountpoint = mountpoint
	w.mountedByBee = mountedByBee
}

func (w *blackboxWorker) captureSnapshots(root string) error {
	if err := captureCommandAtomic(filepath.Join(root, "systemd", "combined.journal.log"), "journalctl", "--no-pager", "--since", w.runtime.bootStarted.Format(time.RFC3339)); err != nil {
		return err
	}
	for _, svc := range supportBundleServices {
		if err := captureCommandAtomic(filepath.Join(root, "systemd", svc+".journal.log"), "journalctl", "--no-pager", "-u", svc, "--since", w.runtime.bootStarted.Format(time.RFC3339)); err != nil {
			return err
		}
		if err := captureCommandAtomic(filepath.Join(root, "systemd", svc+".status.txt"), "systemctl", "status", svc, "--no-pager"); err != nil {
			return err
		}
	}
	if err := captureCommandAtomic(filepath.Join(root, "system", "dmesg.txt"), "dmesg"); err != nil {
		return err
	}
	for _, item := range supportBundleOptionalFiles {
		if err := copyFileIfChanged(item.src, filepath.Join(root, item.name)); err != nil && !errors.Is(err, os.ErrNotExist) {
			return err
		}
	}
	return nil
}

func (rt *blackboxRuntime) persistState() {
	rt.mu.Lock()
	defer rt.mu.Unlock()
	rt.persistStateLocked()
}

func (rt *blackboxRuntime) persistStateLocked() {
	state := BlackboxState{
		Status:           "disabled",
		BootStartedAtUTC: rt.bootStarted.Format(time.RFC3339),
		BootFolder:       rt.bootFolder,
		UpdatedAtUTC:     blackboxNow().Format(time.RFC3339),
		Targets:          make([]BlackboxTargetStatus, 0, len(rt.workers)),
	}
	if len(rt.workers) > 0 {
		state.Status = "running"
	}
	for _, worker := range rt.workers {
		worker.mu.Lock()
		targetState := BlackboxTargetStatus{
			EnrollmentID: worker.enrollmentID,
			Device:       worker.target.Device,
			FS:           worker.target,
			BootFolder:   rt.bootFolder,
			Status:       worker.status,
			FlushPeriod:  worker.flushPeriod.String(),
			LastError:    worker.lastError,
			Mountpoint:   worker.mountpoint,
		}
		if !worker.lastSyncAt.IsZero() {
			targetState.LastSyncAtUTC = worker.lastSyncAt.Format(time.RFC3339)
		}
		if worker.lastDuration > 0 {
			targetState.LastCycleDuration = worker.lastDuration.String()
		}
		if worker.status == "degraded" {
			state.Status = "degraded"
		}
		worker.mu.Unlock()
		state.Targets = append(state.Targets, targetState)
	}
	sort.Slice(state.Targets, func(i, j int) bool {
		return state.Targets[i].EnrollmentID < state.Targets[j].EnrollmentID
	})
	_ = writeJSONAtomic(rt.statePath, state)
}

func bootStartedAtUTC() (time.Time, error) {
	raw, err := os.ReadFile("/proc/stat")
	if err != nil {
		return time.Time{}, err
	}
	for _, line := range strings.Split(string(raw), "\n") {
		line = strings.TrimSpace(line)
		if !strings.HasPrefix(line, "btime ") {
			continue
		}
		parts := strings.Fields(line)
		if len(parts) != 2 {
			break
		}
		sec, err := time.ParseDuration(parts[1] + "s")
		if err != nil {
			break
		}
		return time.Unix(int64(sec/time.Second), 0).UTC(), nil
	}
	return time.Time{}, fmt.Errorf("boot time not found")
}

func newBlackboxEnrollmentID() string {
	var buf [8]byte
	if _, err := rand.Read(buf[:]); err != nil {
		return fmt.Sprintf("bb-%d", time.Now().UnixNano())
	}
	return "bb-" + hex.EncodeToString(buf[:])
}

func sanitizeRemovableTarget(target platform.RemovableTarget) platform.RemovableTarget {
	target.Device = strings.TrimSpace(target.Device)
	target.FSType = strings.TrimSpace(target.FSType)
	target.Size = strings.TrimSpace(target.Size)
	target.Label = strings.TrimSpace(target.Label)
	target.Model = strings.TrimSpace(target.Model)
	target.Mountpoint = strings.TrimSpace(target.Mountpoint)
	return target
}

func ensureMountedTarget(target platform.RemovableTarget, suffix string) (mountpoint string, mountedByBee bool, retErr error) {
	target = sanitizeRemovableTarget(target)
	if target.Mountpoint != "" {
		if err := ensureWritableBlackboxMountpoint(target.Mountpoint); err == nil {
			return target.Mountpoint, false, nil
		}
	}
	mountpoint = filepath.Join("/tmp", "bee-blackbox-"+sanitizeFilename(suffix))
	if err := os.MkdirAll(mountpoint, 0755); err != nil {
		return "", false, err
	}
	if raw, err := blackboxExecCommand("mount", target.Device, mountpoint).CombinedOutput(); err != nil {
		return "", false, formatBlackboxMountTargetError(target, string(raw), err)
	}
	if err := ensureWritableBlackboxMountpoint(mountpoint); err != nil {
		_ = unmountTarget(mountpoint)
		return "", false, err
	}
	return mountpoint, true, nil
}

func unmountTarget(mountpoint string) error {
	_ = blackboxExecCommand("sync").Run()
	raw, err := blackboxExecCommand("umount", mountpoint).CombinedOutput()
	if err != nil {
		msg := strings.TrimSpace(string(raw))
		if msg == "" {
			return err
		}
		return fmt.Errorf("%s: %w", msg, err)
	}
	return nil
}

func readBlackboxMarker(mountpoint string) (BlackboxMarker, bool, error) {
	raw, err := os.ReadFile(filepath.Join(mountpoint, blackboxMarkerName))
	if err != nil {
		if errors.Is(err, os.ErrNotExist) {
			return BlackboxMarker{}, false, os.ErrNotExist
		}
		return BlackboxMarker{}, false, err
	}
	var marker BlackboxMarker
	if err := json.Unmarshal(raw, &marker); err != nil {
		return BlackboxMarker{}, false, err
	}
	return marker, true, nil
}

func writeBlackboxMarker(mountpoint string, marker BlackboxMarker) error {
	if marker.Version == 0 {
		marker.Version = 1
	}
	return writeJSONAtomic(filepath.Join(mountpoint, blackboxMarkerName), marker)
}

func syncDirectoryTree(srcDir, dstDir string) error {
	seen := make(map[string]struct{})
	err := filepath.WalkDir(srcDir, func(path string, d fs.DirEntry, err error) error {
		if err != nil {
			return err
		}
		rel, err := filepath.Rel(srcDir, path)
		if err != nil {
			return err
		}
		rel = filepath.Clean(rel)
		if rel == "." {
			seen["."] = struct{}{}
			return os.MkdirAll(dstDir, 0755)
		}
		seen[rel] = struct{}{}
		dstPath := filepath.Join(dstDir, rel)
		if d.IsDir() {
			info, err := d.Info()
			if err != nil {
				return err
			}
			return os.MkdirAll(dstPath, info.Mode().Perm())
		}
		return copyFileIfChanged(path, dstPath)
	})
	if err != nil {
		return err
	}
	return removeMissingPaths(dstDir, seen)
}

func removeMissingPaths(dstDir string, seen map[string]struct{}) error {
	return filepath.WalkDir(dstDir, func(path string, d fs.DirEntry, err error) error {
		if err != nil {
			return err
		}
		rel, err := filepath.Rel(dstDir, path)
		if err != nil {
			return err
		}
		rel = filepath.Clean(rel)
		if rel == "." {
			return nil
		}
		if _, ok := seen[rel]; ok {
			return nil
		}
		return os.RemoveAll(path)
	})
}

func copyFileIfChanged(src, dst string) error {
	info, err := os.Stat(src)
	if err != nil {
		return err
	}
	if info.IsDir() {
		return os.MkdirAll(dst, info.Mode().Perm())
	}
	srcData, err := os.ReadFile(src)
	if err != nil {
		return err
	}
	if dstData, err := os.ReadFile(dst); err == nil && bytes.Equal(dstData, srcData) {
		return nil
	}
	return writeFileAtomic(dst, srcData, info.Mode().Perm())
}

func captureCommandAtomic(dst string, name string, args ...string) error {
	raw, err := blackboxExecCommand(name, args...).CombinedOutput()
	if len(raw) == 0 {
		if err != nil {
			raw = []byte(err.Error() + "\n")
		} else {
			raw = []byte("no output\n")
		}
	}
	return writeFileAtomic(dst, raw, 0644)
}

func writeJSONAtomic(path string, v any) error {
	raw, err := json.MarshalIndent(v, "", "  ")
	if err != nil {
		return err
	}
	raw = append(raw, '\n')
	return writeFileAtomic(path, raw, 0644)
}

func writeFileAtomic(path string, data []byte, perm os.FileMode) error {
	if err := os.MkdirAll(filepath.Dir(path), 0755); err != nil {
		return err
	}
	if existing, err := os.ReadFile(path); err == nil && bytes.Equal(existing, data) {
		return nil
	}
	tmp := path + ".tmp"
	f, err := os.OpenFile(tmp, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, perm)
	if err != nil {
		return err
	}
	if _, err := f.Write(data); err != nil {
		_ = f.Close()
		return err
	}
	if err := f.Sync(); err != nil {
		_ = f.Close()
		return err
	}
	if err := f.Close(); err != nil {
		return err
	}
	if err := os.Rename(tmp, path); err != nil {
		return err
	}
	return syncFilesystem(filepath.Dir(path))
}

func syncFilesystem(path string) error {
	return blackboxExecCommand("sync").Run()
}

func ensureWritableBlackboxMountpoint(mountpoint string) error {
	probe, err := os.CreateTemp(mountpoint, ".bee-blackbox-write-test-*")
	if err != nil {
		return fmt.Errorf("target filesystem is not writable: %w", err)
	}
	name := probe.Name()
	if closeErr := probe.Close(); closeErr != nil {
		_ = os.Remove(name)
		return closeErr
	}
	if err := os.Remove(name); err != nil {
		return err
	}
	return nil
}

func formatBlackboxMountTargetError(target platform.RemovableTarget, raw string, err error) error {
	msg := strings.TrimSpace(raw)
	fstype := strings.ToLower(strings.TrimSpace(target.FSType))
	if fstype == "exfat" && strings.Contains(strings.ToLower(msg), "unknown filesystem type 'exfat'") {
		return fmt.Errorf("mount %s: exFAT support is missing in this ISO build: %w", target.Device, err)
	}
	if msg == "" {
		return err
	}
	return fmt.Errorf("%s: %w", msg, err)
}