banger/internal/cli/banger.go

package cli

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"syscall"
	"text/tabwriter"
	"time"

	"banger/internal/api"
	"banger/internal/config"
	"banger/internal/model"
	"banger/internal/paths"
	"banger/internal/rpc"
	"banger/internal/system"

	"github.com/spf13/cobra"
)

var (
	bangerdPathFunc = paths.BangerdPath
	daemonExePath   = func(pid int) string {
		return filepath.Join("/proc", fmt.Sprintf("%d", pid), "exe")
	}
)

func NewBangerCommand() *cobra.Command {
	root := &cobra.Command{
		Use:           "banger",
		Short:         "Manage development VMs and images",
		SilenceUsage:  true,
		SilenceErrors: true,
		RunE:          helpNoArgs,
	}
	root.CompletionOptions.DisableDefaultCmd = true
	root.AddCommand(newDaemonCommand(), newVMCommand(), newImageCommand(), newTUICommand())
	return root
}

func newDaemonCommand() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "daemon",
		Short: "Manage the banger daemon",
		RunE:  helpNoArgs,
	}
	cmd.AddCommand(
		&cobra.Command{
			Use:   "status",
			Short: "Show daemon status",
			Args:  noArgsUsage("usage: banger daemon status"),
			RunE: func(cmd *cobra.Command, args []string) error {
				layout, err := paths.Resolve()
				if err != nil {
					return err
				}
				ping, pingErr := rpc.Call[api.PingResult](cmd.Context(), layout.SocketPath, "ping", api.Empty{})
				if pingErr != nil {
					_, err = fmt.Fprintf(cmd.OutOrStdout(), "stopped\nsocket: %s\n", layout.SocketPath)
					return err
				}
				_, err = fmt.Fprintf(cmd.OutOrStdout(), "running\npid: %d\nsocket: %s\n", ping.PID, layout.SocketPath)
				return err
			},
		},
		&cobra.Command{
			Use:   "stop",
			Short: "Stop the daemon",
			Args:  noArgsUsage("usage: banger daemon stop"),
			RunE: func(cmd *cobra.Command, args []string) error {
				if err := system.EnsureSudo(cmd.Context()); err != nil {
					return err
				}
				layout, err := paths.Resolve()
				if err != nil {
					return err
				}
				_, err = rpc.Call[api.ShutdownResult](cmd.Context(), layout.SocketPath, "shutdown", api.Empty{})
				if err != nil {
					if os.IsNotExist(err) || strings.Contains(err.Error(), "connect") {
						_, writeErr := fmt.Fprintln(cmd.OutOrStdout(), "daemon not running")
						return writeErr
					}
					return err
				}
				_, err = fmt.Fprintln(cmd.OutOrStdout(), "stopping")
				return err
			},
		},
		&cobra.Command{
			Use:   "socket",
			Short: "Print the daemon socket path",
			Args:  noArgsUsage("usage: banger daemon socket"),
			RunE: func(cmd *cobra.Command, args []string) error {
				layout, err := paths.Resolve()
				if err != nil {
					return err
				}
				_, err = fmt.Fprintln(cmd.OutOrStdout(), layout.SocketPath)
				return err
			},
		},
	)
	return cmd
}

func newVMCommand() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "vm",
		Short: "Manage virtual machines",
		RunE:  helpNoArgs,
	}
	cmd.AddCommand(
		newVMCreateCommand(),
		newVMListCommand(),
		newVMShowCommand(),
		newVMActionCommand("start", "Start a VM", "vm.start"),
		newVMActionCommand("stop", "Stop a VM", "vm.stop"),
		newVMActionCommand("restart", "Restart a VM", "vm.restart"),
		newVMActionCommand("delete", "Delete a VM", "vm.delete"),
		newVMSetCommand(),
		newVMSSHCommand(),
		newVMLogsCommand(),
		newVMStatsCommand(),
	)
	return cmd
}

