bee/audit/internal/app/app.go

package app

import (
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"os"
	"strings"

	"bee/audit/internal/collector"
	"bee/audit/internal/platform"
	"bee/audit/internal/runtimeenv"
	"bee/audit/internal/schema"
)

var (
	DefaultExportDir                     = "/appdata/bee/export"
	DefaultAuditJSONPath                 = DefaultExportDir + "/bee-audit.json"
	DefaultAuditLogPath                  = DefaultExportDir + "/bee-audit.log"
	DefaultWebLogPath                    = DefaultExportDir + "/bee-web.log"
	DefaultNetworkLogPath                = DefaultExportDir + "/bee-network.log"
	DefaultNvidiaLogPath                 = DefaultExportDir + "/bee-nvidia.log"
	DefaultSSHLogPath                    = DefaultExportDir + "/bee-sshsetup.log"
	DefaultRuntimeJSONPath               = DefaultExportDir + "/runtime-health.json"
	DefaultRuntimeLogPath                = DefaultExportDir + "/runtime-health.log"
	DefaultTechDumpDir                   = DefaultExportDir + "/techdump"
	DefaultSATBaseDir                    = DefaultExportDir + "/bee-sat"
	DefaultBeeBenchBaseDir               = DefaultExportDir + "/bee-bench"
	DefaultBeeBenchAutotuneDir           = DefaultBeeBenchBaseDir + "/autotune"
	DefaultBeeBenchPerfDir               = DefaultBeeBenchBaseDir + "/perf"
	DefaultBeeBenchPowerDir              = DefaultBeeBenchBaseDir + "/power"
	DefaultBeeBenchPowerSourceConfigPath = DefaultBeeBenchBaseDir + "/power-source-autotune.json"
)

type App struct {
	network   networkManager
	services  serviceManager
	exports   exportManager
	tools     toolManager
	sat       satRunner
	runtime   runtimeChecker
	installer installer
	// StatusDB is the unified component health store (nil if unavailable).
	StatusDB *ComponentStatusDB
}

type ActionResult struct {
	Title string
	Body  string
}

type networkManager interface {
	ListInterfaces() ([]platform.InterfaceInfo, error)
	DefaultRoute() string
	DHCPOne(iface string) (string, error)
	DHCPAll() (string, error)
	SetStaticIPv4(cfg platform.StaticIPv4Config) (string, error)
	SetInterfaceState(iface string, up bool) error
	GetInterfaceState(iface string) (bool, error)
	CaptureNetworkSnapshot() (platform.NetworkSnapshot, error)
	RestoreNetworkSnapshot(snapshot platform.NetworkSnapshot) error
}

type serviceManager interface {
	ListBeeServices() ([]string, error)
	ServiceState(name string) string
	ServiceStatus(name string) (string, error)
	ServiceDo(name string, action platform.ServiceAction) (string, error)
}

type exportManager interface {
	ListRemovableTargets() ([]platform.RemovableTarget, error)
	ExportFileToTarget(src string, target platform.RemovableTarget) (string, error)
	ListScenarioFilesOnRemovableMedia() ([]platform.ScenarioFileOnRemovableMedia, error)
	ReadScenarioFromRemovableMedia(name string) ([]byte, error)
	ListAvailableScenarios() ([]platform.ScenarioFileOnRemovableMedia, error)
	ReadScenario(name string) ([]byte, error)
}

type toolManager interface {
	TailFile(path string, lines int) string
	CheckTools(names []string) []platform.ToolStatus
}

type installer interface {
	ListInstallDisks() ([]platform.InstallDisk, error)
	InstallToDisk(ctx context.Context, device string, logFile string) error
	IsLiveMediaInRAM() bool
	LiveBootSource() platform.LiveBootSource
	LiveMediaRAMState() platform.LiveMediaRAMState
	RunInstallToRAM(ctx context.Context, logFunc func(string)) error
}

