server.js

import express from 'express';
import { WebSocketServer } from 'ws';
import http from 'http';
import path from 'path';
import { fileURLToPath } from 'url';
import OpenAI from 'openai';
import dotenv from 'dotenv';

dotenv.config();

const __filename = fileURLToPath(import.meta.url);
const __dirname = path.dirname(__filename);

const app = express();
const server = http.createServer(app);
const wss = new WebSocketServer({ server, path: '/ws' });
const PORT = process.env.PORT || 3000;

// Configuration (env-driven, no hard-coded secrets)
const TETRATE_API_KEY = process.env.TETRATE_API_KEY || '';
const TETRATE_BASE_URL = process.env.TETRATE_BASE_URL || 'https://api.router.tetrate.ai/v1';
const API_KEY = process.env.SENTINEL_API_KEY || ''; // Optional API key for external callers

// Initialize Tetrate AI Router (TARS) - OpenAI compatible
const tetrateAI = new OpenAI({
  apiKey: TETRATE_API_KEY,
  baseURL: TETRATE_BASE_URL,
  timeout: 30000,
  maxRetries: 3
});

// Data stores
let alerts = [];
let clients = new Set();
let metrics = {
  totalAlerts: 0,
  highSeverity: 0,
  lastActivity: new Date().toLocaleTimeString(),
  aiCalls: 0,
  aiErrors: 0
};

// Mesh security state (for mesh.html)
let meshMetrics = {
  anomalyScore: 0,
  blockedRequests: 0,
  mtlsViolations: 0
};

let policies = [];

// Live mesh topology model, seeded with a default view
let topology = {
  nodes: [
    { id: 'frontend', type: 'service', namespace: 'web', traffic: 4500 },
    { id: 'api-gateway', type: 'gateway', namespace: 'istio-system', traffic: 8900 },
    { id: 'user-service', type: 'service', namespace: 'backend', traffic: 3200 },
    { id: 'payment-service', type: 'critical', namespace: 'backend', traffic: 1200 },
    { id: 'database', type: 'database', namespace: 'data', traffic: 2100 },
    { id: 'auth-service', type: 'service', namespace: 'security', traffic: 2800 }
  ],
  edges: [
    { source: 'frontend', target: 'api-gateway', protocol: 'HTTP', mtls: true },
    { source: 'api-gateway', target: 'auth-service', protocol: 'HTTP', mtls: true },
    { source: 'api-gateway', target: 'user-service', protocol: 'HTTP', mtls: true },
    { source: 'api-gateway', target: 'payment-service', protocol: 'HTTPS', mtls: true },
    { source: 'user-service', target: 'database', protocol: 'TCP', mtls: false },
    { source: 'payment-service', target: 'database', protocol: 'TCP', mtls: false }
  ]
};

// Middleware
app.use(express.json({ limit: '1mb' })); // basic body size limit
app.use(express.static('public'));

// Simple per-IP rate limiting for AI-heavy endpoints
const rateLimits = new Map(); // key: ip, value: { count, resetAt }
const MAX_REQUESTS_PER_MINUTE = Number(process.env.RATE_LIMIT_PER_MINUTE || 60);

function rateLimit(req, res, next) {
  const ip = req.ip || req.connection.remoteAddress || 'unknown';
  const now = Date.now();
  const windowMs = 60 * 1000;

  let entry = rateLimits.get(ip);
  if (!entry || entry.resetAt < now) {
    entry = { count: 0, resetAt: now + windowMs };
  }
  entry.count += 1;
  rateLimits.set(ip, entry);

  if (entry.count > MAX_REQUESTS_PER_MINUTE) {
    return res.status(429).json({ error: 'Rate limit exceeded. Please slow down.' });
  }

  next();
}

// Optional API key auth middleware (enabled only if SENTINEL_API_KEY is set)
function requireApiKey(req, res, next) {
  if (!API_KEY) return next(); // disabled in local/dev if no key configured
  const headerKey = req.header('x-api-key');
  if (!headerKey || headerKey !== API_KEY) {
    return res.status(401).json({ error: 'Unauthorized' });
  }
  next();
}