func newVMCreateCommand() *cobra.Command {
	var params api.VMCreateParams
	cmd := &cobra.Command{
		Use:   "create",
		Short: "Create a VM",
		Args:  noArgsUsage("usage: banger vm create"),
		RunE: func(cmd *cobra.Command, args []string) error {
			if err := system.EnsureSudo(cmd.Context()); err != nil {
				return err
			}
			layout, _, err := ensureDaemon(cmd.Context())
			if err != nil {
				return err
			}
			result, err := rpc.Call[api.VMShowResult](cmd.Context(), layout.SocketPath, "vm.create", params)
			if err != nil {
				return err
			}
			return printVMSummary(cmd.OutOrStdout(), result.VM)
		},
	}
	cmd.Flags().StringVar(&params.Name, "name", "", "vm name")
	cmd.Flags().StringVar(&params.ImageName, "image", "", "image name or id")
	cmd.Flags().IntVar(&params.VCPUCount, "vcpu", 0, "vcpu count")
	cmd.Flags().IntVar(&params.MemoryMiB, "memory", 0, "memory in MiB")
	cmd.Flags().StringVar(&params.SystemOverlaySize, "system-overlay-size", "", "system overlay size")
	cmd.Flags().StringVar(&params.WorkDiskSize, "disk-size", "", "work disk size")
	cmd.Flags().BoolVar(&params.NATEnabled, "nat", false, "enable NAT")
	cmd.Flags().BoolVar(&params.NoStart, "no-start", false, "create without starting")
	return cmd
}

func newVMListCommand() *cobra.Command {
	return &cobra.Command{
		Use:   "list",
		Short: "List VMs",
		Args:  noArgsUsage("usage: banger vm list"),
		RunE: func(cmd *cobra.Command, args []string) error {
			layout, _, err := ensureDaemon(cmd.Context())
			if err != nil {
				return err
			}
			result, err := rpc.Call[api.VMListResult](cmd.Context(), layout.SocketPath, "vm.list", api.Empty{})
			if err != nil {
				return err
			}
			w := tabwriter.NewWriter(cmd.OutOrStdout(), 0, 8, 2, ' ', 0)
			fmt.Fprintln(w, "ID\tNAME\tSTATE\tIMAGE\tIP\tVCPU\tMEM\tDISK\tCREATED")
			for _, vm := range result.VMs {
				fmt.Fprintf(
					w,
					"%s\t%s\t%s\t%s\t%s\t%d\t%d MiB\t%s\t%s\n",
					shortID(vm.ID),
					vm.Name,
					vm.State,
					shortID(vm.ImageID),
					vm.Runtime.GuestIP,
					vm.Spec.VCPUCount,
					vm.Spec.MemoryMiB,
					model.FormatSizeBytes(vm.Spec.WorkDiskSizeBytes),
					relativeTime(vm.CreatedAt),
				)
			}
			return w.Flush()
		},
	}
}

func newVMShowCommand() *cobra.Command {
	return &cobra.Command{
		Use:   "show <id-or-name>",
		Short: "Show VM details",
		Args:  exactArgsUsage(1, "usage: banger vm show <id-or-name>"),
		RunE: func(cmd *cobra.Command, args []string) error {
			layout, _, err := ensureDaemon(cmd.Context())
			if err != nil {
				return err
			}
			result, err := rpc.Call[api.VMShowResult](cmd.Context(), layout.SocketPath, "vm.show", api.VMRefParams{IDOrName: args[0]})
			if err != nil {
				return err
			}
			return printJSON(cmd.OutOrStdout(), result.VM)
		},
	}
}

func newVMActionCommand(use, short, method string) *cobra.Command {
	return &cobra.Command{
		Use:   use + " <id-or-name>",
		Short: short,
		Args:  exactArgsUsage(1, fmt.Sprintf("usage: banger vm %s <id-or-name>", use)),
		RunE: func(cmd *cobra.Command, args []string) error {
			if err := system.EnsureSudo(cmd.Context()); err != nil {
				return err
			}
			layout, _, err := ensureDaemon(cmd.Context())
			if err != nil {
				return err
			}
			result, err := rpc.Call[api.VMShowResult](cmd.Context(), layout.SocketPath, method, api.VMRefParams{IDOrName: args[0]})
			if err != nil {
				return err
			}
			return printVMSummary(cmd.OutOrStdout(), result.VM)
		},
	}
}

