rmm2/backend/ws/tunnel.go

package ws

import (
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"sync"
	"time"
)

// TunnelSession repraesentiert eine einzelne TCP-Verbindung durch einen Tunnel
type TunnelSession struct {
	ID        string
	TunnelID  string
	ProxyConn net.Conn
	Active    bool
	Ready     chan struct{} // Wird geschlossen wenn Agent die Session bestaetigt hat
	mutex     sync.Mutex
}

// BackendTunnel repraesentiert einen offenen Tunnel (Proxy-Port -> Agent -> Ziel)
type BackendTunnel struct {
	ID         string
	AgentID    string
	TargetHost string
	TargetPort int
	ProxyPort  int
	Listener   net.Listener
	Active     bool
	CreatedAt  time.Time
	Sessions   map[string]*TunnelSession
	sessMutex  sync.RWMutex
}

// TunnelManager verwaltet alle Tunnel und Sessions
type TunnelManager struct {
	hub     *Hub
	tunnels map[string]*BackendTunnel
	mutex   sync.RWMutex

	// Session-Index fuer schnellen Lookup bei Binary-Messages
	sessions  map[string]*TunnelSession
	sessMutex sync.RWMutex

	portStart int
	portEnd   int
	usedPorts map[int]bool
}

func NewTunnelManager(hub *Hub) *TunnelManager {
	return &TunnelManager{
		hub:       hub,
		tunnels:   make(map[string]*BackendTunnel),
		sessions:  make(map[string]*TunnelSession),
		portStart: 10000,
		portEnd:   20000,
		usedPorts: make(map[int]bool),
	}
}

func (tm *TunnelManager) OpenTunnel(agentID, targetHost string, targetPort int) (*BackendTunnel, error) {
	if !tm.hub.IsAgentConnected(agentID) {
		return nil, ErrAgentNotConnected
	}

	tunnelID := tm.generateID()

	proxyPort, err := tm.allocatePort()
	if err != nil {
		return nil, fmt.Errorf("Proxy-Port allokieren fehlgeschlagen: %w", err)
	}

	listener, err := net.Listen("tcp", fmt.Sprintf(":%d", proxyPort))
	if err != nil {
		tm.releasePort(proxyPort)
		return nil, fmt.Errorf("Proxy-Listener erstellen fehlgeschlagen: %w", err)
	}

	tunnel := &BackendTunnel{
		ID:         tunnelID,
		AgentID:    agentID,
		TargetHost: targetHost,
		TargetPort: targetPort,
		ProxyPort:  proxyPort,
		Listener:   listener,
		Active:     true,
		CreatedAt:  time.Now(),
		Sessions:   make(map[string]*TunnelSession),
	}

	tm.mutex.Lock()
	tm.tunnels[tunnelID] = tunnel
	tm.mutex.Unlock()

	// Proxy-Accept-Loop starten
	go tm.runProxyAccept(tunnel)

	log.Printf("Tunnel %s geoeffnet: Agent %s -> %s:%d (Proxy: :%d)",
		tunnelID, agentID, targetHost, targetPort, proxyPort)

	return tunnel, nil
}

func (tm *TunnelManager) CloseTunnel(tunnelID string) error {
	tm.mutex.Lock()
	tunnel, exists := tm.tunnels[tunnelID]
	if !exists {
		tm.mutex.Unlock()
		return fmt.Errorf("Tunnel %s nicht gefunden", tunnelID)
	}

	tunnel.Active = false
	delete(tm.tunnels, tunnelID)
	tm.mutex.Unlock()

	// Listener schliessen
	if tunnel.Listener != nil {
		tunnel.Listener.Close()
	}

	// Alle Sessions schliessen
	tunnel.sessMutex.Lock()
	for sessID, sess := range tunnel.Sessions {
		sess.Active = false
		if sess.ProxyConn != nil {
			sess.ProxyConn.Close()
		}
		// Session aus globalem Index entfernen
		tm.sessMutex.Lock()
		delete(tm.sessions, sessID)
		tm.sessMutex.Unlock()

		// Agent informieren
		tm.sendSessionDisconnect(tunnel.AgentID, sessID)
	}
	tunnel.Sessions = make(map[string]*TunnelSession)
	tunnel.sessMutex.Unlock()

	tm.releasePort(tunnel.ProxyPort)

	log.Printf("Tunnel %s geschlossen", tunnelID)
	return nil
}

func (tm *TunnelManager) GetTunnel(tunnelID string) (*BackendTunnel, bool) {
	tm.mutex.RLock()
	defer tm.mutex.RUnlock()
	tunnel, exists := tm.tunnels[tunnelID]
	return tunnel, exists
}

