wraith/internal/rdp/service.go

package rdp

import (
	"fmt"
	"sync"
	"time"

	"github.com/google/uuid"
)

// RDPSession represents an active RDP connection with its associated state.
type RDPSession struct {
	ID        string    `json:"id"`
	Config    RDPConfig `json:"config"`
	Backend   RDPBackend `json:"-"`
	Buffer    *PixelBuffer `json:"-"`
	ConnID    int64     `json:"connectionId"`
	Connected time.Time `json:"connected"`
}

// RDPService manages multiple RDP sessions. It uses a backend factory to
// create new RDPBackend instances — during development the factory returns
// MockBackend instances; in production it will return FreeRDP-backed ones.
type RDPService struct {
	sessions       map[string]*RDPSession
	mu             sync.RWMutex
	backendFactory func() RDPBackend
}

// NewRDPService creates a new service with the given backend factory.
// The factory is called once per Connect to create a fresh backend.
func NewRDPService(factory func() RDPBackend) *RDPService {
	return &RDPService{
		sessions:       make(map[string]*RDPSession),
		backendFactory: factory,
	}
}

// Connect creates a new RDP session using the provided configuration.
// Returns the session ID on success.
func (s *RDPService) Connect(config RDPConfig, connectionID int64) (string, error) {
	// Apply defaults
	if config.Port <= 0 {
		config.Port = 3389
	}
	if config.Width <= 0 {
		config.Width = 1920
	}
	if config.Height <= 0 {
		config.Height = 1080
	}
	if config.ColorDepth <= 0 {
		config.ColorDepth = 32
	}
	if config.Security == "" {
		config.Security = "nla"
	}

	backend := s.backendFactory()
	if err := backend.Connect(config); err != nil {
		return "", fmt.Errorf("rdp connect to %s:%d: %w", config.Hostname, config.Port, err)
	}

	sessionID := uuid.NewString()
	session := &RDPSession{
		ID:        sessionID,
		Config:    config,
		Backend:   backend,
		Buffer:    NewPixelBuffer(config.Width, config.Height),
		ConnID:    connectionID,
		Connected: time.Now(),
	}

	s.mu.Lock()
	s.sessions[sessionID] = session
	s.mu.Unlock()

	return sessionID, nil
}

// Disconnect tears down the RDP session and removes it from tracking.
func (s *RDPService) Disconnect(sessionID string) error {
	s.mu.Lock()
	session, ok := s.sessions[sessionID]
	if !ok {
		s.mu.Unlock()
		return fmt.Errorf("session %s not found", sessionID)
	}
	delete(s.sessions, sessionID)
	s.mu.Unlock()

	return session.Backend.Disconnect()
}

// SendMouse forwards a mouse event to the session's backend.
func (s *RDPService) SendMouse(sessionID string, x, y int, flags uint32) error {
	session, err := s.getSession(sessionID)
	if err != nil {
		return err
	}
	return session.Backend.SendMouseEvent(x, y, flags)
}

// SendKey forwards a key event to the session's backend.
func (s *RDPService) SendKey(sessionID string, scancode uint32, pressed bool) error {
	session, err := s.getSession(sessionID)
	if err != nil {
		return err
	}
	return session.Backend.SendKeyEvent(scancode, pressed)
}

// SendClipboard forwards clipboard text to the session's backend.
func (s *RDPService) SendClipboard(sessionID string, data string) error {
	session, err := s.getSession(sessionID)
	if err != nil {
		return err
	}
	return session.Backend.SendClipboard(data)
}

// GetFrame returns the current RGBA pixel buffer for the given session.
func (s *RDPService) GetFrame(sessionID string) ([]byte, error) {
	session, err := s.getSession(sessionID)
	if err != nil {
		return nil, err
	}
	return session.Backend.GetFrame()
}

// GetSessionInfo returns the session metadata without the backend or buffer.
// This is safe to serialize to JSON for the frontend.
func (s *RDPService) GetSessionInfo(sessionID string) (*RDPSession, error) {
	session, err := s.getSession(sessionID)
	if err != nil {
		return nil, err
	}
	return session, nil
}

// ListSessions returns all active RDP sessions.
func (s *RDPService) ListSessions() []*RDPSession {
	s.mu.RLock()
	defer s.mu.RUnlock()

	list := make([]*RDPSession, 0, len(s.sessions))
	for _, sess := range s.sessions {
		list = append(list, sess)
	}
	return list
}

// IsConnected returns whether a specific session is still connected.
func (s *RDPService) IsConnected(sessionID string) bool {
	session, err := s.getSession(sessionID)
	if err != nil {
		return false
	}
	return session.Backend.IsConnected()
}

// getSession is an internal helper that retrieves a session by ID.
func (s *RDPService) getSession(sessionID string) (*RDPSession, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	session, ok := s.sessions[sessionID]
	if !ok {
		return nil, fmt.Errorf("session %s not found", sessionID)
	}
	return session, nil
}