mitm-webserver/frontend/src/components/PacketViewer.tsx

// src/components/PacketViewer.tsx
import { DeleteOutlined, DownloadOutlined, SyncOutlined } from '@ant-design/icons';
import {
  Button,
  Card,
  Col,
  Modal,
  Row,
  Select,
  Space,
  Spin,
  Switch,
  Table,
  Tag,
  Tooltip,
  Typography,
  message,
} from 'antd';
import { ReactElement, useCallback, useEffect, useMemo, useRef, useState } from 'react';
import { clearPackets, fetchPackets } from '../api/apiClient';
import type { PacketRow } from '../types/packets';

const { Text, Title } = Typography;
const { Option } = Select;

const DEFAULT_LIMIT = 200;
const MAX_PACKETS = 2000; // in-memory cap

function base64ToHex(b64: string) {
  try {
    const bin = atob(b64);
    const bytes = new Uint8Array(bin.length);
    for (let i = 0; i < bin.length; i++) bytes[i] = bin.charCodeAt(i);
    return Array.from(bytes)
      .map((b) => b.toString(16).padStart(2, '0'))
      .join(' ');
  } catch {
    return '(invalid base64)';
  }
}

function base64ToBlob(b64: string) {
  const bin = atob(b64);
  const len = bin.length;
  const arr = new Uint8Array(len);
  for (let i = 0; i < len; i++) arr[i] = bin.charCodeAt(i);
  return new Blob([arr.buffer], { type: 'application/octet-stream' });
}

/**
 * Improved deterministic color generation
 *
 * - returns { background: hsl(...), color: '#000' | '#fff' }
 * - uses cached Map so results are stable/fast
 */
const _colorCache = new Map<string, { background: string; color: string }>();
function colorForName(name: string) {
  const key = String(name);
  const cached = _colorCache.get(key);
  if (cached) return cached;

  // simple hash to produce a hue in [0, 359]
  let h = 2166136261 >>> 0; // FNV offset basis
  for (let i = 0; i < key.length; i++) {
    h ^= key.charCodeAt(i);
    h = Math.imul(h, 16777619) >>> 0; // FNV prime
  }
  const hue = h % 360;

  // Choose saturation and lightness to ensure good contrast
  // We vary lightness slightly by name so colors differ more
  const sat = 68; // %
  // map part of hash to lightness in range [38, 58]
  const light = 38 + (h % 21); // 38..58

  // Use comma syntax for compatibility
  const background = `hsl(${hue}, ${sat}%, ${light}%)`;

  // Choose text color for contrast: use white for darker backgrounds
  const textColor = light < 50 ? '#ffffff' : '#0b0b0b';

  const result = { background, color: textColor };
  _colorCache.set(key, result);
  return result;
}

function packetKey(packet: PacketRow) {
  return String(packet.packet_uid ?? packet.id ?? `${packet.iface ?? 'if'}:${packet.timestamp ?? ''}`);
}

function parseIfaces(ifaceField?: string | string[] | null, observedIfaces?: string[] | null) {
  if (Array.isArray(observedIfaces) && observedIfaces.length > 0) return observedIfaces;
  if (!ifaceField) return [];
  if (Array.isArray(ifaceField)) return ifaceField;
  return String(ifaceField)
    .split(/[,\|;]+/)
    .map((s) => s.trim())
    .filter(Boolean);
}

/** Format timestamp to European locale with ms */
function formatTimestamp(ts?: string) {
  if (!ts) return '-';
  try {
    const d = new Date(ts);
    if (Number.isNaN(d.getTime())) return String(ts);
    return (
      d.toLocaleString('de-DE', {
        year: 'numeric',
        month: '2-digit',
        day: '2-digit',
        hour: '2-digit',
        minute: '2-digit',
        second: '2-digit',
      }) + `.${String(d.getMilliseconds()).padStart(3, '0')}`
    );
  } catch {
    return String(ts);
  }
}

/**
 * PacketViewer component
 */