type GPUPresenceResult struct {
	Nvidia bool
	AMD    bool
	// NvidiaInitializing / AMDInitializing report a PCI device of that vendor
	// when DetectGPUVendor did not report the vendor as operational.
	NvidiaInitializing bool
	AMDInitializing    bool
}

// DetectGPUPresence combines the existing operational vendor detection with a
// PCI display-class scan. The latter distinguishes absent hardware from a
// detected PCI function whose runtime is not operational yet.
func (a *App) DetectGPUPresence() GPUPresenceResult {
	vendor := a.sat.DetectGPUVendor()
	res := GPUPresenceResult{
		Nvidia: vendor == "nvidia",
		AMD:    vendor == "amd",
	}
	physNvidia, physAMD := a.sat.PhysicalGPUVendors()
	res.NvidiaInitializing = physNvidia && !res.Nvidia
	res.AMDInitializing = physAMD && !res.AMD
	return res
}

// RuntimeHealthNow collects a fresh runtime-health snapshot (driver / CUDA
// readiness, GSP state). Unlike ReadRuntimeHealth it does not read the
// boot-time JSON; callers that poll for the GPU stack to come up need
// current data.
func (a *App) RuntimeHealthNow() (schema.RuntimeHealth, error) {
	return a.runtime.CollectRuntimeHealth(DefaultExportDir)
}

// TPMPresent reports whether this host has a TPM the checks can talk to.
func (a *App) TPMPresent() bool {
	return a.sat.TPMPresent()
}

func (a *App) IsLiveMediaInRAM() bool {
	return a.installer.IsLiveMediaInRAM()
}

func (a *App) LiveBootSource() platform.LiveBootSource {
	return a.installer.LiveBootSource()
}

func (a *App) LiveMediaRAMState() platform.LiveMediaRAMState {
	return a.installer.LiveMediaRAMState()
}

func (a *App) RunInstallToRAM(ctx context.Context, logFunc func(string)) error {
	return a.installer.RunInstallToRAM(ctx, logFunc)
}