func (tm *TunnelManager) GetAgentTunnels(agentID string) []*BackendTunnel {
	tm.mutex.RLock()
	defer tm.mutex.RUnlock()

	var tunnels []*BackendTunnel
	for _, tunnel := range tm.tunnels {
		if tunnel.AgentID == agentID {
			tunnels = append(tunnels, tunnel)
		}
	}
	return tunnels
}

func (tm *TunnelManager) GetActiveTunnelCount() int {
	tm.mutex.RLock()
	defer tm.mutex.RUnlock()
	return len(tm.tunnels)
}

func (tm *TunnelManager) CloseAgentTunnels(agentID string) {
	tm.mutex.RLock()
	var tunnelIDs []string
	for id, tunnel := range tm.tunnels {
		if tunnel.AgentID == agentID {
			tunnelIDs = append(tunnelIDs, id)
		}
	}
	tm.mutex.RUnlock()

	for _, tunnelID := range tunnelIDs {
		tm.CloseTunnel(tunnelID)
	}
}

// runProxyAccept nimmt eingehende TCP-Verbindungen am Proxy-Port an
func (tm *TunnelManager) runProxyAccept(tunnel *BackendTunnel) {
	defer func() {
		if tunnel.Active {
			tm.CloseTunnel(tunnel.ID)
		}
	}()

	log.Printf("Tunnel %s: Proxy lauscht auf Port %d", tunnel.ID, tunnel.ProxyPort)

	for tunnel.Active {
		conn, err := tunnel.Listener.Accept()
		if err != nil {
			if tunnel.Active {
				log.Printf("Tunnel %s: Accept-Fehler: %v", tunnel.ID, err)
			}
			break
		}

		// Neue Session fuer diese TCP-Verbindung
		go tm.handleNewSession(tunnel, conn)
	}
}

// handleNewSession erstellt eine neue Session fuer eine eingehende Proxy-Verbindung
func (tm *TunnelManager) handleNewSession(tunnel *BackendTunnel, proxyConn net.Conn) {
	sessionID := tm.generateID()

	session := &TunnelSession{
		ID:        sessionID,
		TunnelID:  tunnel.ID,
		ProxyConn: proxyConn,
		Active:    true,
		Ready:     make(chan struct{}),
	}

	// Session registrieren
	tunnel.sessMutex.Lock()
	tunnel.Sessions[sessionID] = session
	tunnel.sessMutex.Unlock()

	tm.sessMutex.Lock()
	tm.sessions[sessionID] = session
	tm.sessMutex.Unlock()

	log.Printf("Tunnel %s: Neue Session %s von %s", tunnel.ID, sessionID, proxyConn.RemoteAddr())

	// Agent auffordern, Ziel-Verbindung fuer diese Session zu oeffnen
	cmd := map[string]interface{}{
		"type":    "command",
		"id":      tm.generateID(),
		"command": "tunnel_connect",
		"params": map[string]interface{}{
			"session_id":  sessionID,
			"tunnel_id":   tunnel.ID,
			"target_host": tunnel.TargetHost,
			"target_port": tunnel.TargetPort,
		},
	}

	cmdJSON, _ := json.Marshal(cmd)
	if err := tm.hub.SendToAgent(tunnel.AgentID, cmdJSON); err != nil {
		log.Printf("Session %s: tunnel_connect senden fehlgeschlagen: %v", sessionID, err)
		tm.cleanupSession(tunnel, sessionID)
		proxyConn.Close()
		return
	}

	// Warten bis Agent die Session bestaetigt hat (max 10s)
	select {
	case <-session.Ready:
		log.Printf("Session %s: Agent bereit, starte Forwarding", sessionID)
	case <-time.After(10 * time.Second):
		log.Printf("Session %s: Timeout beim Warten auf Agent", sessionID)
		tm.cleanupSession(tunnel, sessionID)
		proxyConn.Close()
		return
	}

	// Proxy -> Agent Data-Forwarding (liest von Proxy-Conn, schickt als Binary an Agent)
	defer func() {
		tm.cleanupSession(tunnel, sessionID)
		proxyConn.Close()
		// Agent informieren dass Session beendet ist
		tm.sendSessionDisconnect(tunnel.AgentID, sessionID)
		log.Printf("Tunnel %s: Session %s beendet", tunnel.ID, sessionID)
	}()

	buffer := make([]byte, 32768)
	for session.Active && tunnel.Active {
		n, err := proxyConn.Read(buffer)
		if err != nil {
			if err != io.EOF && session.Active {
				log.Printf("Session %s: Proxy-Read-Fehler: %v", sessionID, err)
			}
			break
		}

		if n > 0 {
			if err := tm.sendSessionData(tunnel.AgentID, sessionID, buffer[:n]); err != nil {
				log.Printf("Session %s: Weiterleitung fehlgeschlagen: %v", sessionID, err)
				break
			}
		}
	}
}

