AetherForge/server/internal/api/pathtracer_handler.go

package api

import (
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"image/color"
	"image/png"
	"log"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/go-chi/chi/v5"
	"github.com/google/uuid"
	qrcode "github.com/skip2/go-qrcode"

	"golang.org/x/crypto/curve25519"
)

// ── types ─────────────────────────────────────────────────────────────────────

// HopStatus tracks one agent's progress in a trace session.
type HopStatus string

const (
	HopPending   HopStatus = "pending"
	HopReady     HopStatus = "ready"
	HopFailed    HopStatus = "failed"
)

// HopInfo records what we know about each hop in the chain.
type HopInfo struct {
	AgentID     string    `json:"agent_id"`
	AgentName   string    `json:"agent_name"`
	ExternalIP  string    `json:"external_ip"`
	Port        int       `json:"port"`
	PublicKey   string    `json:"public_key"`
	PrivateKey  string    `json:"-"` // never sent to client
	LocalAddr   string    `json:"local_addr"`
	Status      HopStatus `json:"status"`
	Error       string    `json:"error,omitempty"`
}

// TraceSession holds all state for one active VPN session.
type TraceSession struct {
	ID        string    `json:"id"`
	AgentIDs  []string  `json:"agent_ids"`
	Hops      []*HopInfo `json:"hops"`
	Ready     bool      `json:"ready"`
	Error     string    `json:"error,omitempty"`
	CreatedAt time.Time `json:"created_at"`
	// Client WireGuard keypair — used to build the QR config.
	clientPrivKey string
	clientPubKey  string
}

const (
	pathTraceSessionTTL      = 2 * time.Hour
	pathTraceCleanupInterval = 5 * time.Minute
)

// PathTracerHandler manages on-demand WireGuard chain sessions.
type PathTracerHandler struct {
	hub      *WSHub
	mu       sync.Mutex
	sessions map[string]*TraceSession
	stopCh   chan struct{}
}

func NewPathTracerHandler(hub *WSHub) *PathTracerHandler {
	h := &PathTracerHandler{
		hub:      hub,
		sessions: make(map[string]*TraceSession),
		stopCh:   make(chan struct{}),
	}
	go h.sessionCleanupLoop()
	return h
}

func (h *PathTracerHandler) sessionCleanupLoop() {
	ticker := time.NewTicker(pathTraceCleanupInterval)
	defer ticker.Stop()
	for {
		select {
		case <-ticker.C:
			h.expireSessions()
		case <-h.stopCh:
			return
		}
	}
}

func (h *PathTracerHandler) expireSessions() {
	now := time.Now()
	var expired []*TraceSession
	h.mu.Lock()
	for id, sess := range h.sessions {
		if now.Sub(sess.CreatedAt) > pathTraceSessionTTL {
			expired = append(expired, sess)
			delete(h.sessions, id)
		}
	}
	h.mu.Unlock()
	for _, sess := range expired {
		log.Printf("[pathtrace] session %s expired after %s", sess.ID[:8], pathTraceSessionTTL)
		for _, hop := range sess.Hops {
			_ = h.hub.writeAgentJSON(hop.AgentID, Message{
				Type:    "command",
				Payload: mustMarshal(map[string]interface{}{"action": "wg_teardown"}),
			})
		}
	}
}

// ── HTTP handlers ─────────────────────────────────────────────────────────────

// POST /api/v1/pathtrace/start
// Body: {"agent_ids": ["id1","id2",...]}
func (h *PathTracerHandler) Start(w http.ResponseWriter, r *http.Request) {
	var req struct {
		AgentIDs []string `json:"agent_ids"`
	}
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil || len(req.AgentIDs) == 0 {
		http.Error(w, "agent_ids required", http.StatusBadRequest)
		return
	}
	if len(req.AgentIDs) > 3 {
		http.Error(w, "max 3 hops supported", http.StatusBadRequest)
		return
	}

	clientPriv, clientPub, err := generateServerWGKeyPair()
	if err != nil {
		http.Error(w, "keygen failed", http.StatusInternalServerError)
		return
	}

	sess := &TraceSession{
		ID:            uuid.New().String(),
		AgentIDs:      req.AgentIDs,
		CreatedAt:     time.Now(),
		clientPrivKey: clientPriv,
		clientPubKey:  clientPub,
	}

	// Build hops with placeholder names (agent name lookup below).
	for i, id := range req.AgentIDs {
		sess.Hops = append(sess.Hops, &HopInfo{
			AgentID:   id,
			AgentName: fmt.Sprintf("hop-%d", i+1),
			LocalAddr: fmt.Sprintf("10.66.0.%d/24", i+2), // .2, .3, .4
			Port:      51820,
			Status:    HopPending,
		})
	}

	h.mu.Lock()
	h.sessions[sess.ID] = sess
	h.mu.Unlock()

	// Orchestrate asynchronously so the HTTP response returns quickly.
	go h.orchestrate(sess)

	writeJSON(w, map[string]interface{}{
		"session_id": sess.ID,
		"hops":       sess.Hops,
	})
}