type satRunner interface {
	RunNvidiaAcceptancePack(baseDir string, logFunc func(string)) (string, error)
	RunNvidiaAcceptancePackWithOptions(ctx context.Context, baseDir string, diagLevel int, gpuIndices []int, logFunc func(string)) (string, error)
	RunNvidiaTargetedStressValidatePack(ctx context.Context, baseDir string, durationSec int, gpuIndices []int, logFunc func(string)) (string, error)
	RunNvidiaBenchmark(ctx context.Context, baseDir string, opts platform.NvidiaBenchmarkOptions, logFunc func(string)) (string, error)
	RunNvidiaPowerBench(ctx context.Context, baseDir string, opts platform.NvidiaBenchmarkOptions, logFunc func(string)) (string, error)
	RunNvidiaPowerSourceAutotune(ctx context.Context, baseDir string, opts platform.NvidiaBenchmarkOptions, benchmarkKind string, logFunc func(string)) (string, error)
	RunNvidiaOfficialComputePack(ctx context.Context, baseDir string, durationSec int, gpuIndices []int, staggerSec int, logFunc func(string)) (string, error)
	RunNvidiaTargetedPowerPack(ctx context.Context, baseDir string, durationSec int, gpuIndices []int, logFunc func(string)) (string, error)
	RunNvidiaPulseTestPack(ctx context.Context, baseDir string, durationSec int, gpuIndices []int, logFunc func(string)) (string, error)
	RunNvidiaBandwidthPack(ctx context.Context, baseDir string, gpuIndices []int, fullMatrix bool, logFunc func(string)) (string, error)
	RunNvidiaStressPack(ctx context.Context, baseDir string, opts platform.NvidiaStressOptions, logFunc func(string)) (string, error)
	ListNvidiaGPUStatuses() ([]platform.NvidiaGPUStatus, error)
	ResetNvidiaGPU(index int) (string, error)
	RunMemoryAcceptancePack(ctx context.Context, baseDir string, sizeMB, passes int, logFunc func(string)) (string, error)
	RunStorageAcceptancePack(ctx context.Context, baseDir string, extended bool, logFunc func(string)) (string, error)
	RunTPMValidationPack(ctx context.Context, baseDir string, logFunc func(string)) (string, error)
	TPMPresent() bool
	RunNvidiaConfigCheckPack(ctx context.Context, baseDir string, logFunc func(string)) (string, error)
	RunPCIeLinkCheckPack(ctx context.Context, baseDir string, logFunc func(string)) (string, error)
	RunNvidiaPCIeBandwidthPack(ctx context.Context, baseDir string, gpuIndices []int, logFunc func(string)) (string, error)
	RunCPUAcceptancePack(ctx context.Context, baseDir string, durationSec int, logFunc func(string)) (string, error)
	ListNvidiaGPUs() ([]platform.NvidiaGPU, error)
	ListNvidiaGPUSettings() ([]platform.NvidiaGPUSetting, error)
	SetNvidiaGPUECC(index int, enabled bool) (string, error)
	SetNvidiaGPUMIG(index int, enabled bool) (string, error)
	SetNvidiaGPUCCMode(index int, enabled bool) (string, error)
	SetNvidiaGPUPowerLimit(index int, watts float64) (string, error)
	ResetNvidiaGPUDefaults() (string, error)
	DetectGPUVendor() string
	PhysicalGPUVendors() (nvidia bool, amd bool)
	ListAMDGPUs() ([]platform.AMDGPUInfo, error)
	RunAMDAcceptancePack(ctx context.Context, baseDir string, logFunc func(string)) (string, error)
	RunAMDMemIntegrityPack(ctx context.Context, baseDir string, logFunc func(string)) (string, error)
	RunAMDMemBandwidthPack(ctx context.Context, baseDir string, logFunc func(string)) (string, error)
	RunAMDStressPack(ctx context.Context, baseDir string, durationSec int, logFunc func(string)) (string, error)
	RunMemoryStressPack(ctx context.Context, baseDir string, durationSec int, logFunc func(string)) (string, error)
	RunSATStressPack(ctx context.Context, baseDir string, durationSec int, logFunc func(string)) (string, error)
	RunFanStressTest(ctx context.Context, baseDir string, opts platform.FanStressOptions) (string, error)
	RunPlatformStress(ctx context.Context, baseDir string, opts platform.PlatformStressOptions, logFunc func(string)) (string, error)
	RunNCCLTests(ctx context.Context, baseDir string, gpuIndices []int, logFunc func(string)) (string, error)
	RunScenario(ctx context.Context, baseDir string, spec platform.ScenarioSpec, logFunc func(string)) (string, error)
}

type runtimeChecker interface {
	CollectRuntimeHealth(exportDir string) (schema.RuntimeHealth, error)
	CaptureTechnicalDump(baseDir string) error
}

func New(sys *platform.System) *App {
	a := &App{
		network:   sys,
		services:  sys,
		exports:   sys,
		tools:     sys,
		sat:       sys,
		runtime:   sys,
		installer: sys,
	}
	if db, err := OpenComponentStatusDB(DefaultExportDir + "/component-status.json"); err == nil {
		a.StatusDB = db
	}
	// Let any running blackbox worker know promptly when a SAT job finishes,
	// instead of waiting out its own adaptive flush period — see
	// requestBlackboxSync/blackboxWorker.run in blackbox.go.
	platform.SetJobBoundaryHook(func(string) {
		requestBlackboxSync(DefaultExportDir)
	})
	// For the actual load step of a diagnostic (nvbandwidth, memtester,
	// stress-ng, dcgmi diag...) — as opposed to the cheap discovery/inventory
	// steps around it — block until blackbox has actually copied the
	// evidence that the load is about to start onto removable media, and
	// again once it finishes. Best-effort: a stuck blackbox target logs a
	// warning (see runSyncBracketHook in sat.go) rather than blocking the
	// diagnostic itself.
	platform.SetSyncBracketHook(func(jobName, phase string) error {
		return requestBlackboxSyncAndWait(DefaultExportDir, DefaultBlackboxStatePath, blackboxSyncBracketTimeout)
	})
	return a
}

