ShieldAI/services/spamshield/src/websocket/alert-server.ts

import { WebSocketServer, WebSocket } from 'ws';
import { createHash } from 'crypto';
import { DecisionResult } from '../engine/decision-engine';

export interface AlertEvent {
  type: 'decision' | 'flag' | 'block' | 'user_feedback' | 'carrier_action';
  data: {
    phoneNumber: string;
    phoneNumberHash?: string;
    decision?: 'BLOCK' | 'FLAG' | 'ALLOW';
    confidence?: number;
    ruleMatches?: string[];
    carrierAction?: string;
    timestamp: Date;
    metadata?: Record<string, any>;
  };
}

export interface ClientSubscription {
  clientId: string;
  subscribedEvents: string[];
  connectedAt: Date;
  lastActivity: Date;
  ws?: WebSocket;
}

export interface AlertServerConfig {
  port?: number;
  host?: string;
  heartbeatIntervalMs?: number;
  maxClients?: number;
  enableLogging?: boolean;
  enableAuth?: boolean;
  jwtSecret?: string;
  allowedOrigins?: string[];
}

const DEFAULT_CONFIG: Required<AlertServerConfig> = {
  port: 8080,
  host: '0.0.0.0',
  heartbeatIntervalMs: 30000,
  maxClients: 100,
  enableLogging: true,
  enableAuth: true,
  jwtSecret: process.env.SPAMSHIELD_JWT_SECRET || '',
  allowedOrigins: ['http://localhost:3000'],
};

export class AlertServer {
  private readonly config: Required<AlertServerConfig>;
  private readonly wss: WebSocketServer;
  private readonly clients: Map<string, ClientSubscription> = new Map();
  private heartbeatInterval?: NodeJS.Timeout;
  private isRunning = false;

  constructor(config?: AlertServerConfig) {
    this.config = { ...DEFAULT_CONFIG, ...config };
    this.wss = new WebSocketServer({
      port: this.config.port,
      host: this.config.host,
    });
    
    this.setupWebSocketHandlers();
  }

  private setupWebSocketHandlers(): void {
    this.wss.on('connection', async (ws: WebSocket, req: any) => {
      const origin = req.headers.origin;
      if (origin && this.config.allowedOrigins.length > 0 && !this.config.allowedOrigins.includes(origin)) {
        ws.close(1008, 'Origin not allowed');
        return;
      }

      if (this.config.enableAuth && this.config.jwtSecret) {
        const authHeader = req.headers.authorization;
        if (!authHeader || !authHeader.startsWith('Bearer ')) {
          ws.close(4001, 'Missing or invalid JWT token');
          return;
        }
        const token = authHeader.substring(7);
        const valid = await this.verifyJWT(token);
        if (!valid) {
          ws.close(4002, 'Invalid or expired JWT token');
          return;
        }
      }

      if (this.clients.size >= this.config.maxClients) {
        ws.close(1013, 'Too many clients');
        return;
      }

      const clientId = req.headers['x-client-id'] as string || `client-${Date.now()}-${Math.random()}`;

      const subscription: ClientSubscription = {
        clientId,
        subscribedEvents: ['decision', 'flag', 'block', 'user_feedback', 'carrier_action'],
        connectedAt: new Date(),
        lastActivity: new Date(),
        ws,
      };

      this.clients.set(clientId, subscription);

      ws.on('message', (data: Buffer) => {
        try {
          const message = JSON.parse(data.toString()) as { eventTypes?: string[] };
          if (message.eventTypes) {
            subscription.subscribedEvents = message.eventTypes;
          }
          subscription.lastActivity = new Date();
        } catch (error) {
          console.error('[AlertServer] Error parsing client message:', error);
        }
      });

      ws.on('close', () => {
        this.clients.delete(clientId);
        if (this.config.enableLogging) {
          console.log(`[AlertServer] Client ${clientId} disconnected. Active clients: ${this.clients.size}`);
        }
      });

      ws.on('error', (error: Error) => {
        console.error(`[AlertServer] WebSocket error for client ${clientId}:`, error);
      });

      ws.send(JSON.stringify({
        type: 'connected',
        data: {
          clientId,
          subscribedEvents: subscription.subscribedEvents,
          connectedAt: subscription.connectedAt,
        },
      }));

      if (this.config.enableLogging) {
        console.log(`[AlertServer] Client ${clientId} connected. Total clients: ${this.clients.size}`);
      }
    });

    this.wss.on('error', (error: Error) => {
      console.error('[AlertServer] Server error:', error);
    });
  }

