agent-helm/helmd/hub.go

package main

import (
	"crypto/subtle"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"sort"
	"strings"
	"sync"
	"time"
)

const (
	workPollTimeout = 25 * time.Second // hur länge en agent-poll får hänga
	proxyTimeout    = 90 * time.Second // hur länge hubben väntar på ett jobbsvar
	offlineAfter    = 15 * time.Second // ingen poll på så länge ⇒ offline
)

// hubJob är ett HTTP-anrop som hubben ber agenten utföra mot sitt
// lokala API. Body är base64 i JSON ([]byte).
type hubJob struct {
	ID          string `json:"id"`
	Method      string `json:"method"`
	Path        string `json:"path"` // t.ex. /api/sessions/x/screen?ansi=1
	ContentType string `json:"content_type,omitempty"`
	Body        []byte `json:"body,omitempty"`
}

type hubResult struct {
	ID      string              `json:"id"`
	Status  int                 `json:"status"`
	Headers map[string][]string `json:"headers,omitempty"`
	Body    []byte              `json:"body,omitempty"`
}

// hubAgent är en registrerad agent. Agenter ser aldrig varandra: hela
// /hub/*-ytan är avgränsad till den egna pollnyckeln.
type hubAgent struct {
	Name      string    `json:"name"`
	Host      string    `json:"host"`
	Version   string    `json:"version"`
	Online    bool      `json:"online"`
	Connected time.Time `json:"connected"`
	LastSeen  time.Time `json:"last_seen"`

	pollKey string
	jobs    chan *hubJob
	pollers int
}

type hubEvent struct {
	Agent   string      `json:"agent,omitempty"`
	Type    string      `json:"type"`
	Session string      `json:"session,omitempty"`
	Data    interface{} `json:"data,omitempty"`
}

type Hub struct {
	cfg  *HubConfig
	auth Authenticator
	sess *sessionStore

	mu      sync.Mutex
	agents  map[string]*hubAgent
	results map[string]chan *hubResult

	subMu sync.Mutex
	subs  map[chan hubEvent]struct{}
}

func NewHub(cfg *HubConfig) *Hub {
	h := &Hub{
		cfg:     cfg,
		auth:    &localAuth{cfg: cfg},
		sess:    newSessionStore(time.Duration(cfg.SessionTTLHours) * time.Hour),
		agents:  map[string]*hubAgent{},
		results: map[string]chan *hubResult{},
		subs:    map[chan hubEvent]struct{}{},
	}
	go h.watchOffline()
	return h
}

// ---------- events ----------

func (h *Hub) subscribe() chan hubEvent {
	ch := make(chan hubEvent, 32)
	h.subMu.Lock()
	h.subs[ch] = struct{}{}
	h.subMu.Unlock()
	return ch
}

func (h *Hub) unsubscribe(ch chan hubEvent) {
	h.subMu.Lock()
	delete(h.subs, ch)
	h.subMu.Unlock()
}

func (h *Hub) emit(ev hubEvent) {
	h.subMu.Lock()
	for ch := range h.subs {
		select {
		case ch <- ev:
		default: // långsam klient: släpp hellre än blockera
		}
	}
	h.subMu.Unlock()
}

// watchOffline flaggar agenter som slutat polla och skickar events när
// online-läget ändras.
func (h *Hub) watchOffline() {
	for {
		time.Sleep(5 * time.Second)
		h.mu.Lock()
		for _, a := range h.agents {
			online := a.pollers > 0 || time.Since(a.LastSeen) < offlineAfter
			if online != a.Online {
				a.Online = online
				state := "offline"
				if online {
					state = "online"
				}
				go h.emit(hubEvent{Agent: a.Name, Type: "agent", Data: state})
				log.Printf("agent %s är %s", a.Name, state)
			}
		}
		h.mu.Unlock()
	}
}

// ---------- agentsidan (/hub/*) ----------

// checkAgent letar upp agenten och verifierar dess pollnyckel.
func (h *Hub) checkAgent(r *http.Request) *hubAgent {
	name := r.URL.Query().Get("name")
	key := r.URL.Query().Get("key")
	h.mu.Lock()
	defer h.mu.Unlock()
	a := h.agents[name]
	if a == nil || subtle.ConstantTimeCompare([]byte(key), []byte(a.pollKey)) != 1 {
		return nil
	}
	a.LastSeen = time.Now()
	return a
}