// ApplySATOverlay parses a raw audit JSON, overlays the latest SAT results,
// and returns the updated JSON. Used by the web UI to serve always-fresh status.
func ApplySATOverlay(auditJSON []byte) ([]byte, error) {
	snap, err := readAuditSnapshot(auditJSON)
	if err != nil {
		return nil, err
	}
	applyLatestSATStatuses(&snap.Hardware, DefaultSATBaseDir, nil)
	return json.MarshalIndent(snap, "", "  ")
}

func readAuditSnapshot(auditJSON []byte) (schema.HardwareIngestRequest, error) {
	var snap schema.HardwareIngestRequest
	if err := json.Unmarshal(auditJSON, &snap); err != nil {
		return schema.HardwareIngestRequest{}, err
	}
	collector.NormalizeSnapshot(&snap.Hardware, snap.CollectedAt)
	return snap, nil
}

func (a *App) RunAudit(runtimeMode runtimeenv.Mode, output string) (string, error) {
	if runtimeMode == runtimeenv.ModeLiveCD {
		if err := a.runtime.CaptureTechnicalDump(DefaultTechDumpDir); err != nil {
			slog.Warn("capture technical dump", "err", err)
		}
	}
	result := collector.Run(runtimeMode)
	applyLatestSATStatuses(&result.Hardware, DefaultSATBaseDir, a.StatusDB)
	writePSUStatusesToDB(a.StatusDB, result.Hardware.PowerSupplies)
	writePCIeGPUStatusesToDB(a.StatusDB, result.Hardware.PCIeDevices)
	if health, err := ReadRuntimeHealth(DefaultRuntimeJSONPath); err == nil {
		result.Runtime = &health
	}
	data, err := json.MarshalIndent(result, "", "  ")
	if err != nil {
		return "", err
	}

	switch {
	case output == "stdout":
		_, err := os.Stdout.Write(append(data, '\n'))
		return "stdout", err
	case strings.HasPrefix(output, "file:"):
		path := strings.TrimPrefix(output, "file:")
		if err := atomicWriteFile(path, append(data, '\n'), 0644); err != nil {
			return "", err
		}
		return path, nil
	default:
		return "", fmt.Errorf("unknown output destination %q — use stdout or file:<path>", output)
	}
}

func (a *App) RunRuntimePreflight(output string) (string, error) {
	health, err := a.runtime.CollectRuntimeHealth(DefaultExportDir)
	if err != nil {
		return "", err
	}
	data, err := json.MarshalIndent(health, "", "  ")
	if err != nil {
		return "", err
	}

	switch {
	case output == "stdout":
		_, err := os.Stdout.Write(append(data, '\n'))
		return "stdout", err
	case strings.HasPrefix(output, "file:"):
		path := strings.TrimPrefix(output, "file:")
		if err := atomicWriteFile(path, append(data, '\n'), 0644); err != nil {
			return "", err
		}
		return path, nil
	default:
		return "", fmt.Errorf("unknown output destination %q — use stdout or file:<path>", output)
	}
}