// GET /api/v1/pathtrace/{id}/status
func (h *PathTracerHandler) Status(w http.ResponseWriter, r *http.Request) {
	sess := h.getSession(chi.URLParam(r, "id"))
	if sess == nil {
		http.Error(w, "session not found", http.StatusNotFound)
		return
	}
	h.mu.Lock()
	defer h.mu.Unlock()
	writeJSON(w, map[string]interface{}{
		"session_id": sess.ID,
		"ready":      sess.Ready,
		"error":      sess.Error,
		"hops":       sess.Hops,
	})
}

// GET /api/v1/pathtrace/{id}/qr
func (h *PathTracerHandler) QR(w http.ResponseWriter, r *http.Request) {
	sess := h.getSession(chi.URLParam(r, "id"))
	if sess == nil {
		http.Error(w, "session not found", http.StatusNotFound)
		return
	}
	h.mu.Lock()
	ready := sess.Ready
	h.mu.Unlock()
	if !ready {
		http.Error(w, "session not ready yet", http.StatusAccepted)
		return
	}

	cfg := h.buildClientConfig(sess)

	// Return format: ?format=png → PNG image, default → JSON with text+png.
	if r.URL.Query().Get("format") == "png" {
		qr, err := qrcode.New(cfg, qrcode.High)
		if err != nil {
			http.Error(w, "qr generation failed", http.StatusInternalServerError)
			return
		}
		qr.BackgroundColor = color.Black
		qr.ForegroundColor = color.RGBA{R: 0, G: 255, B: 170, A: 255} // neon green
		img := qr.Image(400)
		w.Header().Set("Content-Type", "image/png")
		_ = png.Encode(w, img)
		return
	}

	qr, err := qrcode.Encode(cfg, qrcode.High, 300)
	if err != nil {
		http.Error(w, "qr generation failed", http.StatusInternalServerError)
		return
	}
	writeJSON(w, map[string]interface{}{
		"config":     cfg,
		"qr_png_b64": base64.StdEncoding.EncodeToString(qr),
	})
}

// DELETE /api/v1/pathtrace/{id}
func (h *PathTracerHandler) Delete(w http.ResponseWriter, r *http.Request) {
	id := chi.URLParam(r, "id")
	sess := h.getSession(id)
	if sess == nil {
		http.Error(w, "session not found", http.StatusNotFound)
		return
	}

	// Tell all agents to tear down.
	for _, hop := range sess.Hops {
		_ = h.hub.writeAgentJSON(hop.AgentID, Message{
			Type:    "command",
			Payload: mustMarshal(map[string]interface{}{"action": "wg_teardown"}),
		})
	}

	h.mu.Lock()
	delete(h.sessions, id)
	h.mu.Unlock()

	writeJSON(w, map[string]interface{}{"ok": true})
}

// ── orchestration ─────────────────────────────────────────────────────────────

