sshpiper/plugin/api/api.go

//go:build full || e2e

package main

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"time"

	log "github.com/sirupsen/logrus"

	"github.com/tg123/sshpiper/libplugin"
	"github.com/tg123/sshpiper/libplugin/skel"
)

// ServiceResponse represents the API response for a service lookup
type ServiceResponse struct {
	Service struct {
		Name        string `json:"name"`
		ContainerIP string `json:"container_ip"`
		SSHPort     int    `json:"ssh_port"`
		State       string `json:"state"`
		Node        string `json:"node"` // "cammy" or "ai"
	} `json:"service"`
	Error string `json:"error,omitempty"`
}

type plugin struct {
	APIURL      string
	UpstreamKey string
	DefaultPool string
	CammyHost   string
	AIHost      string

	httpClient *http.Client
}

func newAPIPlugin() *plugin {
	return &plugin{
		httpClient: &http.Client{
			Timeout: 5 * time.Second,
		},
	}
}

// lookupService queries the API for service information
func (p *plugin) lookupService(serviceName string) (*ServiceResponse, error) {
	url := fmt.Sprintf("%s/internal/services/by-name/%s", p.APIURL, serviceName)

	resp, err := p.httpClient.Get(url)
	if err != nil {
		return nil, fmt.Errorf("API request failed: %w", err)
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("failed to read response: %w", err)
	}

	if resp.StatusCode == 404 {
		return nil, fmt.Errorf("service not found: %s", serviceName)
	}

	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("API returned status %d: %s", resp.StatusCode, string(body))
	}

	var serviceResp ServiceResponse
	if err := json.Unmarshal(body, &serviceResp); err != nil {
		return nil, fmt.Errorf("failed to parse response: %w", err)
	}

	return &serviceResp, nil
}

// getPoolHost returns the SSHPiper host for a given pool
func (p *plugin) getPoolHost(pool string) string {
	switch pool {
	case "ai":
		return p.AIHost
	case "cammy":
		return p.CammyHost
	default:
		log.Warnf("Unknown pool %q, using default %q", pool, p.DefaultPool)
		if p.DefaultPool == "ai" {
			return p.AIHost
		}
		return p.CammyHost
	}
}

// apiPipe implements skel.SkelPipe for dynamic API-based routing
type apiPipe struct {
	plugin   *plugin
	username string
	host     string
}

// apiPipeFrom implements skel.SkelPipeFrom
type apiPipeFrom struct {
	plugin   *plugin
	username string
}

// apiPipeTo implements skel.SkelPipeTo
type apiPipeTo struct {
	plugin   *plugin
	username string
	host     string
}

func (p *apiPipe) From() []skel.SkelPipeFrom {
	return []skel.SkelPipeFrom{
		&apiPipeFrom{
			plugin:   p.plugin,
			username: p.username,
		},
	}
}

// MatchConn checks if this pipe matches the connection
func (f *apiPipeFrom) MatchConn(conn libplugin.ConnMetadata) (skel.SkelPipeTo, error) {
	username := conn.User()

	log.Infof("API lookup for service: %s", username)

	// Query the API for this service
	serviceResp, err := f.plugin.lookupService(username)
	if err != nil {
		log.Warnf("API lookup failed for %s: %v, using default pool", username, err)
		// Fall back to default pool
		host := f.plugin.getPoolHost(f.plugin.DefaultPool)
		return &apiPipeTo{
			plugin:   f.plugin,
			username: username,
			host:     host,
		}, nil
	}

	// Determine which pool SSHPiper to route to
	pool := serviceResp.Service.Node
	if pool == "" {
		pool = f.plugin.DefaultPool
	}

	host := f.plugin.getPoolHost(pool)
	log.Infof("Routing %s to pool %s (%s)", username, pool, host)

	return &apiPipeTo{
		plugin:   f.plugin,
		username: username,
		host:     host,
	}, nil
}

// Implement SkelPipeFromPublicKey interface
func (f *apiPipeFrom) AuthorizedKeys(conn libplugin.ConnMetadata) ([]byte, error) {
	return nil, nil
}

func (f *apiPipeFrom) TrustedUserCAKeys(conn libplugin.ConnMetadata) ([]byte, error) {
	return nil, nil
}

// Implement SkelPipeFromPublicKeyAcceptAny interface - accept any key
func (f *apiPipeFrom) AcceptAnyKey() bool {
	return true
}

// apiPipeTo methods
func (t *apiPipeTo) User(conn libplugin.ConnMetadata) string {
	return t.username
}

func (t *apiPipeTo) Host(conn libplugin.ConnMetadata) string {
	return t.host
}

func (t *apiPipeTo) IgnoreHostKey(conn libplugin.ConnMetadata) bool {
	return true
}

func (t *apiPipeTo) KnownHosts(conn libplugin.ConnMetadata) ([]byte, error) {
	return nil, nil
}

// Implement private key auth for upstream
func (t *apiPipeTo) PrivateKey(conn libplugin.ConnMetadata) ([]byte, []byte, error) {
	keyData, err := os.ReadFile(t.plugin.UpstreamKey)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to read upstream key: %w", err)
	}
	return keyData, nil, nil
}

// listPipe returns a catch-all pipe that handles all connections via API lookup
func (p *plugin) listPipe(conn libplugin.ConnMetadata) ([]skel.SkelPipe, error) {
	// Return a single catch-all pipe
	// The actual routing decision happens in MatchConn via API lookup
	return []skel.SkelPipe{
		&apiPipe{
			plugin:   p,
			username: ".*", // catch-all
		},
	}, nil
}