func (h *Hub) register(w http.ResponseWriter, r *http.Request) {
	var req struct {
		Name    string `json:"name"`
		Host    string `json:"host"`
		Version string `json:"version"`
		Key     string `json:"key"`
	}
	if err := readJSON(r, &req); err != nil {
		jsonErr(w, 400, err.Error())
		return
	}
	if subtle.ConstantTimeCompare([]byte(req.Key), []byte(h.cfg.AgentKey)) != 1 {
		time.Sleep(400 * time.Millisecond)
		jsonErr(w, 401, "fel agent-nyckel — hämta rätt med `helmd hub agentkey` på hubben")
		return
	}
	if req.Name == "" || strings.ContainsAny(req.Name, " \t/?#%") {
		jsonErr(w, 400, fmt.Sprintf("ogiltigt agentnamn %q", req.Name))
		return
	}
	a := &hubAgent{
		Name: req.Name, Host: req.Host, Version: req.Version,
		Online: true, Connected: time.Now(), LastSeen: time.Now(),
		pollKey: newToken(), jobs: make(chan *hubJob, 64),
	}
	h.mu.Lock()
	if old := h.agents[req.Name]; old != nil {
		log.Printf("agent %s omregistrerad (ersätter anslutning från %s)", req.Name, old.Host)
	} else {
		log.Printf("agent %s registrerad (host %s, version %s)", req.Name, req.Host, req.Version)
	}
	h.agents[req.Name] = a
	h.mu.Unlock()
	h.emit(hubEvent{Agent: a.Name, Type: "agent", Data: "online"})
	jsonOut(w, 200, map[string]string{"key": a.pollKey})
}

// work är agentens long-poll: häng tills ett jobb finns eller timeout.
func (h *Hub) work(w http.ResponseWriter, r *http.Request) {
	a := h.checkAgent(r)
	if a == nil {
		jsonErr(w, 401, "okänd agent eller fel nyckel — registrera om")
		return
	}
	h.mu.Lock()
	a.pollers++
	h.mu.Unlock()
	defer func() {
		h.mu.Lock()
		a.pollers--
		a.LastSeen = time.Now()
		h.mu.Unlock()
	}()
	select {
	case job := <-a.jobs:
		jsonOut(w, 200, job)
	case <-time.After(workPollTimeout):
		w.WriteHeader(204)
	case <-r.Context().Done():
	}
}

func (h *Hub) result(w http.ResponseWriter, r *http.Request) {
	a := h.checkAgent(r)
	if a == nil {
		jsonErr(w, 401, "okänd agent eller fel nyckel — registrera om")
		return
	}
	var res hubResult
	if err := json.NewDecoder(http.MaxBytesReader(w, r.Body, 96<<20)).Decode(&res); err != nil {
		jsonErr(w, 400, err.Error())
		return
	}
	h.mu.Lock()
	ch := h.results[res.ID]
	h.mu.Unlock()
	if ch != nil {
		select {
		case ch <- &res:
		default:
		}
	}
	jsonOut(w, 200, map[string]bool{"ok": true})
}

func (h *Hub) agentEvents(w http.ResponseWriter, r *http.Request) {
	a := h.checkAgent(r)
	if a == nil {
		jsonErr(w, 401, "okänd agent eller fel nyckel — registrera om")
		return
	}
	var req struct {
		Events []Event `json:"events"`
	}
	if err := readJSON(r, &req); err != nil {
		jsonErr(w, 400, err.Error())
		return
	}
	for _, ev := range req.Events {
		h.emit(hubEvent{Agent: a.Name, Type: ev.Type, Session: ev.Session, Data: ev.Data})
	}
	jsonOut(w, 200, map[string]bool{"ok": true})
}

// ---------- klientsidan (/api/*, kräver login) ----------