func newVMSetCommand() *cobra.Command {
	var (
		vcpu     int
		memory   int
		diskSize string
		nat      bool
		noNat    bool
	)
	cmd := &cobra.Command{
		Use:   "set <id-or-name>",
		Short: "Update stopped VM settings",
		Args:  exactArgsUsage(1, "usage: banger vm set [--vcpu N] [--memory MiB] [--disk-size SIZE] [--nat|--no-nat] <id-or-name>"),
		RunE: func(cmd *cobra.Command, args []string) error {
			params, err := vmSetParamsFromFlags(args[0], vcpu, memory, diskSize, nat, noNat)
			if err != nil {
				return err
			}
			if err := system.EnsureSudo(cmd.Context()); err != nil {
				return err
			}
			layout, _, err := ensureDaemon(cmd.Context())
			if err != nil {
				return err
			}
			result, err := rpc.Call[api.VMShowResult](cmd.Context(), layout.SocketPath, "vm.set", params)
			if err != nil {
				return err
			}
			return printVMSummary(cmd.OutOrStdout(), result.VM)
		},
	}
	cmd.Flags().IntVar(&vcpu, "vcpu", -1, "vcpu count")
	cmd.Flags().IntVar(&memory, "memory", -1, "memory in MiB")
	cmd.Flags().StringVar(&diskSize, "disk-size", "", "new work disk size")
	cmd.Flags().BoolVar(&nat, "nat", false, "enable NAT")
	cmd.Flags().BoolVar(&noNat, "no-nat", false, "disable NAT")
	return cmd
}

func newVMSSHCommand() *cobra.Command {
	return &cobra.Command{
		Use:   "ssh <id-or-name> [ssh args...]",
		Short: "SSH into a running VM",
		Args:  minArgsUsage(1, "usage: banger vm ssh <id-or-name> [ssh args...]"),
		RunE: func(cmd *cobra.Command, args []string) error {
			layout, cfg, err := ensureDaemon(cmd.Context())
			if err != nil {
				return err
			}
			result, err := rpc.Call[api.VMSSHResult](cmd.Context(), layout.SocketPath, "vm.ssh", api.VMRefParams{IDOrName: args[0]})
			if err != nil {
				return err
			}
			sshArgs, err := sshCommandArgs(cfg, result.GuestIP, args[1:])
			if err != nil {
				return err
			}
			sshCmd := exec.CommandContext(cmd.Context(), "ssh", sshArgs...)
			sshCmd.Stdout = cmd.OutOrStdout()
			sshCmd.Stderr = cmd.ErrOrStderr()
			sshCmd.Stdin = cmd.InOrStdin()
			return sshCmd.Run()
		},
	}
}

func newVMLogsCommand() *cobra.Command {
	var follow bool
	cmd := &cobra.Command{
		Use:   "logs <id-or-name>",
		Short: "Show VM logs",
		Args:  exactArgsUsage(1, "usage: banger vm logs [-f] <id-or-name>"),
		RunE: func(cmd *cobra.Command, args []string) error {
			layout, _, err := ensureDaemon(cmd.Context())
			if err != nil {
				return err
			}
			result, err := rpc.Call[api.VMLogsResult](cmd.Context(), layout.SocketPath, "vm.logs", api.VMRefParams{IDOrName: args[0]})
			if err != nil {
				return err
			}
			if result.LogPath == "" {
				return errors.New("vm has no log path")
			}
			return system.CopyStream(cmd.OutOrStdout(), system.TailCommand(result.LogPath, follow))
		},
	}
	cmd.Flags().BoolVarP(&follow, "follow", "f", false, "follow logs")
	return cmd
}

func newVMStatsCommand() *cobra.Command {
	return &cobra.Command{
		Use:   "stats <id-or-name>",
		Short: "Show VM stats",
		Args:  exactArgsUsage(1, "usage: banger vm stats <id-or-name>"),
		RunE: func(cmd *cobra.Command, args []string) error {
			layout, _, err := ensureDaemon(cmd.Context())
			if err != nil {
				return err
			}
			result, err := rpc.Call[api.VMStatsResult](cmd.Context(), layout.SocketPath, "vm.stats", api.VMRefParams{IDOrName: args[0]})
			if err != nil {
				return err
			}
			return printJSON(cmd.OutOrStdout(), result)
		},
	}
}

