path: root/internal/cli/cli.go

package cli

import (
	"flag"
	"fmt"
	"os"
	"os/exec"

	"git.theodohertyfamily.com/tools/wg-wrap/internal/config"
	"git.theodohertyfamily.com/tools/wg-wrap/internal/namespace"
)


type App struct {
	Args           []string
	ConfigDir      string // Optional override for profile storage location
	RuntimeBaseDir string // Optional override for namespace/PID tracking
}

func NewApp(args []string) *App {
	return &App{Args: args}
}

func (a *App) Route() error {
	// 1. Validate arguments for null bytes to prevent exec failures in the C launcher
	for i, arg := range a.Args {
		for j := 0; j < len(arg); j++ {
			if arg[j] == 0 {
				return fmt.Errorf("argument %d contains null byte at position %d", i, j)
			}
		}
	}

	// Handle the internal diagnostic commands that should run on the HOST
	if len(a.Args) > 1 {
		switch a.Args[1] {
		case "show-config":
			return a.showConfig()
		}
	}

	// Handle subcommands first (profile list, import, configure, delete, stop)
	if len(a.Args) > 1 && a.Args[1] == "profile" {
		return a.handleProfileCmd()
	}

	// If we reach here, we are either wrapping a process or running a command
	// that requires isolation (e.g., test-ns, test-args).
	return a.Run()
}

func (a *App) Run() error {
	// Handle the internal diagnostic commands that require ISOLATION
	if len(a.Args) > 1 {
		switch a.Args[1] {
		case "test-ns":
			if !namespace.IsIsolated() {
				if err := namespace.Bootstrap(); err != nil {
					return fmt.Errorf("bootstrap failed: %w", err)
				}
			}
			ok, msg := namespace.VerifyIsolation()
			if !ok {
				return fmt.Errorf("isolation check failed: %s", msg)
			}
			fmt.Println("Isolation Verified: OK")
			return nil
		case "test-args":
			if !namespace.IsIsolated() {
				if err := namespace.Bootstrap(); err != nil {
					return fmt.Errorf("bootstrap failed: %w", err)
				}
			}
			return namespace.VerifyArguments(a.Args)
		}
	}

	cfg := &config.Config{}

	fs := flag.NewFlagSet("wg-wrap", flag.ExitOnError)
	fs.StringVar(&cfg.Profile, "profile", "", "WireGuard profile to use (filename without extension in ~/.config/wg-wrap/profiles/)")
	fs.StringVar(&cfg.DNSServer, "dns-server", "", "Override DNS server to use")

	args := a.Args[1:]
	sepIdx := -1
	for i, arg := range args {
		if arg == "--" {
			sepIdx = i
			break
		}
	}

	var flagsToParse []string
	if sepIdx != -1 {
		flagsToParse = args[:sepIdx]
		cfg.Command = args[sepIdx+1:]
	} else {
		flagsToParse = args
	}

	err := fs.Parse(flagsToParse)
	if err != nil {
		return fmt.Errorf("error parsing flags: %w", err)
	}

	if sepIdx == -1 {
		cfg.Command = fs.Args()
	}

	if len(cfg.Command) == 0 {
		return fmt.Errorf("no command provided. use --help for usage")
	}

	if cfg.Profile == "" {
		cfg.Profile = "default"
	}

	// If we are already isolated, we enter the execution phase.
	if namespace.IsIsolated() {
		return a.ExecuteCommand(cfg)
	}

	// If we are not isolated, we bootstrap.
	// The Bootstrap process will replace this process and restart it.
	if err := namespace.Bootstrap(); err != nil {
		return fmt.Errorf("bootstrap failed: %w", err)
	}

	// This point is never reached because Bootstrap uses syscall.Exec
	return nil
}