func (h *Hub) requireAuth(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		tok := r.Header.Get("X-Api-Token")
		if tok == "" {
			if hdr := r.Header.Get("Authorization"); strings.HasPrefix(hdr, "Bearer ") {
				tok = strings.TrimPrefix(hdr, "Bearer ")
			}
		}
		if tok == "" {
			tok = r.URL.Query().Get("token") // SSE/EventSource + <img>-länkar
		}
		if _, ok := h.sess.Lookup(tok); !ok {
			jsonErr(w, http.StatusUnauthorized, "inte inloggad")
			return
		}
		next(w, r)
	}
}

func (h *Hub) login(w http.ResponseWriter, r *http.Request) {
	var req struct {
		Username string `json:"username"`
		Password string `json:"password"`
	}
	if err := readJSON(r, &req); err != nil {
		jsonErr(w, 400, err.Error())
		return
	}
	if !h.auth.Authenticate(req.Username, req.Password) {
		time.Sleep(400 * time.Millisecond) // dämpa gissningsförsök
		jsonErr(w, 401, "fel användarnamn eller lösenord")
		return
	}
	tok := h.sess.Create(req.Username)
	log.Printf("användare %s loggade in", req.Username)
	jsonOut(w, 200, map[string]interface{}{
		"token": tok, "user": req.Username, "ttl_hours": h.cfg.SessionTTLHours,
	})
}

func (h *Hub) listAgents() []*hubAgent {
	h.mu.Lock()
	defer h.mu.Unlock()
	out := make([]*hubAgent, 0, len(h.agents))
	for _, a := range h.agents {
		a.Online = a.pollers > 0 || time.Since(a.LastSeen) < offlineAfter
		out = append(out, a)
	}
	sort.Slice(out, func(i, j int) bool { return out[i].Name < out[j].Name })
	return out
}

// proxy vidarebefordrar /api/agents/{agent}/{rest...} till agentens
// lokala /api/{rest} genom jobb-kanalen.
func (h *Hub) proxy(w http.ResponseWriter, r *http.Request) {
	name := r.PathValue("agent")
	rest := r.PathValue("rest")
	if rest == "events" || strings.HasPrefix(rest, "events/") {
		jsonErr(w, 404, "events proxas inte per agent — använd hubbens /api/events")
		return
	}
	h.mu.Lock()
	a := h.agents[name]
	h.mu.Unlock()
	if a == nil {
		jsonErr(w, 404, "ingen sådan agent")
		return
	}

	var body []byte
	if r.Body != nil {
		var err error
		body, err = io.ReadAll(http.MaxBytesReader(w, r.Body, 64<<20))
		if err != nil {
			jsonErr(w, 400, err.Error())
			return
		}
	}
	path := "/api/" + rest
	if r.URL.RawQuery != "" {
		path += "?" + r.URL.RawQuery
	}
	job := &hubJob{
		ID: newToken(), Method: r.Method, Path: path,
		ContentType: r.Header.Get("Content-Type"), Body: body,
	}

	ch := make(chan *hubResult, 1)
	h.mu.Lock()
	h.results[job.ID] = ch
	h.mu.Unlock()
	defer func() {
		h.mu.Lock()
		delete(h.results, job.ID)
		h.mu.Unlock()
	}()

	select {
	case a.jobs <- job:
	default:
		jsonErr(w, 503, "agentens jobbkö är full")
		return
	}

	select {
	case res := <-ch:
		// Behåll bara ofarliga headers — särskilt CSP måste följa med
		// så att delad HTML förblir sandboxad även genom hubben.
		for _, k := range []string{"Content-Type", "Content-Security-Policy", "Content-Disposition", "Cache-Control"} {
			if vs := res.Headers[k]; len(vs) > 0 {
				w.Header()[k] = vs
			}
		}
		w.WriteHeader(res.Status)
		w.Write(res.Body)
	case <-time.After(proxyTimeout):
		jsonErr(w, 504, "agenten svarade inte i tid")
	case <-r.Context().Done():
	}
}