// Security: Request logging (metadata only, no keys)
app.use((req, res, next) => {
  console.log(`${new Date().toISOString()} [${req.method}] ${req.path} - Org: default`);
  next();
});

// WebSocket handling
wss.on('connection', (ws, req) => {
  const clientId = Math.random().toString(36).substr(2, 9);
  console.log(`✅ Client connected [${clientId}]`);
  clients.add(ws);
  
  ws.send(JSON.stringify({
    type: 'init',
    alerts: alerts.slice(0, 50),
    metrics: metrics
  }));
  
  ws.on('close', () => {
    console.log(`❌ Client disconnected [${clientId}]`);
    clients.delete(ws);
  });
  
  ws.on('error', (err) => {
    console.error(`WebSocket error [${clientId}]:`, err.message);
    clients.delete(ws);
  });
});

// Broadcast to all connected clients
function broadcast(data) {
  const message = JSON.stringify(data);
  clients.forEach(client => {
    if (client.readyState === 1) {
      client.send(message);
    }
  });
}

// Update metrics and broadcast
function updateMetrics() {
  metrics = {
    totalAlerts: alerts.length,
    highSeverity: alerts.filter(a => a.severity === 'HIGH' || a.severity === 'CRITICAL').length,
    lastActivity: new Date().toLocaleTimeString(),
    aiCalls: metrics.aiCalls,
    aiErrors: metrics.aiErrors
  };
  
  broadcast({
    type: 'metrics',
    metrics: metrics
  });
}

// ---- Mesh topology helpers (Istio-driven) ----

function upsertNode(id, type = 'service', namespace = 'default') {
  if (!id) return;
  let node = topology.nodes.find(n => n.id === id);
  if (!node) {
    node = { id, type, namespace, traffic: 0 };
    topology.nodes.push(node);
  }
  return node;
}

function upsertEdge(source, target, protocol = 'HTTP', mtls = true) {
  if (!source || !target) return;
  let edge = topology.edges.find(
    e => e.source === source && e.target === target
  );
  if (!edge) {
    edge = { source, target, protocol, mtls };
    topology.edges.push(edge);
  } else {
    edge.protocol = protocol || edge.protocol;
    edge.mtls = mtls;
  }
  return edge;
}

// REAL AI Analysis via Tetrate TARS (No direct OpenAI)
async function analyzeWithAI(content, severity) {
  metrics.aiCalls++;
  
  try {
    console.log(`🤖 Calling Tetrate TARS for analysis...`);
    
    const response = await tetrateAI.chat.completions.create({
      model: "gpt-4",  // FIXED: Comma added here
      messages: [
        {
          role: "system",
          content: `You are a cybersecurity expert analyzing Istio service mesh traffic. 
          Analyze the security alert and return ONLY valid JSON with:
          - threatType: specific classification (e.g., "SQL Injection", "Lateral Movement", "Data Exfiltration", "Credential Stuffing")
          - confidence: number 0-100
          - recommendation: specific action (e.g., "Block source IP immediately", "Rate limit destination service", "Investigate manually")
          - indicators: array of evidence indicators`
        },
        {
          role: "user",
          content: `Analyze this service mesh security event:
          Severity: ${severity}
          Content: ${content}
          
          Return ONLY JSON format.`
        }
      ],
      temperature: 0.1,
      max_tokens: 500
    });

    // Parse Tetrate response (best-effort JSON parsing)
    let analysis = {};
    const rawContent = response.choices?.[0]?.message?.content || '';
    try {
      analysis = JSON.parse(rawContent);
    } catch (parseErr) {
      console.warn('⚠️ Failed to parse AI JSON, using defaults. Raw content:', rawContent);
      analysis = {};
    }
    
    console.log(`✅ Tetrate TARS Response: model=${response.model}, tokens=${response.usage?.total_tokens || 'unknown'}`);
    
    return {
      threatType: analysis.threatType || 'Unknown Threat',
      confidence: Math.min(100, Math.max(0, analysis.confidence || 85)),
      recommendation: analysis.recommendation || 'Monitor and investigate',
      indicators: analysis.indicators || ['Pattern match detected'],
      provider: 'Tetrate TARS',
      model: response.model,
      timestamp: new Date().toISOString()
    };
    
  } catch (error) {
    metrics.aiErrors++;
    console.error('❌ Tetrate TARS Error:', error.message);
    
    // Fail secure - require manual review on AI failure
    return {
      threatType: 'AI Router Unavailable - Manual Review Required',
      confidence: 100, // Force high severity on failure
      recommendation: 'Tetrate TARS connection failed - escalate to security team immediately',
      indicators: ['TARS endpoint error', error.message, 'Manual analysis required'],
      provider: 'Tetrate TARS (Failed)',
      error: true,
      timestamp: new Date().toISOString()
    };
  }
}