// ExecuteCommand handles the isolated execution of the target application.
// This is called after the bootstrap loop has successfully isolated the process.
func (a *App) ExecuteCommand(cfg *config.Config) error {
	if !namespace.IsIsolated() {
		return fmt.Errorf("ExecuteCommand called without namespace isolation")
	}

	// 1. Prepare the namespace
	baseDir := a.RuntimeBaseDir
	if baseDir == "" {
		// Use XDG_RUNTIME_DIR or default via the namespace package
		// Since the namespace package now handles the default in GetProfileNamespacePath,
		// we can pass empty string if no override is present.
		baseDir = "" 
	}

	namespace.PruneStalePids(baseDir, cfg.Profile)
	if err := namespace.RegisterProcess(baseDir, cfg.Profile); err != nil {
		return fmt.Errorf("failed to register process: %w", err)
	}

	// Ensure we unregister and check for cleanup on exit
	defer func() {
		namespace.UnregisterProcess(baseDir, cfg.Profile)
		if last, err := namespace.IsLastProcess(baseDir, cfg.Profile); err == nil && last {
			fmt.Printf("Last process exiting. Cleaning up profile %s...\n", cfg.Profile)
			// Here we would call namespace.UnpinNamespace(baseDir, cfg.Profile) 
			// and terminate the userspace WG process.
		}
	}()

	// 2. VPN Setup (Stubbed)
	fmt.Printf("Initializing WireGuard tunnel for profile %s...\n", cfg.Profile)
	// TODO: Integrate with internal/wireguard to set up TUN and WG-Go

	// 3. Execute the target command
	cmd := exec.Command(cfg.Command[0], cfg.Command[1:]...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Env = os.Environ()

	if err := cmd.Run(); err != nil {
		return fmt.Errorf("command execution failed: %w", err)
	}

	return nil
}

func (a *App) handleProfileCmd() error {
	if len(a.Args) < 3 {
		return fmt.Errorf("usage: wg-wrap profile <list|import|configure|delete|stop> [args]")
	}

	subCmd := a.Args[2]
	switch subCmd {
	case "list":
		fmt.Println("Listing profiles...")
		return fmt.Errorf("profile list not yet implemented")
	case "import":
		if len(a.Args) < 4 {
			return fmt.Errorf("usage: wg-wrap profile import <path>")
		}
		fmt.Printf("Importing profile from %s...\n", a.Args[3])
		return fmt.Errorf("profile import not yet implemented")
	case "configure":
		if len(a.Args) < 4 {
			return fmt.Errorf("usage: wg-wrap profile configure <name>")
		}
		fmt.Printf("Configuring profile %s...\n", a.Args[3])
		return fmt.Errorf("profile configure not yet implemented")
	case "delete":
		if len(a.Args) < 4 {
			return fmt.Errorf("usage: wg-wrap profile delete <name>")
		}
		fmt.Printf("Deleting profile %s...\n", a.Args[3])
		return fmt.Errorf("profile delete not yet implemented")
	case "stop":
		if len(a.Args) < 4 {
			return fmt.Errorf("usage: wg-wrap profile stop <name>")
		}
		fmt.Printf("Stopping profile %s and unpinning namespace...\n", a.Args[3])
		return fmt.Errorf("profile stop not yet implemented")
	default:
		return fmt.Errorf("unknown profile subcommand: %s", subCmd)
	}
}

func (a *App) showConfig() error {
	cfg := &config.Config{}
	fs := flag.NewFlagSet("wg-wrap", flag.ExitOnError)
	fs.StringVar(&cfg.Profile, "profile", "default", "WireGuard profile to use")
	fs.StringVar(&cfg.DNSServer, "dns-server", "", "Override DNS server to use")

	// Parse the arguments that follow 'show-config'
	if len(a.Args) > 2 {
		_ = fs.Parse(a.Args[2:])
	}

	// Determine runtime base directory
	runtimeBase := a.RuntimeBaseDir
	if runtimeBase == "" {
		runtimeBase = os.Getenv("XDG_RUNTIME_DIR")
		if runtimeBase == "" {
			runtimeBase = fmt.Sprintf("/run/user/%d", os.Getuid())
		}
	}

	// Resolve paths
	profilePath := namespace.GetProfileNamespacePath(runtimeBase, cfg.Profile)
	pidsPath := namespace.GetPidsDirPath(runtimeBase, cfg.Profile)

	fmt.Printf("Configuration:\n")
	fmt.Printf("  Profile:       %s\n", cfg.Profile)
	fmt.Printf("  DNS Server:    %s\n", cfg.DNSServer)
	fmt.Printf("  Runtime Base:  %s\n", runtimeBase)
	fmt.Printf("  Profile Path:  %s\n", profilePath)
	fmt.Printf("  PIDs Path:     %s\n", pidsPath)
	fmt.Printf("  Isolated:      %v\n", namespace.IsIsolated())
	fmt.Printf("  UID:           %d\n", os.Getuid())

	return nil
}