// ProcessBinaryMessage verarbeitet eingehende Binary-Messages vom Agent (Ziel -> Proxy)
func (tm *TunnelManager) ProcessBinaryMessage(agentID string, data []byte) error {
	if len(data) < 1 {
		return fmt.Errorf("Binary-Message zu kurz")
	}

	idLen := int(data[0])
	if len(data) < 1+idLen {
		return fmt.Errorf("Session-ID zu kurz")
	}

	sessionID := string(data[1 : 1+idLen])
	payload := data[1+idLen:]

	// Session finden
	tm.sessMutex.RLock()
	session, exists := tm.sessions[sessionID]
	tm.sessMutex.RUnlock()

	if !exists || !session.Active {
		return fmt.Errorf("Session %s nicht gefunden oder inaktiv", sessionID)
	}

	// Daten an Proxy-Verbindung schreiben
	session.mutex.Lock()
	defer session.mutex.Unlock()

	if session.ProxyConn == nil {
		return fmt.Errorf("Session %s: keine Proxy-Verbindung", sessionID)
	}

	_, err := session.ProxyConn.Write(payload)
	return err
}

// sendSessionData schickt Daten als Binary-Message an den Agent
func (tm *TunnelManager) sendSessionData(agentID, sessionID string, data []byte) error {
	idBytes := []byte(sessionID)
	if len(idBytes) > 255 {
		return fmt.Errorf("session_id zu lang")
	}

	message := make([]byte, 1+len(idBytes)+len(data))
	message[0] = byte(len(idBytes))
	copy(message[1:], idBytes)
	copy(message[1+len(idBytes):], data)

	return tm.hub.SendBinaryToAgent(agentID, message)
}

// sendSessionDisconnect informiert den Agent dass eine Session beendet ist
func (tm *TunnelManager) sendSessionDisconnect(agentID, sessionID string) {
	if !tm.hub.IsAgentConnected(agentID) {
		return
	}

	cmd := map[string]interface{}{
		"type":    "command",
		"id":      tm.generateID(),
		"command": "tunnel_disconnect",
		"params": map[string]interface{}{
			"session_id": sessionID,
		},
	}

	cmdJSON, _ := json.Marshal(cmd)
	tm.hub.SendToAgent(agentID, cmdJSON)
}

func (tm *TunnelManager) cleanupSession(tunnel *BackendTunnel, sessionID string) {
	tunnel.sessMutex.Lock()
	if sess, ok := tunnel.Sessions[sessionID]; ok {
		sess.Active = false
		delete(tunnel.Sessions, sessionID)
	}
	tunnel.sessMutex.Unlock()

	tm.sessMutex.Lock()
	delete(tm.sessions, sessionID)
	tm.sessMutex.Unlock()
}

// ConfirmSession markiert eine Session als bereit (Agent hat Ziel-Verbindung hergestellt)
func (tm *TunnelManager) ConfirmSession(sessionID string) {
	tm.sessMutex.RLock()
	session, exists := tm.sessions[sessionID]
	tm.sessMutex.RUnlock()

	if exists && session.Active {
		select {
		case <-session.Ready:
			// Bereits geschlossen
		default:
			close(session.Ready)
		}
	}
}

func (tm *TunnelManager) allocatePort() (int, error) {
	tm.mutex.Lock()
	defer tm.mutex.Unlock()

	for port := tm.portStart; port <= tm.portEnd; port++ {
		if !tm.usedPorts[port] {
			tm.usedPorts[port] = true
			return port, nil
		}
	}
	return 0, fmt.Errorf("keine freien Ports verfuegbar")
}

func (tm *TunnelManager) releasePort(port int) {
	tm.mutex.Lock()
	defer tm.mutex.Unlock()
	delete(tm.usedPorts, port)
}

func (tm *TunnelManager) generateID() string {
	bytes := make([]byte, 8)
	rand.Read(bytes)
	return hex.EncodeToString(bytes)
}

// CreateHTTPProxy - HTTP-Reverse-Proxy (noch nicht implementiert)
func (tm *TunnelManager) CreateHTTPProxy(tunnelID string, w http.ResponseWriter, r *http.Request) error {
	http.Error(w, "HTTP-Proxy noch nicht implementiert", http.StatusNotImplemented)
	return nil
}