  async broadcastDecision(phoneNumber: string, decision: DecisionResult): Promise<void> {
    const event: AlertEvent = {
      type: 'decision',
      data: {
        phoneNumber,
        phoneNumberHash: this.hashPhoneNumber(phoneNumber),
        decision: decision.decision,
        confidence: decision.confidence,
        ruleMatches: decision.reasons,
        timestamp: decision.executedAt,
        metadata: {
          scoring: decision.scoring,
        },
      },
    };

    await this.broadcast(event, ['decision']);
  }

  async broadcastBlock(phoneNumber: string, callSid: string): Promise<void> {
    const event: AlertEvent = {
      type: 'block',
      data: {
        phoneNumber,
        timestamp: new Date(),
        metadata: {
          callSid,
          action: 'carrier_block',
        },
      },
    };

    await this.broadcast(event, ['block', 'carrier_action']);
  }

  async broadcastFlag(phoneNumber: string, reasons: string[]): Promise<void> {
    const event: AlertEvent = {
      type: 'flag',
      data: {
        phoneNumber,
        timestamp: new Date(),
        metadata: {
          reasons,
        },
      },
    };

    await this.broadcast(event, ['flag']);
  }

  async broadcastUserFeedback(
    phoneNumber: string,
    isSpam: boolean,
    userId: string
  ): Promise<void> {
    const event: AlertEvent = {
      type: 'user_feedback',
      data: {
        phoneNumber,
        timestamp: new Date(),
        metadata: {
          isSpam,
          userId,
        },
      },
    };

    await this.broadcast(event, ['user_feedback']);
  }

  private async broadcast(event: AlertEvent, eventTypes: string[]): Promise<void> {
    const eventData = JSON.stringify(event);
    const now = new Date();

    for (const [clientId, subscription] of this.clients.entries()) {
      const shouldSend = subscription.subscribedEvents.some(et => eventTypes.includes(et));
      
      if (shouldSend && subscription.ws?.readyState === WebSocket.OPEN) {
        try {
          subscription.ws.send(eventData);
          subscription.lastActivity = now;
        } catch (error) {
          if (this.config.enableLogging) {
            console.error(`[AlertServer] Failed to send to client ${clientId}:`, error);
          }
        }
      }
    }
  }

  subscribe(clientId: string, eventTypes: string[]): void {
    const subscription = this.clients.get(clientId);
    if (subscription) {
      subscription.subscribedEvents = eventTypes;
      subscription.lastActivity = new Date();
    }
  }

  unsubscribe(clientId: string): void {
    this.clients.delete(clientId);
    if (this.config.enableLogging) {
      console.log(`[AlertServer] Client ${clientId} unsubscribed. Active clients: ${this.clients.size}`);
    }
  }

  getClientCount(): number {
    return this.clients.size;
  }

  getActiveClients(): Array<{ clientId: string; subscribedEvents: string[]; connectedAt: Date }> {
    return Array.from(this.clients.values()).map(({ clientId, subscribedEvents, connectedAt }) => ({
      clientId,
      subscribedEvents,
      connectedAt,
    }));
  }

  startHeartbeat(): void {
    this.heartbeatInterval = setInterval(() => {
      const heartbeat: AlertEvent = {
        type: 'decision',
        data: {
          phoneNumber: '',
          timestamp: new Date(),
          metadata: {
            heartbeat: true,
            activeClients: this.clients.size,
          },
        },
      };

      const eventData = JSON.stringify(heartbeat);
      for (const subscription of this.clients.values()) {
        if (subscription.subscribedEvents.includes('decision') && subscription.ws?.readyState === WebSocket.OPEN) {
          subscription.ws.send(eventData);
        }
      }
    }, this.config.heartbeatIntervalMs);

    this.isRunning = true;
  }

  stopHeartbeat(): void {
    if (this.heartbeatInterval) {
      clearInterval(this.heartbeatInterval);
      this.heartbeatInterval = undefined;
    }
    this.isRunning = false;
  }

  async shutdown(): Promise<void> {
    this.stopHeartbeat();
    
    return new Promise((resolve) => {
      this.wss.close(() => {
        for (const subscription of this.clients.values()) {
          subscription.ws?.terminate();
        }
        this.clients.clear();
        resolve();
      });
    });
  }

  getConfig(): Required<AlertServerConfig> {
    return { ...this.config };
  }

  private hashPhoneNumber(phoneNumber: string): string {
    return createHash('sha256').update(phoneNumber).digest('hex');
  }

  private async verifyJWT(token: string): Promise<boolean> {
    try {
      const { jwtVerify } = await import('jose');
      await jwtVerify(token, new TextEncoder().encode(this.config.jwtSecret), {
        algorithms: ['HS256'],
      });
      return true;
    } catch {
      if (this.config.enableLogging) {
        console.log('[AlertServer] JWT verification failed');
      }
      return false;
    }
  }
}