func (a *App) RuntimeHealthResult() ActionResult {
	health, err := ReadRuntimeHealth(DefaultRuntimeJSONPath)
	if err != nil {
		return ActionResult{Title: "Runtime issues", Body: "No runtime health found."}
	}
	driverLabel := "Driver ready"
	accelLabel := "CUDA ready"
	switch a.sat.DetectGPUVendor() {
	case "amd":
		driverLabel = "AMDGPU ready"
		accelLabel = "ROCm SMI ready"
	case "nvidia":
		driverLabel = "NVIDIA ready"
	}
	var body strings.Builder
	fmt.Fprintf(&body, "Status: %s\n", firstNonEmpty(health.Status, "UNKNOWN"))
	fmt.Fprintf(&body, "Export dir: %s\n", firstNonEmpty(health.ExportDir, DefaultExportDir))
	fmt.Fprintf(&body, "%s: %t\n", driverLabel, health.DriverReady)
	fmt.Fprintf(&body, "%s: %t\n", accelLabel, health.CUDAReady)
	fmt.Fprintf(&body, "Network: %s", firstNonEmpty(health.NetworkStatus, "UNKNOWN"))
	if len(health.Issues) > 0 {
		body.WriteString("\n\nIssues:\n")
		for _, issue := range health.Issues {
			fmt.Fprintf(&body, "- %s: %s\n", issue.Code, issue.Description)
		}
	}
	return ActionResult{Title: "Runtime issues", Body: strings.TrimSpace(body.String())}
}

func (a *App) RunAuditNow(runtimeMode runtimeenv.Mode) (ActionResult, error) {
	path, err := a.RunAudit(runtimeMode, "file:"+DefaultAuditJSONPath)
	body := "Audit completed."
	if path != "" {
		body = "Audit output: " + path
	}
	return ActionResult{Title: "Run audit", Body: body}, err
}

func (a *App) HealthSummaryResult() ActionResult {
	raw, err := os.ReadFile(DefaultAuditJSONPath)
	if err != nil {
		return ActionResult{Title: "Health summary", Body: "No audit JSON found."}
	}
	var snapshot schema.HardwareIngestRequest
	if err := json.Unmarshal(raw, &snapshot); err != nil {
		return ActionResult{Title: "Health summary", Body: "Audit JSON is unreadable."}
	}
	collector.NormalizeSnapshot(&snapshot.Hardware, snapshot.CollectedAt)

	summary := collector.BuildHealthSummary(snapshot.Hardware)
	var body strings.Builder
	status := summary.Status
	if status == "" {
		status = "Unknown"
	}
	fmt.Fprintf(&body, "Overall: %s\n", status)
	fmt.Fprintf(&body, "Storage: warn=%d fail=%d\n", summary.StorageWarn, summary.StorageFail)
	fmt.Fprintf(&body, "PCIe: warn=%d fail=%d\n", summary.PCIeWarn, summary.PCIeFail)
	fmt.Fprintf(&body, "PSU: warn=%d fail=%d\n", summary.PSUWarn, summary.PSUFail)
	fmt.Fprintf(&body, "Memory: warn=%d fail=%d\n", summary.MemoryWarn, summary.MemoryFail)
	for _, item := range latestSATSummaries() {
		fmt.Fprintf(&body, "\n\n%s", item)
	}
	if len(summary.Failures) > 0 {
		fmt.Fprintf(&body, "\n\nFailures:\n- %s", strings.Join(summary.Failures, "\n- "))
	}
	if len(summary.Warnings) > 0 {
		fmt.Fprintf(&body, "\n\nWarnings:\n- %s", strings.Join(summary.Warnings, "\n- "))
	}
	return ActionResult{Title: "Health summary", Body: strings.TrimSpace(body.String())}
}

func (a *App) MainBanner() string {
	raw, err := os.ReadFile(DefaultAuditJSONPath)
	if err != nil {
		return ""
	}

	var snapshot schema.HardwareIngestRequest
	if err := json.Unmarshal(raw, &snapshot); err != nil {
		return ""
	}
	collector.NormalizeSnapshot(&snapshot.Hardware, snapshot.CollectedAt)

	var lines []string
	if system := formatSystemLine(snapshot.Hardware.Board); system != "" {
		lines = append(lines, system)
	}
	if cpu := formatCPULine(snapshot.Hardware.CPUs); cpu != "" {
		lines = append(lines, cpu)
	}
	if memory := formatMemoryLine(snapshot.Hardware.Memory); memory != "" {
		lines = append(lines, memory)
	}
	if storage := formatStorageLine(snapshot.Hardware.Storage); storage != "" {
		lines = append(lines, storage)
	}
	if gpu := formatGPULine(snapshot.Hardware.PCIeDevices); gpu != "" {
		lines = append(lines, gpu)
	}
	if ip := formatIPLine(a.network.ListInterfaces); ip != "" {
		lines = append(lines, ip)
	}

	return strings.TrimSpace(strings.Join(lines, "\n"))
}