// API: Analyze new alert (UI-facing, no API key required for now)
app.post('/api/analyze', rateLimit, async (req, res) => {
  try {
    const { alertId, severity, content } = req.body;
    
    if (typeof content !== 'string' || !content.trim() || typeof severity !== 'string') {
      return res.status(400).json({ error: 'Missing required fields' });
    }
    
    console.log(`🚨 New alert received: ${alertId || 'auto-id'} (${severity})`);
    
    // AI Analysis via Tetrate (NO direct OpenAI calls)
    const analysis = await analyzeWithAI(content, severity);
    
    const alert = {
      id: alertId || `ALT-${Date.now()}`,
      severity: severity.toUpperCase(),
      content: content,
      analysis: analysis,
      timestamp: new Date().toLocaleTimeString(),
      createdAt: Date.now()
    };
    
    // Add to beginning of array
    alerts.unshift(alert);
    
    // Keep only last 100
    if (alerts.length > 100) alerts.pop();
    
    // Broadcast to all connected dashboards
    broadcast({
      type: 'new-alert',
      alert: alert
    });
    
    updateMetrics();
    
    console.log(`✅ Alert processed: ${alert.id} | Threat: ${analysis.threatType} | Confidence: ${analysis.confidence}%`);
    res.json(alert);
    
  } catch (error) {
    console.error('Analysis endpoint error:', error);
    res.status(500).json({ error: 'Analysis failed', details: error.message });
  }
});