func newImageCommand() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "image",
		Short: "Manage images",
		RunE:  helpNoArgs,
	}
	cmd.AddCommand(
		newImageBuildCommand(),
		newImageListCommand(),
		newImageShowCommand(),
		newImageDeleteCommand(),
	)
	return cmd
}

func newImageBuildCommand() *cobra.Command {
	var params api.ImageBuildParams
	cmd := &cobra.Command{
		Use:   "build",
		Short: "Build an image",
		Args:  noArgsUsage("usage: banger image build"),
		RunE: func(cmd *cobra.Command, args []string) error {
			if err := system.EnsureSudo(cmd.Context()); err != nil {
				return err
			}
			layout, _, err := ensureDaemon(cmd.Context())
			if err != nil {
				return err
			}
			result, err := rpc.Call[api.ImageShowResult](cmd.Context(), layout.SocketPath, "image.build", params)
			if err != nil {
				return err
			}
			return printImageSummary(cmd.OutOrStdout(), result.Image)
		},
	}
	cmd.Flags().StringVar(&params.Name, "name", "", "image name")
	cmd.Flags().StringVar(&params.BaseRootfs, "base-rootfs", "", "base rootfs path")
	cmd.Flags().StringVar(&params.Size, "size", "", "output image size")
	cmd.Flags().StringVar(&params.KernelPath, "kernel", "", "kernel path")
	cmd.Flags().StringVar(&params.InitrdPath, "initrd", "", "initrd path")
	cmd.Flags().StringVar(&params.ModulesDir, "modules", "", "modules dir")
	cmd.Flags().BoolVar(&params.Docker, "docker", false, "install docker")
	return cmd
}

func newImageListCommand() *cobra.Command {
	return &cobra.Command{
		Use:   "list",
		Short: "List images",
		Args:  noArgsUsage("usage: banger image list"),
		RunE: func(cmd *cobra.Command, args []string) error {
			layout, _, err := ensureDaemon(cmd.Context())
			if err != nil {
				return err
			}
			result, err := rpc.Call[api.ImageListResult](cmd.Context(), layout.SocketPath, "image.list", api.Empty{})
			if err != nil {
				return err
			}
			w := tabwriter.NewWriter(cmd.OutOrStdout(), 0, 8, 2, ' ', 0)
			fmt.Fprintln(w, "ID\tNAME\tMANAGED\tROOTFS\tCREATED")
			for _, image := range result.Images {
				fmt.Fprintf(w, "%s\t%s\t%t\t%s\t%s\n", shortID(image.ID), image.Name, image.Managed, image.RootfsPath, relativeTime(image.CreatedAt))
			}
			return w.Flush()
		},
	}
}

func newImageShowCommand() *cobra.Command {
	return &cobra.Command{
		Use:   "show <id-or-name>",
		Short: "Show image details",
		Args:  exactArgsUsage(1, "usage: banger image show <id-or-name>"),
		RunE: func(cmd *cobra.Command, args []string) error {
			layout, _, err := ensureDaemon(cmd.Context())
			if err != nil {
				return err
			}
			result, err := rpc.Call[api.ImageShowResult](cmd.Context(), layout.SocketPath, "image.show", api.ImageRefParams{IDOrName: args[0]})
			if err != nil {
				return err
			}
			return printJSON(cmd.OutOrStdout(), result.Image)
		},
	}
}

func newImageDeleteCommand() *cobra.Command {
	return &cobra.Command{
		Use:   "delete <id-or-name>",
		Short: "Delete an image",
		Args:  exactArgsUsage(1, "usage: banger image delete <id-or-name>"),
		RunE: func(cmd *cobra.Command, args []string) error {
			if err := system.EnsureSudo(cmd.Context()); err != nil {
				return err
			}
			layout, _, err := ensureDaemon(cmd.Context())
			if err != nil {
				return err
			}
			result, err := rpc.Call[api.ImageShowResult](cmd.Context(), layout.SocketPath, "image.delete", api.ImageRefParams{IDOrName: args[0]})
			if err != nil {
				return err
			}
			return printImageSummary(cmd.OutOrStdout(), result.Image)
		},
	}
}