func (h *PathTracerHandler) orchestrate(sess *TraceSession) {
	log.Printf("[pathtrace] session %s: orchestrating %d hop(s)", sess.ID[:8], len(sess.Hops))

	// Phase 1: send wg_setup to every hop in parallel, collect keypairs + IPs.
	type setupResp struct {
		hop *HopInfo
		pub string
		ip  string
		port int
		err string
	}
	results := make(chan setupResp, len(sess.Hops))

	for _, hop := range sess.Hops {
		hop := hop
		ch := h.hub.AwaitCommandResult(hop.AgentID, "wg_setup")
		err := h.hub.writeAgentJSON(hop.AgentID, Message{
			Type:    "command",
			Payload: mustMarshal(map[string]interface{}{"action": "wg_setup"}),
		})
		if err != nil {
			h.hub.CancelAwait(hop.AgentID, "wg_setup")
			results <- setupResp{hop: hop, err: "agent not connected: " + err.Error()}
			continue
		}
		go func() {
			select {
			case payload := <-ch:
				msgStr, _ := payload["message"].(string)
				// message is JSON-encoded WGSetupResult
				var res struct {
					PublicKey    string `json:"public_key"`
					ExternalIP   string `json:"external_ip"`
					ExternalPort int    `json:"external_port"`
					Error        string `json:"error"`
				}
				if jerr := json.Unmarshal([]byte(msgStr), &res); jerr != nil {
					results <- setupResp{hop: hop, err: "parse error: " + jerr.Error()}
					return
				}
				if res.Error != "" {
					results <- setupResp{hop: hop, err: res.Error}
					return
				}
				p := res.ExternalPort
				if p == 0 {
					p = 51820
				}
				// If UPnP failed, fall back to the agent's last known IP in the DB.
				ip := res.ExternalIP
				if ip == "" {
					if ag, err := h.hub.db.GetAgent(hop.AgentID); err == nil && strings.TrimSpace(ag.IP) != "" {
						ip = ag.IP
					}
				}
				results <- setupResp{hop: hop, pub: res.PublicKey, ip: ip, port: p}
			case <-time.After(20 * time.Second):
				results <- setupResp{hop: hop, err: "timeout waiting for wg_setup response"}
			}
		}()
	}

	// Collect Phase 1 results.
	for range sess.Hops {
		r := <-results
		h.mu.Lock()
		if r.err != "" {
			r.hop.Status = HopFailed
			r.hop.Error = r.err
			sess.Error = "hop " + r.hop.AgentID[:8] + " failed: " + r.err
		} else {
			r.hop.PublicKey = r.pub
			r.hop.ExternalIP = r.ip
			r.hop.Port = r.port
			r.hop.Status = HopReady
		}
		h.mu.Unlock()
	}

	// If any hop failed at setup, abort.
	h.mu.Lock()
	anyFailed := false
	for _, hop := range sess.Hops {
		if hop.Status == HopFailed {
			anyFailed = true
			break
		}
	}
	h.mu.Unlock()
	if anyFailed {
		log.Printf("[pathtrace] session %s: setup failed", sess.ID[:8])
		return
	}

	// Phase 2: send wg_configure to each hop.
	// Topology:
	//   - Hop 1 always peers to the client (10.66.0.1/32) so return traffic works.
	//   - Relay hops also peer forward to the next hop (0.0.0.0/0).
	//   - Middle/exit hops peer back to the previous hop for reverse routing.
	// The CLIENT config always points to the FIRST hop.

	type cfgResp struct {
		hop *HopInfo
		err string
	}
	cfgResults := make(chan cfgResp, len(sess.Hops))

	for i, hop := range sess.Hops {
		hop := hop
		i := i

		payload := map[string]interface{}{
			"session_id":           sess.ID,
			"private_key":          "", // agent uses its own generated key
			"local_address":        hop.LocalAddr,
			"listen_port":          hop.Port,
			"enable_ip_forwarding": true,
			"peers":                buildHopPeers(sess, i),
		}

		dataJSON, _ := json.Marshal(payload)

		ch := h.hub.AwaitCommandResult(hop.AgentID, "wg_configure")
		if err := h.hub.writeAgentJSON(hop.AgentID, Message{
			Type: "command",
			Payload: mustMarshal(map[string]interface{}{
				"action": "wg_configure",
				"data":   string(dataJSON),
			}),
		}); err != nil {
			h.hub.CancelAwait(hop.AgentID, "wg_configure")
			cfgResults <- cfgResp{hop: hop, err: "agent not connected: " + err.Error()}
			continue
		}

		go func() {
			select {
			case p := <-ch:
				success, _ := p["success"].(bool)
				if !success {
					msg, _ := p["message"].(string)
					cfgResults <- cfgResp{hop: hop, err: msg}
					return
				}
				cfgResults <- cfgResp{hop: hop}
			case <-time.After(60 * time.Second):
				cfgResults <- cfgResp{hop: hop, err: "timeout waiting for wg_configure"}
			}
		}()
	}

	for range sess.Hops {
		r := <-cfgResults
		h.mu.Lock()
		if r.err != "" {
			r.hop.Status = HopFailed
			r.hop.Error = r.err
			if sess.Error == "" {
				sess.Error = "configure failed on " + r.hop.AgentID[:8] + ": " + r.err
			}
		}
		h.mu.Unlock()
	}

	h.mu.Lock()
	allReady := true
	for _, hop := range sess.Hops {
		if hop.Status != HopReady {
			allReady = false
		}
	}
	if allReady {
		sess.Ready = true
	}
	h.mu.Unlock()

	log.Printf("[pathtrace] session %s: orchestration complete, ready=%v", sess.ID[:8], allReady)
}