func (h *Hub) routes() *http.ServeMux {
	mux := http.NewServeMux()

	// agentsidan — delad agent-nyckel vid registrering, pollnyckel därefter
	mux.HandleFunc("POST /hub/register", h.register)
	mux.HandleFunc("GET /hub/work", h.work)
	mux.HandleFunc("POST /hub/result", h.result)
	mux.HandleFunc("POST /hub/events", h.agentEvents)

	// oautentiserat: läges-detektering för webbklienten
	mux.HandleFunc("GET /api/mode", func(w http.ResponseWriter, r *http.Request) {
		jsonOut(w, 200, map[string]string{"mode": "hub", "version": version})
	})
	mux.HandleFunc("POST /api/login", h.login)

	mux.HandleFunc("POST /api/logout", func(w http.ResponseWriter, r *http.Request) {
		tok := r.Header.Get("X-Api-Token")
		if hdr := r.Header.Get("Authorization"); strings.HasPrefix(hdr, "Bearer ") {
			tok = strings.TrimPrefix(hdr, "Bearer ")
		}
		h.sess.Delete(tok)
		jsonOut(w, 200, map[string]bool{"ok": true})
	})

	mux.HandleFunc("GET /api/me", h.requireAuth(func(w http.ResponseWriter, r *http.Request) {
		tok := r.Header.Get("X-Api-Token")
		if hdr := r.Header.Get("Authorization"); strings.HasPrefix(hdr, "Bearer ") {
			tok = strings.TrimPrefix(hdr, "Bearer ")
		}
		user, _ := h.sess.Lookup(tok)
		jsonOut(w, 200, map[string]string{"user": user})
	}))

	mux.HandleFunc("GET /api/health", h.requireAuth(func(w http.ResponseWriter, r *http.Request) {
		jsonOut(w, 200, map[string]interface{}{
			"ok": true, "version": version, "agents": len(h.listAgents()),
		})
	}))

	mux.HandleFunc("GET /api/agents", h.requireAuth(func(w http.ResponseWriter, r *http.Request) {
		jsonOut(w, 200, h.listAgents())
	}))

	// agent-nyckeln — för att koppla in nya agenter (kravet: nås från
	// hub-config ELLER efter inloggning)
	mux.HandleFunc("GET /api/agentkey", h.requireAuth(func(w http.ResponseWriter, r *http.Request) {
		jsonOut(w, 200, map[string]string{"agent_key": h.cfg.AgentKey})
	}))

	// per metod — en metodlös route krockar med webbens "GET /"
	for _, m := range []string{"GET", "POST", "PUT", "DELETE"} {
		mux.HandleFunc(m+" /api/agents/{agent}/{rest...}", h.requireAuth(h.proxy))
	}

	mux.HandleFunc("GET /api/events", h.requireAuth(func(w http.ResponseWriter, r *http.Request) {
		fl, ok := w.(http.Flusher)
		if !ok {
			jsonErr(w, 500, "streaming stöds inte")
			return
		}
		w.Header().Set("Content-Type", "text/event-stream")
		w.Header().Set("Cache-Control", "no-cache")
		ch := h.subscribe()
		defer h.unsubscribe(ch)
		fmt.Fprintf(w, "event: hello\ndata: {\"version\":%q,\"mode\":\"hub\"}\n\n", version)
		fl.Flush()
		ping := time.NewTicker(25 * time.Second)
		defer ping.Stop()
		for {
			select {
			case <-r.Context().Done():
				return
			case <-ping.C:
				fmt.Fprint(w, ": ping\n\n")
				fl.Flush()
			case ev := <-ch:
				data, _ := json.Marshal(ev)
				fmt.Fprintf(w, "event: %s\ndata: %s\n\n", ev.Type, data)
				fl.Flush()
			}
		}
	}))

	return mux
}

func hubServe(cfg *HubConfig) error {
	if len(cfg.Users) == 0 {
		pw := os.Getenv("HELMD_HUB_ADMIN_PASSWORD")
		generated := pw == ""
		if generated {
			pw = newToken()[:16]
		}
		if err := cfg.SetPassword("admin", pw); err != nil {
			return err
		}
		if err := cfg.Save(); err != nil {
			return err
		}
		if generated {
			log.Printf("skapade användaren admin med lösenord: %s  (byt med `helmd hub setpass admin`)", pw)
		} else {
			log.Printf("skapade användaren admin (lösenord från HELMD_HUB_ADMIN_PASSWORD)")
		}
	}
	h := NewHub(cfg)
	mux := h.routes()
	mountWeb(mux)
	log.Printf("helm-hub %s lyssnar på http://%s (config %s)", version, cfg.Listen, cfg.path)
	return http.ListenAndServe(cfg.Listen, mux)
}