func helpNoArgs(cmd *cobra.Command, args []string) error {
	if len(args) != 0 {
		return fmt.Errorf("unknown arguments: %s", strings.Join(args, " "))
	}
	return cmd.Help()
}

func noArgsUsage(usage string) cobra.PositionalArgs {
	return func(cmd *cobra.Command, args []string) error {
		if len(args) != 0 {
			return errors.New(usage)
		}
		return nil
	}
}

func exactArgsUsage(n int, usage string) cobra.PositionalArgs {
	return func(cmd *cobra.Command, args []string) error {
		if len(args) != n {
			return errors.New(usage)
		}
		return nil
	}
}

func minArgsUsage(n int, usage string) cobra.PositionalArgs {
	return func(cmd *cobra.Command, args []string) error {
		if len(args) < n {
			return errors.New(usage)
		}
		return nil
	}
}

func ensureDaemon(ctx context.Context) (paths.Layout, model.DaemonConfig, error) {
	layout, err := paths.Resolve()
	if err != nil {
		return paths.Layout{}, model.DaemonConfig{}, err
	}
	cfg, err := config.Load(layout)
	if err != nil {
		return paths.Layout{}, model.DaemonConfig{}, err
	}
	if ping, err := rpc.Call[api.PingResult](ctx, layout.SocketPath, "ping", api.Empty{}); err == nil {
		if daemonOutdated(ping.PID) {
			if err := restartDaemon(ctx, layout, ping.PID); err != nil {
				return paths.Layout{}, model.DaemonConfig{}, err
			}
			return layout, cfg, nil
		}
		return layout, cfg, nil
	}
	if err := startDaemon(ctx, layout); err != nil {
		return paths.Layout{}, model.DaemonConfig{}, err
	}
	return layout, cfg, nil
}

func daemonOutdated(pid int) bool {
	if pid <= 0 {
		return false
	}
	daemonBin, err := bangerdPathFunc()
	if err != nil {
		return false
	}
	currentInfo, err := os.Stat(daemonBin)
	if err != nil {
		return false
	}
	runningInfo, err := os.Stat(daemonExePath(pid))
	if err != nil {
		return false
	}
	return !os.SameFile(currentInfo, runningInfo)
}

func restartDaemon(ctx context.Context, layout paths.Layout, pid int) error {
	stopCtx, cancel := context.WithTimeout(ctx, 2*time.Second)
	defer cancel()

	_, _ = rpc.Call[api.ShutdownResult](stopCtx, layout.SocketPath, "shutdown", api.Empty{})
	if waitForPIDExit(pid, 2*time.Second) {
		return startDaemon(ctx, layout)
	}
	if proc, err := os.FindProcess(pid); err == nil {
		_ = proc.Signal(syscall.SIGTERM)
	}
	if !waitForPIDExit(pid, 2*time.Second) {
		return fmt.Errorf("timed out restarting stale daemon pid %d", pid)
	}
	return startDaemon(ctx, layout)
}

func waitForPIDExit(pid int, timeout time.Duration) bool {
	deadline := time.Now().Add(timeout)
	for time.Now().Before(deadline) {
		if !pidRunning(pid) {
			return true
		}
		time.Sleep(50 * time.Millisecond)
	}
	return !pidRunning(pid)
}

func pidRunning(pid int) bool {
	if pid <= 0 {
		return false
	}
	proc, err := os.FindProcess(pid)
	if err != nil {
		return false
	}
	return proc.Signal(syscall.Signal(0)) == nil
}

