HostKeeper/cmd/hostkeeper/connect.go

package main

import (
	"context"
	"fmt"
	"os"
	"os/exec"
	"strings"
	"time"

	"github.com/spf13/cobra"

	"git.tukangketik.id/swanadiva/hostkeeper/internal/errors"
	"git.tukangketik.id/swanadiva/hostkeeper/internal/models"
	"git.tukangketik.id/swanadiva/hostkeeper/pkg/config"
	"git.tukangketik.id/swanadiva/hostkeeper/pkg/ssh"
	"git.tukangketik.id/swanadiva/hostkeeper/pkg/storage"
)

var (
	connectTimeout int
	connectNative  bool
)

// connectCmd represents the connect command
var connectCmd = &cobra.Command{
	Use:   "connect <host-name-or-id>",
	Short: "Connect to a saved SSH host",
	Long: `Connect to a saved SSH host using stored credentials.

Examples:
  # Connect to a host by name
  hostkeeper connect myserver

  # Connect with a specific timeout
  hostkeeper connect myserver --timeout 60

  # Use native system SSH instead of Go SSH client
  hostkeeper connect myserver --native`,
	Args: cobra.ExactArgs(1),
	RunE: runConnect,
}

func init() {
	connectCmd.Flags().IntVar(&connectTimeout, "timeout", 30, "Connection timeout in seconds")
	connectCmd.Flags().BoolVar(&connectNative, "native", false, "Use native system SSH instead of Go SSH client")

	rootCmd.AddCommand(connectCmd)
}

func runConnect(cmd *cobra.Command, args []string) error {
	hostIdentifier := args[0]

	cfg := appCfg
	if cfg == nil {
		var err error
		cfg, err = config.New()
		if err != nil {
			return fmt.Errorf("failed to initialize config: %w", err)
		}
	}

	// Initialize storage
	store, err := storage.NewJSONStorage(cfg.GetDataDir())
	if err != nil {
		return fmt.Errorf("failed to initialize storage: %w", err)
	}

	ctx := context.Background()

	// Find host by name or ID
	host, err := findHost(ctx, store, hostIdentifier)
	if err != nil {
		return err
	}

	fmt.Printf("Connecting to %s (%s@%s:%d)...\n", host.Name, host.Username, host.Hostname, host.Port)

	// Choose connection method
	if connectNative {
		return connectWithNativeSSH(host)
	}

	return connectDirectSSH(host)
}

// findHost finds a host by ID first, then by name
func findHost(ctx context.Context, store storage.Storage, identifier string) (*models.Host, error) {
	// Try to find by ID first
	host, err := store.GetHost(ctx, identifier)
	if err == nil {
		return host, nil
	}

	// Try to find by name
	hosts, err := store.ListHosts(ctx)
	if err != nil {
		return nil, fmt.Errorf("failed to list hosts: %w", err)
	}

	for _, h := range hosts {
		if h.Name == identifier {
			return h, nil
		}
	}

	// Try to find by hostname
	for _, h := range hosts {
		if h.Hostname == identifier {
			return h, nil
		}
	}

	// Try to find by ID prefix (short ID match)
	for _, h := range hosts {
		if len(h.ID) >= 8 && h.ID[:8] == identifier {
			return h, nil
		}
	}

	// Host not found, provide helpful error
	msg := fmt.Sprintf("host '%s' not found. Use 'hostkeeper list' to see available hosts", identifier)
	if strings.Contains(identifier, " ") {
		msg += fmt.Sprintf("\nHint: if the host name contains spaces, quote it: connect \"%s\"", identifier)
	}
	return nil, fmt.Errorf("%s", msg)
}

// connectWithNativeSSH uses the system SSH client
func connectWithNativeSSH(host *models.Host) error {
	sshArgs := buildSSHArgs(host)

	cmd := exec.Command("ssh", sshArgs...)

	// Set up standard I/O for interactive session
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

	// Execute SSH command
	if err := cmd.Run(); err != nil {
		return errors.HandleSSHError(err)
	}

	return nil
}

// connectDirectSSH uses Go SSH client
func connectDirectSSH(host *models.Host) error {
	timeout := time.Duration(connectTimeout) * time.Second
	client := ssh.NewClient(host, timeout)

	ctx := context.Background()
	if err := client.Connect(ctx); err != nil {
		return err
	}
	defer client.Close()

	fmt.Printf("Connected to %s (%s@%s:%d)\n", host.Name, host.Username, host.Hostname, host.Port)

	if err := client.Shell(); err != nil {
		return fmt.Errorf("shell session failed: %w", err)
	}

	return nil
}

// buildSSHArgs builds SSH command arguments for the system SSH client
func buildSSHArgs(host *models.Host) []string {
	var args []string

	// Add port if not default
	if host.Port != 22 && host.Port != 0 {
		args = append(args, "-p", fmt.Sprintf("%d", host.Port))
	}

	// Add connection string
	connectionString := fmt.Sprintf("%s@%s", host.Username, host.Hostname)
	args = append(args, connectionString)

	return args
}