// buildHopPeers returns WireGuard peer entries for hop index i in the chain.
func buildHopPeers(sess *TraceSession, i int) []map[string]interface{} {
	var peers []map[string]interface{}

	// Hop 1 is the client entry point — always accept the phone/client tunnel.
	if i == 0 {
		peers = append(peers, map[string]interface{}{
			"public_key":           sess.clientPubKey,
			"allowed_ips":          "10.66.0.1/32",
			"persistent_keepalive": 25,
		})
	}

	// Forward peer: route outbound traffic to the next hop in the chain.
	if i < len(sess.Hops)-1 {
		nextHop := sess.Hops[i+1]
		peers = append(peers, map[string]interface{}{
			"public_key":           nextHop.PublicKey,
			"endpoint":             fmt.Sprintf("%s:%d", nextHop.ExternalIP, nextHop.Port),
			"allowed_ips":          "0.0.0.0/0",
			"persistent_keepalive": 25,
		})
	}

	// Reverse peer: return traffic toward the client via the previous hop.
	if i > 0 {
		prevHop := sess.Hops[i-1]
		peers = append(peers, map[string]interface{}{
			"public_key":           prevHop.PublicKey,
			"allowed_ips":          "10.66.0.1/32",
			"persistent_keepalive": 25,
		})
	}

	return peers
}

// buildClientConfig generates the WireGuard config text the user scans/imports.
func (h *PathTracerHandler) buildClientConfig(sess *TraceSession) string {
	h.mu.Lock()
	defer h.mu.Unlock()

	var sb strings.Builder
	sb.WriteString("[Interface]\n")
	sb.WriteString("PrivateKey = " + sess.clientPrivKey + "\n")
	sb.WriteString("Address = 10.66.0.1/24\n")
	sb.WriteString("DNS = 1.1.1.1\n\n")

	// Phone always connects to the first hop.
	first := sess.Hops[0]
	sb.WriteString("[Peer]\n")
	sb.WriteString("PublicKey = " + first.PublicKey + "\n")
	sb.WriteString(fmt.Sprintf("Endpoint = %s:%d\n", first.ExternalIP, first.Port))
	sb.WriteString("AllowedIPs = 0.0.0.0/0\n")
	sb.WriteString("PersistentKeepalive = 25\n")

	return sb.String()
}

// ── helpers ───────────────────────────────────────────────────────────────────

func (h *PathTracerHandler) getSession(id string) *TraceSession {
	h.mu.Lock()
	defer h.mu.Unlock()
	return h.sessions[id]
}

// getAgentConnByID returns the AgentConnection for the given ID (nil if offline).
func (h *WSHub) getAgentConnByID(id string) *AgentConnection {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return h.agents[id]
}

func generateServerWGKeyPair() (privB64, pubB64 string, err error) {
	var priv [32]byte
	if _, err = rand.Read(priv[:]); err != nil {
		return
	}
	priv[0] &= 248
	priv[31] &= 127
	priv[31] |= 64
	var pub [32]byte
	curve25519.ScalarBaseMult(&pub, &priv)
	privB64 = base64.StdEncoding.EncodeToString(priv[:])
	pubB64 = base64.StdEncoding.EncodeToString(pub[:])
	return
}