// writePSUStatusesToDB records PSU statuses collected during audit into the
// component-status DB so they are visible in the Hardware Summary card.
// PSU status is sourced from IPMI (ipmitool fru + sdr) during audit.
func writePSUStatusesToDB(db *ComponentStatusDB, psus []schema.HardwarePowerSupply) {
	if db == nil || len(psus) == 0 {
		return
	}
	const source = "audit:ipmi"
	worstStatus := "OK"
	for _, psu := range psus {
		if psu.Status == nil {
			continue
		}
		slot := "?"
		if psu.Slot != nil {
			slot = *psu.Slot
		}
		st := *psu.Status
		detail := ""
		if psu.ErrorDescription != nil {
			detail = *psu.ErrorDescription
		}
		db.Record("psu:"+slot, source, st, detail)
		switch st {
		case "Critical":
			worstStatus = "Critical"
		case "Warning":
			if worstStatus != "Critical" {
				worstStatus = "Warning"
			}
		}
	}
	db.Record("psu:all", source, worstStatus, "")
}

// writePCIeGPUStatusesToDB records GPU PCIe-link statuses (e.g. link speed
// degraded to Gen1 while capable of Gen4) collected during audit into the
// component-status DB. Without this, the DB's pcie:gpu:<vendor> key is only
// ever written by SAT diag results (ApplySATResultToDB), and none of the
// nvidia/nvidia-config/nvidia-interconnect/nvidia-bandwidth SAT jobs check
// PCIe link speed — so a real Gen1/Gen4 degradation the collector already
// flagged as Warning in the hardware snapshot never surfaces in the DB-backed
// "check passed" status the Hardware Summary/webui reads.
func writePCIeGPUStatusesToDB(db *ComponentStatusDB, devices []schema.HardwarePCIeDevice) {
	if db == nil || len(devices) == 0 {
		return
	}
	const source = "audit:pcie"
	worst := map[string]string{}
	detail := map[string]string{}
	for _, dev := range devices {
		if dev.Status == nil {
			continue
		}
		st := *dev.Status
		for _, vendor := range []string{"nvidia", "amd"} {
			if !matchesGPUVendor(dev, vendor) {
				continue
			}
			if _, seen := worst[vendor]; !seen {
				worst[vendor] = "OK"
			}
			switch st {
			case "Critical":
				worst[vendor] = "Critical"
				if dev.ErrorDescription != nil {
					detail[vendor] = *dev.ErrorDescription
				}
			case "Warning":
				if worst[vendor] != "Critical" {
					worst[vendor] = "Warning"
					if dev.ErrorDescription != nil {
						detail[vendor] = *dev.ErrorDescription
					}
				}
			}
		}
	}
	for _, vendor := range []string{"nvidia", "amd"} {
		status, present := worst[vendor]
		if !present {
			continue
		}
		db.Record("pcie:gpu:"+vendor, source, status, detail[vendor])
	}
}

func ReadRuntimeHealth(path string) (schema.RuntimeHealth, error) {
	raw, err := os.ReadFile(path)
	if err != nil {
		return schema.RuntimeHealth{}, err
	}
	var health schema.RuntimeHealth
	if err := json.Unmarshal(raw, &health); err != nil {
		return schema.RuntimeHealth{}, err
	}
	return health, nil
}