prx/internal/cli/root.go

// SPDX-License-Identifier: MIT

// Package cli implements the prx command line interface.
package cli

import (
	"context"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"os"
	"os/signal"
	"syscall"

	"github.com/spf13/cobra"

	"prx/internal/client"
	"prx/internal/config"
	"prx/internal/proto"
)

// app carries the state shared by every command.
type app struct {
	configPath string
	verbose    bool

	out    io.Writer
	errOut io.Writer
	logger *slog.Logger
}

// Execute builds and runs the root command, cancelling in-flight work when the
// process receives SIGINT or SIGTERM.
func Execute() error {
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()

	a := &app{out: os.Stdout, errOut: os.Stderr}
	root := a.newRootCmd()
	return root.ExecuteContext(ctx)
}

func (a *app) newRootCmd() *cobra.Command {
	root := &cobra.Command{
		Use:   "prx",
		Short: "Safe proxy routing manager",
		Long: "prx enables and disables proxy routing while protecting inbound\n" +
			"connections (such as SSH) from being pulled into the tunnel.",
		SilenceUsage:  true,
		SilenceErrors: true,
		PersistentPreRunE: func(_ *cobra.Command, _ []string) error {
			level := slog.LevelInfo
			if a.verbose {
				level = slog.LevelDebug
			}
			a.logger = slog.New(slog.NewTextHandler(a.errOut, &slog.HandlerOptions{Level: level}))
			return nil
		},
	}
	root.PersistentFlags().StringVar(&a.configPath, "config", config.DefaultPath(), "path to config file")
	root.PersistentFlags().BoolVarP(&a.verbose, "verbose", "v", false, "verbose logging")

	root.AddCommand(
		a.newOnCmd(),
		a.newOffCmd(),
		a.newToggleCmd(),
		a.newStatusCmd(),
		a.newDoctorCmd(),
		a.newBypassCmd(),
		a.newLogsCmd(),
		a.newInitCmd(),
		a.newVersionCmd(),
	)
	return root
}

func (a *app) loadConfig() (*config.Config, error) {
	return config.Load(a.configPath)
}

func (a *app) buildPlan(cfg *config.Config, action proto.Action) proto.Plan {
	return proto.Plan{
		Action:          action,
		WANInterface:    cfg.Network.WANInterface,
		WANGateway4:     cfg.Network.WANGateway,
		WANGateway6:     cfg.Network.WANGateway6,
		Mark:            cfg.Marks.InboundMark,
		TableV4:         cfg.Routing.TableV4,
		TableV6:         cfg.Routing.TableV6,
		PriorityV4:      cfg.Routing.PriorityV4,
		PriorityV6:      cfg.Routing.PriorityV6,
		BypassTableV4:   cfg.Routing.BypassTableV4,
		BypassTableV6:   cfg.Routing.BypassTableV6,
		BypassPrioV4:    cfg.Routing.BypassPrioV4,
		BypassPrioV6:    cfg.Routing.BypassPrioV6,
		BypassCIDRs:     cfg.Bypass.ExtraCIDRs,
		ManageEngine:    cfg.Engine.Manage,
		ServiceName:     cfg.Engine.ServiceName,
		EngineConfig:    cfg.Engine.ConfigPath,
		StopEngineOnOff: cfg.Engine.StopOnOff,
		HealthProbe:     cfg.Health.Probe,
		StatePath:       cfg.State.Path,
	}
}

// dispatch sends a plan to the helper, prints the result and turns a failed
// result into an error. mutate, when non-nil, adjusts the plan before it is
// sent.
func (a *app) dispatch(ctx context.Context, cfg *config.Config, action proto.Action, mutate func(*proto.Plan)) error {
	plan := a.buildPlan(cfg, action)
	if mutate != nil {
		mutate(&plan)
	}
	if a.logger != nil {
		a.logger.Debug("dispatching plan", "action", string(action), "state", plan.StatePath)
	}
	res, err := client.Dispatch(ctx, plan)
	if err != nil {
		return err
	}
	a.printResult(res)
	if !res.OK {
		return formatResultError(res)
	}
	return nil
}

func (a *app) printResult(res proto.Result) {
	for _, m := range res.Messages {
		fmt.Fprintln(a.out, "  "+m)
	}
	if len(res.Checks) == 0 {
		return
	}
	fmt.Fprintln(a.out, "  checks:")
	for _, c := range res.Checks {
		line := fmt.Sprintf("    [%s] %s", checkMark(c), c.Name)
		if c.Detail != "" {
			line += ": " + c.Detail
		}
		fmt.Fprintln(a.out, line)
	}
}

// checkMark renders a diagnostic check as a short status marker.
func checkMark(c proto.Check) string {
	switch {
	case c.OK:
		return "ok"
	case c.Warn:
		return "warn"
	default:
		return "FAIL"
	}
}

func formatResultError(res proto.Result) error {
	if res.Error == "" {
		return errors.New("operation failed")
	}
	return errors.New(res.Error)
}