// API: Tetrate mesh log ingestion endpoint (UI + external)
app.post('/api/tetrate/ingest', rateLimit, async (req, res) => {
  try {
    const { logs, source, cluster } = req.body;
    
    if (!logs || !Array.isArray(logs)) {
      return res.status(400).json({ error: 'Invalid logs format' });
    }
    
    console.log(`📥 Ingesting ${logs.length} logs from ${source || 'unknown'} (cluster: ${cluster || 'unknown'})`);
    
    const processedAlerts = [];
    let topologyChanged = false;
    
    for (const log of logs) {
      if (!log || typeof log !== 'object') continue;

      const responseCode = Number(log.response_code);
      const destinationPort = Number(log.destination_port || 0);
      const method = log.request_method || 'GET';
      const pathValue = log.path || '/';
      const srcWorkload = log.source_workload || 'unknown-src';
      const dstWorkload = log.destination_workload || 'unknown-dst';
      const mtls = Boolean(log.mtls);
      const protocol = log.protocol || 'HTTP';

      // Auto-detect severity based on response codes
      let severity = 'LOW';
      if (responseCode >= 500) severity = 'CRITICAL';
      else if (responseCode >= 400) severity = 'HIGH';
      else if (!mtls && destinationPort !== 443) severity = 'MEDIUM';
      
      // Analyze via Tetrate
      const analysis = await analyzeWithAI(
        `${method} ${pathValue} from ${srcWorkload} to ${dstWorkload} (code: ${responseCode})`,
        severity
      );
      
      if (analysis.confidence > 75 || analysis.error) {
        const id = `MESH-${Date.now()}-${Math.random().toString(36).substr(2, 4)}`;

        const alert = {
          id: id,
          severity: severity,
          content: `Istio traffic: ${srcWorkload} → ${dstWorkload}: ${pathValue}`,
          analysis: analysis,
          timestamp: new Date().toLocaleTimeString(),
          source: 'istio-mesh',
          cluster: cluster
        };
        
        alerts.unshift(alert);
        processedAlerts.push(alert);

        // Update live topology from Istio-style telemetry
        const srcNode = upsertNode(srcWorkload, 'service', log.source_namespace || 'default');
        const dstType =
          dstWorkload.toLowerCase().includes('db') || dstWorkload.toLowerCase().includes('database')
            ? 'database'
            : 'service';
        const dstNode = upsertNode(dstWorkload, dstType, log.destination_namespace || 'default');

        // Bump simple traffic counters
        const count = Number(log.request_count || 1);
        srcNode.traffic = (srcNode.traffic || 0) + count;
        dstNode.traffic = (dstNode.traffic || 0) + count;

        upsertEdge(srcWorkload, dstWorkload, protocol, mtls);
        topologyChanged = true;

        // Update mesh metrics based on this suspicious event
        meshMetrics.anomalyScore = Math.min(
          100,
          meshMetrics.anomalyScore + (analysis.confidence || 10) / 5
        );
        if (!mtls) {
          meshMetrics.mtlsViolations += 1;
        }

        // Broadcast high-level "new threat" event for mesh.html
        const threat = {
          id: id,
          source: log.source_workload,
          destination: log.destination_workload,
          type: analysis.threatType,
          riskScore: analysis.confidence,
          evidence: Array.isArray(analysis.indicators)
            ? analysis.indicators.join(', ')
            : '',
          recommendation: analysis.recommendation,
          autoRemediate: severity === 'CRITICAL' || (analysis.confidence || 0) >= 90
        };

        broadcast({
          type: 'new-threat',
          data: threat,
          metrics: meshMetrics
        });

        // Auto-remediate critical threats by proposing a simple "block" policy
        if (threat.autoRemediate && log.destination_workload) {
          const policy = {
            id: `POL-${Date.now()}-${Math.random().toString(36).substr(2, 4)}`,
            service: log.destination_workload,
            action: 'block',
            reason: `Auto-blocked due to ${analysis.threatType} (confidence ${analysis.confidence}%)`,
            createdAt: new Date().toISOString(),
            source: 'auto-remediation'
          };

          policies.unshift(policy);

          broadcast({
            type: 'policy-applied',
            policy: policy
          });
        }
      }
    }
    
    if (processedAlerts.length > 0) {
      broadcast({
        type: 'batch-alerts',
        alerts: processedAlerts,
        count: processedAlerts.length
      });
      updateMetrics();
    }

    // If mesh topology changed based on these logs, push an update
    if (topologyChanged) {
      broadcast({
        type: 'topology-update',
        topology
      });
    }
    
    res.json({ 
      ingested: logs.length, 
      threatsDetected: processedAlerts.length,
      source: source 
    });
    
  } catch (error) {
    console.error('Ingestion error:', error);
    res.status(500).json({ error: 'Ingestion failed' });
  }
});

// API: Get all alerts
app.get('/api/alerts', (req, res) => {
  res.json(alerts);
});

// API: Get current mesh topology (Istio-driven)
app.get('/api/topology', (req, res) => {
  res.json(topology);
});

// API: Delete alert
app.delete('/api/alerts/:id', (req, res) => {
  const id = req.params.id;
  alerts = alerts.filter(a => a.id !== id);
  updateMetrics();
  broadcast({ type: 'alert-deleted', id: id });
  res.json({ success: true });
});

// API: Clear all
app.delete('/api/alerts', (req, res) => {
  alerts = [];
  updateMetrics();
  broadcast({ type: 'alerts-cleared' });
  res.json({ success: true });
});