func startDaemon(ctx context.Context, layout paths.Layout) error {
	if err := paths.Ensure(layout); err != nil {
		return err
	}
	logFile, err := os.OpenFile(layout.DaemonLog, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o644)
	if err != nil {
		return err
	}
	defer logFile.Close()

	daemonBin, err := paths.BangerdPath()
	if err != nil {
		return err
	}
	cmd := exec.CommandContext(ctx, daemonBin)
	cmd.Stdout = logFile
	cmd.Stderr = logFile
	cmd.Stdin = nil
	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
	if err := cmd.Start(); err != nil {
		return err
	}
	if err := rpc.WaitForSocket(layout.SocketPath, 5*time.Second); err != nil {
		return fmt.Errorf("daemon failed to start; inspect %s: %w", layout.DaemonLog, err)
	}
	return nil
}

func vmSetParamsFromFlags(idOrName string, vcpu, memory int, diskSize string, nat, noNat bool) (api.VMSetParams, error) {
	if nat && noNat {
		return api.VMSetParams{}, errors.New("use only one of --nat or --no-nat")
	}
	params := api.VMSetParams{IDOrName: idOrName, WorkDiskSize: diskSize}
	if vcpu >= 0 {
		params.VCPUCount = &vcpu
	}
	if memory >= 0 {
		params.MemoryMiB = &memory
	}
	if nat || noNat {
		value := nat && !noNat
		params.NATEnabled = &value
	}
	if params.VCPUCount == nil && params.MemoryMiB == nil && params.WorkDiskSize == "" && params.NATEnabled == nil {
		return api.VMSetParams{}, errors.New("no VM settings changed")
	}
	return params, nil
}

func sshCommandArgs(cfg model.DaemonConfig, guestIP string, extra []string) ([]string, error) {
	if guestIP == "" {
		return nil, errors.New("vm has no guest IP")
	}
	args := []string{}
	if cfg.RepoRoot != "" {
		args = append(args, "-i", filepath.Join(cfg.RepoRoot, "id_ed25519"))
	}
	args = append(args, "-o", "StrictHostKeyChecking=no", "-o", "UserKnownHostsFile=/dev/null", "root@"+guestIP)
	args = append(args, extra...)
	return args, nil
}

func printJSON(out anyWriter, v any) error {
	data, err := json.MarshalIndent(v, "", "  ")
	if err != nil {
		return err
	}
	_, err = fmt.Fprintln(out, string(data))
	return err
}

func printVMSummary(out anyWriter, vm model.VMRecord) error {
	_, err := fmt.Fprintf(
		out,
		"%s\t%s\t%s\t%s\t%s\t%s\n",
		shortID(vm.ID),
		vm.Name,
		vm.State,
		vm.Runtime.GuestIP,
		model.FormatSizeBytes(vm.Spec.WorkDiskSizeBytes),
		vm.Runtime.DNSName,
	)
	return err
}

func printImageSummary(out anyWriter, image model.Image) error {
	_, err := fmt.Fprintf(out, "%s\t%s\t%t\t%s\n", shortID(image.ID), image.Name, image.Managed, image.RootfsPath)
	return err
}

type anyWriter interface {
	Write(p []byte) (n int, err error)
}

func shortID(id string) string {
	if len(id) <= 12 {
		return id
	}
	return id[:12]
}

func relativeTime(t time.Time) string {
	if t.IsZero() {
		return "-"
	}
	delta := time.Since(t)
	switch {
	case delta < 30*time.Second:
		return "moments ago"
	case delta < time.Minute:
		return fmt.Sprintf("%d seconds ago", int(delta.Seconds()))
	case delta < 2*time.Minute:
		return "1 minute ago"
	case delta < time.Hour:
		return fmt.Sprintf("%d minutes ago", int(delta.Minutes()))
	case delta < 2*time.Hour:
		return "1 hour ago"
	case delta < 24*time.Hour:
		return fmt.Sprintf("%d hours ago", int(delta.Hours()))
	case delta < 48*time.Hour:
		return "1 day ago"
	case delta < 7*24*time.Hour:
		return fmt.Sprintf("%d days ago", int(delta.Hours()/24))
	case delta < 14*24*time.Hour:
		return "1 week ago"
	default:
		return fmt.Sprintf("%d weeks ago", int(delta.Hours()/(24*7)))
	}
}