export default function PacketViewer(): ReactElement {
  const [packets, setPackets] = useState<PacketRow[]>([]);
  const [loading, setLoading] = useState<boolean>(true);
  const [statusLoading, setStatusLoading] = useState<boolean>(false);
  const [limit, setLimit] = useState<number>(DEFAULT_LIMIT);
  const [paused, setPaused] = useState<boolean>(false);
  const wsRef = useRef<WebSocket | null>(null);
  const [hexModalOpen, setHexModalOpen] = useState(false);
  const [rawHex, setRawHex] = useState<string | null>(null);
  const [rawB64ForDownload, setRawB64ForDownload] = useState<string | null>(null);
  const [selectedPacket, setSelectedPacket] = useState<PacketRow | null>(null);
  const queuedDuringPause = useRef<PacketRow[]>([]);

  const pushNew = useCallback((newPkts: PacketRow[] | PacketRow) => {
    setPackets((prev) => {
      const arr = Array.isArray(newPkts) ? newPkts : [newPkts];
      const next = [...prev];

      for (const packet of arr) {
        const key = packetKey(packet);
        const index = next.findIndex((item) => packetKey(item) === key);
        if (index >= 0) {
          next[index] = { ...next[index], ...packet };
        } else {
          next.unshift(packet);
        }
      }

      next.sort((left, right) => {
        const leftId = Number(left.id ?? 0);
        const rightId = Number(right.id ?? 0);
        if (!Number.isNaN(leftId) && !Number.isNaN(rightId) && leftId !== rightId) {
          return rightId - leftId;
        }
        return String(right.timestamp ?? '').localeCompare(String(left.timestamp ?? ''));
      });

      if (next.length > MAX_PACKETS) return next.slice(0, MAX_PACKETS);
      return next;
    });
  }, []);

  // fetch historical packets
  const fetchHistory = useCallback(async (limitVal: number) => {
    setStatusLoading(true);
    try {
      const res = await fetchPackets(limitVal);
      const list: PacketRow[] = res.packets ?? [];
      setPackets(list);
    } catch (err: any) {
      console.error('fetchHistory error', err);
      message.error(err?.message ?? 'Failed to fetch packet history');
    } finally {
      setStatusLoading(false);
      setLoading(false);
    }
  }, []);

  // open websocket
  const openWs = useCallback(() => {
    if (wsRef.current) {
      try {
        wsRef.current.close();
      } catch {}
      wsRef.current = null;
    }

    const loc = window.location;
    const protocol = loc.protocol === 'https:' ? 'wss' : 'ws';
    const wsUrl = `ws://mitm.lan/api/packets/ws/packets?subscribe_recent=20`;
    const ws = new WebSocket(wsUrl);
    wsRef.current = ws;

    ws.onopen = () => {};

    ws.onmessage = (ev) => {
      try {
        const data = JSON.parse(ev.data);
        if (data && data.type === 'recent' && Array.isArray(data.packets)) {
          if (paused) {
            queuedDuringPause.current.unshift(...(data.packets as PacketRow[]));
          } else {
            pushNew(data.packets as PacketRow[]);
          }
        } else if (data && (data.iface || data.raw_b64 || data.id || data.timestamp)) {
          if (paused) {
            queuedDuringPause.current.unshift(data as PacketRow);
          } else {
            pushNew(data as PacketRow);
          }
        } else {
          console.debug('ws unknown message', data);
        }
      } catch (err) {
        console.error('ws message parse error', err);
      }
    };

    ws.onerror = (e) => {
      console.error('ws error', e);
    };

    ws.onclose = () => {
      wsRef.current = null;
    };
  }, [paused, pushNew]);

  // pause handling: when unpausing, flush queuedDuringPause into list
  useEffect(() => {
    if (!paused) {
      const queued = queuedDuringPause.current;
      if (queued.length > 0) {
        pushNew(queued);
        queuedDuringPause.current = [];
      }
    }
  }, [paused, pushNew]);

  // start up: fetch history and open ws
  useEffect(() => {
    setLoading(true);
    fetchHistory(limit).then(() => {
      openWs();
    });
    return () => {
      if (wsRef.current) {
        try {
          wsRef.current.close();
        } catch {}
        wsRef.current = null;
      }
    };
  }, [fetchHistory, limit, openWs]);

  const handleRefresh = async () => {
    setLoading(true);
    try {
      await fetchHistory(limit);
    } finally {
      setLoading(false);
    }
  };

  const handlePauseToggle = (val: boolean) => {
    setPaused(val);
  };

  const openRawModal = (row: PacketRow) => {
    if (!row.raw_b64) {
      message.info('No raw bytes available for this packet');
      return;
    }
    setRawHex(base64ToHex(row.raw_b64));
    setRawB64ForDownload(row.raw_b64 ?? null);
    setSelectedPacket(row);
    setHexModalOpen(true);
  };

  const downloadRaw = () => {
    if (!rawB64ForDownload) return;
    const blob = base64ToBlob(rawB64ForDownload);
    const url = URL.createObjectURL(blob);
    const a = document.createElement('a');
    a.href = url;
    a.download = `packet_${selectedPacket?.id ?? 'pkt'}.bin`;
    a.click();
    URL.revokeObjectURL(url);
  };

  const columns = useMemo(
    () => [
      {
        title: 'Time',
        dataIndex: 'timestamp',
        key: 'timestamp',
        width: 120,
        render: (val: any) => <Text>{formatTimestamp(val)}</Text>,
      },
      {
        title: 'Ifaces',
        dataIndex: 'iface',
        key: 'iface',
        width: 100,
        render: (_: any, rec: PacketRow) => {
          const ifaces = parseIfaces(rec.iface, rec.observed_ifaces);
          if (ifaces.length === 0) return <Text type="secondary">-</Text>;
          return (
            <Space wrap size={[6, 6]}>
              {ifaces.map((name) => {
                const { background, color } = colorForName(name);
                // tag style: use background and computed text color for accessibility
                return (
                  <Tooltip key={name} title={name}>
                    <Tag style={{ background, color, fontWeight: 600, borderRadius: 6, fontSize: 15 }}>{name}</Tag>
                  </Tooltip>
                );
              })}
            </Space>
          );
        },
      },
      {
        title: 'Path',
        key: 'path',
        width: 150,
        render: (_: any, rec: PacketRow) => (
          <Space direction="vertical" size={0}>
            <Text>{rec.ingress_if ?? '-'}</Text>
            <Text type="secondary">{rec.egress_if ?? '-'}</Text>
          </Space>
        ),
      },
      {
        title: 'Verdict',
        key: 'verdict',
        width: 120,
        render: (_: any, rec: PacketRow) => {
          const verdict = rec.verdict ?? '-';
          const confidence = rec.verdict_confidence ? ` (${rec.verdict_confidence})` : '';
          return (
            <Space direction="vertical" size={0}>
              <Text>{verdict}</Text>
              <Text type="secondary" style={{ fontSize: 12 }}>
                {(rec.verdict_reason ?? '') + confidence}
              </Text>
            </Space>
          );
        },
      },
      {
        title: 'MAC (src → dst)',
        key: 'macs',
        width: 180,
        render: (_: any, rec: PacketRow) => (
          <div>
            <div>
              <Text>{rec.src_mac ?? '—'}</Text>
            </div>
            <div style={{ marginTop: 6 }}>
              <Text>{rec.dst_mac ?? '—'}</Text>
            </div>
          </div>
        ),
      },
      {
        title: 'Protocol',
        dataIndex: 'ip_proto',
        key: 'ip_proto',
        width: 110,
        render: (v: any) => <Text>{v ?? '-'}</Text>,
      },
      {
        title: 'App',
        key: 'app',
        width: 170,
        render: (_: any, rec: PacketRow) => {
          const app = rec.app_protocol ?? rec.app_master_protocol ?? '-';
          return (
            <Space direction="vertical" size={0}>
              <Text>{app}</Text>
              <Text type="secondary" style={{ fontSize: 12 }}>
                {rec.app_confidence ?? rec.app_category ?? ''}
              </Text>
            </Space>
          );
        },
      },
      {
        title: 'Host',
        dataIndex: 'app_hostname',
        key: 'app_hostname',
        width: 180,
        render: (v: any) => <Text>{v ?? '-'}</Text>,
      },
      {
        title: 'Src IP:Port',
        key: 'src',
        width: 180,
        render: (_: any, rec: PacketRow) =>
          rec.src_ip ? (
            <Text copyable={{ text: `${rec.src_ip}${rec.src_port ? `:${rec.src_port}` : ''}` }}>
              {rec.src_ip}
              {rec.src_port ? `:${rec.src_port}` : ''}
            </Text>
          ) : (
            <Text type="secondary">-</Text>
          ),
      },
      {
        title: 'Dst IP:Port',
        key: 'dst',
        width: 180,
        render: (_: any, rec: PacketRow) =>
          rec.dst_ip ? (
            <Text copyable={{ text: `${rec.dst_ip}${rec.dst_port ? `:${rec.dst_port}` : ''}` }}>
              {rec.dst_ip}
              {rec.dst_port ? `:${rec.dst_port}` : ''}
            </Text>
          ) : (
            <Text type="secondary">-</Text>
          ),
      },
      {
        title: 'Len',
        dataIndex: 'length',
        key: 'length',
        width: 80,
        render: (n: any) => (typeof n === 'number' ? n.toLocaleString('de-DE') : '-'),
      },
      {
        title: 'Risk',
        dataIndex: 'app_risk_score',
        key: 'app_risk_score',
        width: 70,
        render: (n: any) => (typeof n === 'number' ? n : '-'),
      },
      {
        title: 'Actions',
        key: 'actions',
        width: 180,
        render: (_: any, rec: PacketRow) => (
          <Space>
            <Button size="small" onClick={() => openRawModal(rec)}>
              View raw
            </Button>
            {rec.raw_b64 && (
              <Button
                size="small"
                icon={<DownloadOutlined />}
                onClick={() => {
                  const blob = base64ToBlob(rec.raw_b64!);
                  const url = URL.createObjectURL(blob);
                  const a = document.createElement('a');
                  a.href = url;
                  a.download = `packet_${rec.id ?? 'pkt'}.bin`;
                  a.click();
                  URL.revokeObjectURL(url);
                }}
              >
                Download
              </Button>
            )}
          </Space>
        ),
      },
    ],
    [],
  );

  return (
    <div>
      <Row justify="space-between" align="middle" style={{ marginBottom: 12 }}>
        <Col>
          <Title level={3}>Packets</Title>
          <div>
            <Text type="secondary">Live packet viewer — history + live stream</Text>
          </div>
        </Col>

        <Col>
          <Space>
            <Text>History</Text>
            <Select value={limit} onChange={(v) => setLimit(Number(v))} style={{ width: 120 }}>
              <Option value={50}>50</Option>
              <Option value={100}>100</Option>
              <Option value={200}>200</Option>
              <Option value={500}>500</Option>
            </Select>

            <Text>Live</Text>
            <Switch checked={!paused} onChange={(checked) => handlePauseToggle(!checked ? true : false)} />

            <Button icon={<SyncOutlined />} onClick={handleRefresh} loading={statusLoading} />
            <Button
              onClick={async () => {
                await clearPackets();
                setPackets([]);
              }}
              danger
              icon={<DeleteOutlined />}
            />
          </Space>
        </Col>
      </Row>

      <Card>
        {loading ? (
          <div style={{ textAlign: 'center', padding: 40 }}>
            <Spin />
          </div>
        ) : packets.length === 0 ? (
          <div style={{ padding: 12 }}>
            <Text type="secondary">No packets captured yet.</Text>
          </div>
        ) : (
          <Table
            rowKey={(rec: PacketRow) => packetKey(rec)}
            dataSource={packets}
            columns={columns as any}
            pagination={{ pageSize: 25 }}
            size="small"
            bordered
          />
        )}
      </Card>

      <Modal
        title={`Raw packet ${selectedPacket?.id ?? ''}`}
        open={hexModalOpen}
        onCancel={() => setHexModalOpen(false)}
        footer={
          <Space>
            <Button onClick={() => setHexModalOpen(false)}>Close</Button>
            <Button icon={<DownloadOutlined />} onClick={downloadRaw} type="primary">
              Download raw
            </Button>
          </Space>
        }
        width={900}
      >
        <pre style={{ whiteSpace: 'pre-wrap', wordBreak: 'break-word', fontSize: 12 }}>{rawHex}</pre>
      </Modal>
    </div>
  );
}