// API: Recommend an Istio/Tetrate policy for a given alert (YAML)
app.post('/api/policies/recommend', rateLimit, async (req, res) => {
  try {
    const { alertId, alert: alertOverride } = req.body || {};

    if (!alertId) {
      return res.status(400).json({ error: 'alertId is required' });
    }

    const alert = alertOverride || alerts.find(a => a.id === alertId);
    if (!alert) {
      return res.status(404).json({ error: 'Alert not found' });
    }

    // Ask Tetrate AI Router to generate a concrete policy YAML
    console.log(`🛠️ Generating policy recommendation for alert ${alertId}...`);

    const response = await tetrateAI.chat.completions.create({
      model: 'gpt-4',
      messages: [
        {
          role: 'system',
          content:
            'You are an Istio and Tetrate service mesh security engineer. ' +
            'Given a Sentinel AI security alert (JSON), generate a single, production-ready ' +
            'Kubernetes YAML policy for Istio (e.g., AuthorizationPolicy) or equivalent Tetrate policy ' +
            'that mitigates the described threat. ' +
            'Use best practices, keep the policy minimal but safe. ' +
            'Output ONLY the YAML, with no explanations or backticks.'
        },
        {
          role: 'user',
          content:
            'Here is the alert JSON:\n\n' +
            JSON.stringify(alert, null, 2) +
            '\n\nGenerate one YAML policy that best mitigates this threat.'
        }
      ],
      temperature: 0.2,
      max_tokens: 600
    });

    const yaml = response.choices?.[0]?.message?.content?.trim() || '';

    if (!yaml) {
      return res.status(500).json({ error: 'Policy generation returned empty response' });
    }

    console.log(`✅ Policy YAML generated for alert ${alertId}`);

    res.json({
      alertId,
      policy: yaml,
      model: response.model
    });
  } catch (error) {
    console.error('Policy recommend error:', error);
    res.status(500).json({ error: 'Failed to generate policy', details: error.message });
  }
});

// API: Apply mesh/network security policy (used by unified console)
app.post('/api/policy/apply', rateLimit, (req, res) => {
  try {
    const { service, action, reason } = req.body || {};

    if (!service || !action) {
      return res.status(400).json({ error: 'service and action are required' });
    }

    const policy = {
      id: `POL-${Date.now()}-${Math.random().toString(36).substr(2, 4)}`,
      service,
      action,
      reason: reason || 'Policy applied via Sentinel mesh UI',
      createdAt: new Date().toISOString(),
      source: 'manual'
    };

    policies.unshift(policy);

    // Notify all mesh dashboards
    broadcast({
      type: 'policy-applied',
      policy
    });

    res.json(policy);
  } catch (error) {
    console.error('Policy apply error:', error);
    res.status(500).json({ error: 'Failed to apply policy' });
  }
});

// API: Get metrics
app.get('/api/metrics', (req, res) => {
  res.json(metrics);
});

// Health check
app.get('/api/health', (req, res) => {
  res.json({ 
    status: 'healthy',
    provider: 'Tetrate TARS',
    clients: clients.size,
    uptime: process.uptime(),
    aiCalls: metrics.aiCalls,
    aiErrors: metrics.aiErrors
  });
});

// Test Tetrate connection on startup
async function testTetrateConnection() {
  try {
    console.log('🔍 Testing Tetrate TARS connection...');
    const response = await tetrateAI.models.list();
    console.log(`✅ Tetrate TARS Connected: ${response.data.length} models available`);
    return true;
  } catch (error) {
    console.error('❌ Tetrate TARS Connection Failed:', error.message);
    console.log('⚠️  AI analysis will return fallback responses until connection restored');
    return false;
  }
}

server.listen(PORT, async () => {
  console.log(`🛡️  Sentinel AI (Tetrate Edition) on port ${PORT}`);
  console.log(`📊  Dashboard: http://localhost:${PORT}/console.html`);
  console.log(`🔌  WebSocket: ws://localhost:${PORT}/ws`);
  console.log(`🤖  AI Provider: Tetrate TARS (https://api.router.tetrate.ai)`);
  
  // Test connection
  await testTetrateConnection();
});