Refactor bee CLI and LiveCD integration
This commit is contained in:
Mikhail Chusavitin
@@ -1,167 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"flag"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"bee/audit/internal/collector"
|
||||
)
|
||||
|
||||
// Version is the audit binary version.
|
||||
// Injected at release build time via:
|
||||
//
|
||||
// -ldflags "-X main.Version=1.2"
|
||||
//
|
||||
// Defaults to "dev" in local builds.
|
||||
var Version = "dev"
|
||||
|
||||
func main() {
|
||||
output := flag.String("output", "stdout", `output destination:
|
||||
stdout — print JSON to stdout (default)
|
||||
file:<path> — write JSON to file
|
||||
usb — auto-detect removable media, write JSON there`)
|
||||
showVersion := flag.Bool("version", false, "print version and exit")
|
||||
flag.Parse()
|
||||
|
||||
if *showVersion {
|
||||
fmt.Println(Version)
|
||||
return
|
||||
}
|
||||
|
||||
slog.SetDefault(slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{
|
||||
Level: slog.LevelInfo,
|
||||
})))
|
||||
|
||||
result := collector.Run()
|
||||
|
||||
data, err := json.MarshalIndent(result, "", " ")
|
||||
if err != nil {
|
||||
slog.Error("marshal result", "err", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
if err := writeOutput(*output, data); err != nil {
|
||||
slog.Error("write output", "destination", *output, "err", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
}
|
||||
|
||||
func writeOutput(dest string, data []byte) error {
|
||||
switch {
|
||||
case dest == "stdout":
|
||||
_, err := os.Stdout.Write(append(data, '\n'))
|
||||
return err
|
||||
|
||||
case strings.HasPrefix(dest, "file:"):
|
||||
path := strings.TrimPrefix(dest, "file:")
|
||||
return os.WriteFile(path, append(data, '\n'), 0644)
|
||||
|
||||
case dest == "usb":
|
||||
return writeToUSB(data)
|
||||
|
||||
default:
|
||||
return fmt.Errorf("unknown output destination %q — use stdout, file:<path>, or usb", dest)
|
||||
}
|
||||
}
|
||||
|
||||
// writeToUSB auto-detects the first removable block device, mounts it,
|
||||
// and writes the audit JSON. Falls back to /tmp on any failure.
|
||||
func writeToUSB(data []byte) error {
|
||||
boardSerial := extractBoardSerial(data)
|
||||
filename := auditFilename(boardSerial, time.Now().UTC())
|
||||
|
||||
device, err := firstRemovableDevice()
|
||||
if err != nil {
|
||||
slog.Warn("usb output: no removable device, writing to /tmp", "err", err)
|
||||
return writeAuditToPath(filepath.Join("/tmp", filename), data)
|
||||
}
|
||||
|
||||
mountpoint := "/tmp/bee-usb"
|
||||
if err := os.MkdirAll(mountpoint, 0755); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if err := exec.Command("mount", device, mountpoint).Run(); err != nil {
|
||||
slog.Warn("usb output: mount failed, writing to /tmp", "device", device, "err", err)
|
||||
return writeAuditToPath(filepath.Join("/tmp", filename), data)
|
||||
}
|
||||
defer func() {
|
||||
if err := exec.Command("umount", mountpoint).Run(); err != nil {
|
||||
slog.Warn("usb output: umount failed", "mountpoint", mountpoint, "err", err)
|
||||
}
|
||||
}()
|
||||
|
||||
path := filepath.Join(mountpoint, filename)
|
||||
if err := writeAuditToPath(path, data); err != nil {
|
||||
slog.Warn("usb output: write failed, falling back to /tmp", "path", path, "err", err)
|
||||
return writeAuditToPath(filepath.Join("/tmp", filename), data)
|
||||
}
|
||||
|
||||
slog.Info("usb output: written", "path", path)
|
||||
return nil
|
||||
}
|
||||
|
||||
func writeAuditToPath(path string, data []byte) error {
|
||||
if err := os.WriteFile(path, append(data, '\n'), 0644); err != nil {
|
||||
return err
|
||||
}
|
||||
slog.Info("audit output written", "path", path)
|
||||
return nil
|
||||
}
|
||||
|
||||
func extractBoardSerial(data []byte) string {
|
||||
var doc struct {
|
||||
Hardware struct {
|
||||
Board struct {
|
||||
SerialNumber string `json:"serial_number"`
|
||||
} `json:"board"`
|
||||
} `json:"hardware"`
|
||||
}
|
||||
if err := json.Unmarshal(data, &doc); err != nil {
|
||||
return "unknown"
|
||||
}
|
||||
serial := strings.TrimSpace(doc.Hardware.Board.SerialNumber)
|
||||
if serial == "" {
|
||||
return "unknown"
|
||||
}
|
||||
return serial
|
||||
}
|
||||
|
||||
func auditFilename(boardSerial string, now time.Time) string {
|
||||
boardSerial = strings.TrimSpace(boardSerial)
|
||||
if boardSerial == "" {
|
||||
boardSerial = "unknown"
|
||||
}
|
||||
return fmt.Sprintf("audit-%s-%s.json", boardSerial, now.Format("20060102-150405"))
|
||||
}
|
||||
|
||||
func firstRemovableDevice() (string, error) {
|
||||
entries, err := os.ReadDir("/sys/block")
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
sort.Slice(entries, func(i, j int) bool { return entries[i].Name() < entries[j].Name() })
|
||||
|
||||
for _, e := range entries {
|
||||
name := e.Name()
|
||||
if strings.HasPrefix(name, "loop") || strings.HasPrefix(name, "ram") {
|
||||
continue
|
||||
}
|
||||
removableFlag, err := os.ReadFile(filepath.Join("/sys/block", name, "removable"))
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
if strings.TrimSpace(string(removableFlag)) == "1" {
|
||||
return filepath.Join("/dev", name), nil
|
||||
}
|
||||
}
|
||||
return "", fmt.Errorf("no removable block device found")
|
||||
}
|
||||
185
audit/cmd/bee/main.go
Normal file
185
audit/cmd/bee/main.go
Normal file
@@ -0,0 +1,185 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"flag"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"os"
|
||||
"strings"
|
||||
|
||||
"bee/audit/internal/app"
|
||||
"bee/audit/internal/platform"
|
||||
"bee/audit/internal/runtimeenv"
|
||||
"bee/audit/internal/tui"
|
||||
)
|
||||
|
||||
var Version = "dev"
|
||||
|
||||
func main() {
|
||||
os.Exit(run(os.Args[1:], os.Stdout, os.Stderr))
|
||||
}
|
||||
|
||||
func run(args []string, stdout, stderr io.Writer) int {
|
||||
slog.SetDefault(slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{
|
||||
Level: slog.LevelInfo,
|
||||
})))
|
||||
|
||||
if len(args) == 0 {
|
||||
printRootUsage(stderr)
|
||||
return 1
|
||||
}
|
||||
|
||||
switch args[0] {
|
||||
case "help", "--help", "-h":
|
||||
printRootUsage(stdout)
|
||||
return 0
|
||||
case "audit":
|
||||
return runAudit(args[1:], stdout, stderr)
|
||||
case "tui":
|
||||
return runTUI(args[1:], stdout, stderr)
|
||||
case "export":
|
||||
return runExport(args[1:], stdout, stderr)
|
||||
case "sat":
|
||||
return runSAT(args[1:], stdout, stderr)
|
||||
case "version", "--version", "-version":
|
||||
fmt.Fprintln(stdout, Version)
|
||||
return 0
|
||||
default:
|
||||
fmt.Fprintf(stderr, "bee: unknown command %q\n\n", args[0])
|
||||
printRootUsage(stderr)
|
||||
return 1
|
||||
}
|
||||
}
|
||||
|
||||
func printRootUsage(w io.Writer) {
|
||||
fmt.Fprintln(w, `bee commands:
|
||||
bee audit --runtime auto|local|livecd --output stdout|file:<path>
|
||||
bee tui --runtime auto|local|livecd
|
||||
bee export --target <device>
|
||||
bee sat nvidia
|
||||
bee version`)
|
||||
}
|
||||
|
||||
func runAudit(args []string, stdout, stderr io.Writer) int {
|
||||
fs := flag.NewFlagSet("audit", flag.ContinueOnError)
|
||||
fs.SetOutput(stderr)
|
||||
output := fs.String("output", "stdout", "output destination: stdout or file:<path>")
|
||||
runtimeFlag := fs.String("runtime", "auto", "runtime environment: auto, local, livecd")
|
||||
showVersion := fs.Bool("version", false, "print version and exit")
|
||||
fs.Usage = func() {
|
||||
fmt.Fprintln(stderr, "usage: bee audit [--runtime auto|local|livecd] [--output stdout|file:<path>]")
|
||||
fs.PrintDefaults()
|
||||
}
|
||||
if err := fs.Parse(args); err != nil {
|
||||
return 2
|
||||
}
|
||||
if *showVersion {
|
||||
fmt.Fprintln(stdout, Version)
|
||||
return 0
|
||||
}
|
||||
|
||||
runtimeInfo, err := runtimeenv.Detect(*runtimeFlag)
|
||||
if err != nil {
|
||||
slog.Error("resolve runtime", "err", err)
|
||||
return 1
|
||||
}
|
||||
slog.Info("runtime resolved", "mode", runtimeInfo.Mode, "reason", runtimeInfo.Reason)
|
||||
|
||||
application := app.New(platform.New())
|
||||
path, err := application.RunAudit(runtimeInfo.Mode, *output)
|
||||
if err != nil {
|
||||
slog.Error("run audit", "err", err)
|
||||
return 1
|
||||
}
|
||||
if path != "stdout" {
|
||||
slog.Info("audit output written", "path", path)
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func runTUI(args []string, stdout, stderr io.Writer) int {
|
||||
fs := flag.NewFlagSet("tui", flag.ContinueOnError)
|
||||
fs.SetOutput(stderr)
|
||||
runtimeFlag := fs.String("runtime", "auto", "runtime environment: auto, local, livecd")
|
||||
fs.Usage = func() {
|
||||
fmt.Fprintln(stderr, "usage: bee tui [--runtime auto|local|livecd]")
|
||||
fs.PrintDefaults()
|
||||
}
|
||||
if err := fs.Parse(args); err != nil {
|
||||
return 2
|
||||
}
|
||||
|
||||
runtimeInfo, err := runtimeenv.Detect(*runtimeFlag)
|
||||
if err != nil {
|
||||
slog.Error("resolve runtime", "err", err)
|
||||
return 1
|
||||
}
|
||||
|
||||
application := app.New(platform.New())
|
||||
if err := tui.Run(application, runtimeInfo.Mode); err != nil {
|
||||
slog.Error("run tui", "err", err)
|
||||
return 1
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func runExport(args []string, stdout, stderr io.Writer) int {
|
||||
fs := flag.NewFlagSet("export", flag.ContinueOnError)
|
||||
fs.SetOutput(stderr)
|
||||
targetDevice := fs.String("target", "", "removable device path, e.g. /dev/sdb1")
|
||||
fs.Usage = func() {
|
||||
fmt.Fprintln(stderr, "usage: bee export --target <device>")
|
||||
fs.PrintDefaults()
|
||||
}
|
||||
if err := fs.Parse(args); err != nil {
|
||||
return 2
|
||||
}
|
||||
if strings.TrimSpace(*targetDevice) == "" {
|
||||
fmt.Fprintln(stderr, "bee export: --target is required")
|
||||
fs.Usage()
|
||||
return 2
|
||||
}
|
||||
|
||||
application := app.New(platform.New())
|
||||
targets, err := application.ListRemovableTargets()
|
||||
if err != nil {
|
||||
slog.Error("list removable targets", "err", err)
|
||||
return 1
|
||||
}
|
||||
|
||||
for _, target := range targets {
|
||||
if target.Device == *targetDevice {
|
||||
path, err := application.ExportLatestAudit(target)
|
||||
if err != nil {
|
||||
slog.Error("export latest audit", "err", err)
|
||||
return 1
|
||||
}
|
||||
slog.Info("audit exported", "path", path)
|
||||
return 0
|
||||
}
|
||||
}
|
||||
|
||||
slog.Error("target device not found among removable filesystems", "device", *targetDevice)
|
||||
return 1
|
||||
}
|
||||
|
||||
func runSAT(args []string, stdout, stderr io.Writer) int {
|
||||
if len(args) == 0 || args[0] == "help" || args[0] == "--help" || args[0] == "-h" {
|
||||
fmt.Fprintln(stderr, "usage: bee sat nvidia")
|
||||
return 2
|
||||
}
|
||||
if args[0] != "nvidia" {
|
||||
fmt.Fprintf(stderr, "bee sat: unknown target %q\n", args[0])
|
||||
fmt.Fprintln(stderr, "usage: bee sat nvidia")
|
||||
return 2
|
||||
}
|
||||
application := app.New(platform.New())
|
||||
archive, err := application.RunNvidiaAcceptancePack("")
|
||||
if err != nil {
|
||||
slog.Error("run nvidia sat", "err", err)
|
||||
return 1
|
||||
}
|
||||
slog.Info("nvidia sat archive written", "path", archive)
|
||||
return 0
|
||||
}
|
||||
115
audit/cmd/bee/main_test.go
Normal file
115
audit/cmd/bee/main_test.go
Normal file
@@ -0,0 +1,115 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestRunRootHelp(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
var stdout, stderr bytes.Buffer
|
||||
rc := run([]string{"help"}, &stdout, &stderr)
|
||||
if rc != 0 {
|
||||
t.Fatalf("rc=%d want 0", rc)
|
||||
}
|
||||
if !strings.Contains(stdout.String(), "bee commands:") {
|
||||
t.Fatalf("stdout missing root usage:\n%s", stdout.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunNoArgsPrintsUsage(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
var stdout, stderr bytes.Buffer
|
||||
rc := run(nil, &stdout, &stderr)
|
||||
if rc != 1 {
|
||||
t.Fatalf("rc=%d want 1", rc)
|
||||
}
|
||||
if !strings.Contains(stderr.String(), "bee commands:") {
|
||||
t.Fatalf("stderr missing root usage:\n%s", stderr.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunUnknownCommand(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
var stdout, stderr bytes.Buffer
|
||||
rc := run([]string{"wat"}, &stdout, &stderr)
|
||||
if rc != 1 {
|
||||
t.Fatalf("rc=%d want 1", rc)
|
||||
}
|
||||
if !strings.Contains(stderr.String(), `unknown command "wat"`) {
|
||||
t.Fatalf("stderr missing unknown command message:\n%s", stderr.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunVersion(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
old := Version
|
||||
Version = "test-version"
|
||||
t.Cleanup(func() { Version = old })
|
||||
|
||||
var stdout, stderr bytes.Buffer
|
||||
rc := run([]string{"version"}, &stdout, &stderr)
|
||||
if rc != 0 {
|
||||
t.Fatalf("rc=%d want 0", rc)
|
||||
}
|
||||
if strings.TrimSpace(stdout.String()) != "test-version" {
|
||||
t.Fatalf("stdout=%q want %q", strings.TrimSpace(stdout.String()), "test-version")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunExportRequiresTarget(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
var stdout, stderr bytes.Buffer
|
||||
rc := run([]string{"export"}, &stdout, &stderr)
|
||||
if rc != 2 {
|
||||
t.Fatalf("rc=%d want 2", rc)
|
||||
}
|
||||
if !strings.Contains(stderr.String(), "--target is required") {
|
||||
t.Fatalf("stderr missing target error:\n%s", stderr.String())
|
||||
}
|
||||
if !strings.Contains(stderr.String(), "usage: bee export --target <device>") {
|
||||
t.Fatalf("stderr missing export usage:\n%s", stderr.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunSATUsage(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
var stdout, stderr bytes.Buffer
|
||||
rc := run([]string{"sat"}, &stdout, &stderr)
|
||||
if rc != 2 {
|
||||
t.Fatalf("rc=%d want 2", rc)
|
||||
}
|
||||
if !strings.Contains(stderr.String(), "usage: bee sat nvidia") {
|
||||
t.Fatalf("stderr missing sat usage:\n%s", stderr.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunSATUnknownTarget(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
var stdout, stderr bytes.Buffer
|
||||
rc := run([]string{"sat", "amd"}, &stdout, &stderr)
|
||||
if rc != 2 {
|
||||
t.Fatalf("rc=%d want 2", rc)
|
||||
}
|
||||
if !strings.Contains(stderr.String(), `unknown target "amd"`) {
|
||||
t.Fatalf("stderr missing sat target error:\n%s", stderr.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunAuditInvalidRuntime(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
var stdout, stderr bytes.Buffer
|
||||
rc := run([]string{"audit", "--runtime", "bad"}, &stdout, &stderr)
|
||||
if rc != 1 {
|
||||
t.Fatalf("rc=%d want 1", rc)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user