Assessor_2.0/src/features/mitre-graph/mitre-graph.ts

import type { FindingRecord, Severity } from '../../domain/finding';
import { mitreBadgeLabel, normalizeSeverity } from '../../profiles/purple-knight/presentation';
import { cisBenchmarkCategories } from '../../profiles/cis-cat/cis-benchmark-hierarchy';
import { recordTitleFor } from '../dataset-viewer/record-title';

export interface MitreTechnique {
  id: string;
  name: string;
  tactics: string[];
  platforms: string[];
  parent: string | null;
  url: string;
  sourceCount: number;
}

export interface MitreEntity {
  id: string;
  name: string;
  type: 'group' | 'campaign' | 'software';
  techniques: string[];
}

export interface MitreGraphData {
  version: string;
  tacticOrder: string[];
  techniques: MitreTechnique[];
  entities: MitreEntity[];
}

export interface FindingMitreFocus {
  techniqueIds: Set<string>;
  tactics: Set<string>;
  findingCount: number;
}

export interface FindingPathItem {
  id: string;
  record: FindingRecord;
  title: string;
  category: string;
  categories: string[];
  defenses: string[];
  targets: string[];
  tactics: string[];
  primaryTactic: string;
  techniqueIds: string[];
  severity: Severity;
  weight: number;
  score: number;
  risk: number;
  mappingOrigin: 'native' | 'derived' | 'unmapped';
}

const STORAGE_KEY = 'assessor_custom_mitre_matrix';

export function validateMitreMatrixData(raw: unknown): MitreGraphData {
  if (!raw || typeof raw !== 'object') throw new Error('Invalid MITRE matrix format: expected a JSON object.');
  const data = raw as Partial<MitreGraphData>;
  if (!Array.isArray(data.tacticOrder) || !data.tacticOrder.length) throw new Error('Invalid MITRE matrix: "tacticOrder" array is missing or empty.');
  if (!Array.isArray(data.techniques)) throw new Error('Invalid MITRE matrix: "techniques" array is missing.');
  return {
    version: String(data.version ?? 'Updated'),
    tacticOrder: data.tacticOrder.map(String),
    techniques: data.techniques.map(t => ({
      id: String(t.id),
      name: String(t.name),
      tactics: Array.isArray(t.tactics) ? t.tactics.map(String) : [],
      platforms: Array.isArray(t.platforms) ? t.platforms.map(String) : [],
      parent: t.parent ? String(t.parent) : null,
      url: String(t.url ?? ''),
      sourceCount: Number(t.sourceCount ?? 0)
    })),
    entities: Array.isArray(data.entities) ? data.entities : []
  };
}

export function loadStoredMitreMatrix(): MitreGraphData | null {
  try {
    const raw = globalThis.localStorage?.getItem(STORAGE_KEY);
    if (!raw) return null;
    return validateMitreMatrixData(JSON.parse(raw));
  } catch {
    return null;
  }
}

export function saveStoredMitreMatrix(data: MitreGraphData): void {
  globalThis.localStorage?.setItem(STORAGE_KEY, JSON.stringify(data));
}

export function clearStoredMitreMatrix(): void {
  globalThis.localStorage?.removeItem(STORAGE_KEY);
}

const simplify = (value: string) => value
  .trim()
  .toLocaleLowerCase('en-US')
  .replace(/&/g, 'and')
  .replace(/[^a-z0-9]+/g, ' ')
  .trim();

const numberField = (record: FindingRecord, pattern: RegExp) => {
  const value = record.fields.find(field => pattern.test(field.label.trim()))?.value ?? '';
  const parsed = Number(value.replace(',', '.').match(/-?\d+(?:\.\d+)?/)?.[0] ?? 0);
  return Number.isFinite(parsed) ? parsed : 0;
};

const severityRisk: Record<Severity, number> = {
  critical: 500,
  high: 400,
  medium: 300,
  low: 200,
  info: 100,
  unknown: 0
};

export function extractFindingMitreFocus(records: FindingRecord[], tacticOrder: string[]): FindingMitreFocus {
  const techniqueIds = new Set<string>();
  const tactics = new Set<string>();
  const tacticLookup = new Map(tacticOrder.map(tactic => [simplify(tactic), tactic]));
  let findingCount = 0;

  for (const record of records) {
    const fields = record.fields.filter(field => mitreBadgeLabel(field.label) && field.value);
    if (!fields.length) continue;
    findingCount += 1;
    for (const field of fields) {
      for (const match of field.value.toUpperCase().matchAll(/\bT\d{4}(?:\.\d{3})?\b/g)) techniqueIds.add(match[0]);
      const fragments = field.value.split(/[,;|/]+/);
      for (const fragment of fragments) {
        const normalized = simplify(fragment.replace(/\bT\d{4}(?:\.\d{3})?\b/gi, ''));
        const tactic = tacticLookup.get(normalized);
        if (tactic) tactics.add(tactic);
      }
    }
  }

  return { techniqueIds, tactics, findingCount };
}

export function buildFindingPathItems(records: FindingRecord[], data: MitreGraphData, titleColumn = '', fallbackToolId?: string): FindingPathItem[] {
  const tacticLookup = new Map(data.tacticOrder.map(tactic => [simplify(tactic), tactic]));
  const techniqueLookup = new Map(data.techniques.map(technique => [technique.id, technique]));

  return records.map(record => {
    const mitreFields = record.fields.filter(field => mitreBadgeLabel(field.label) && field.value);
    const attackFields = mitreFields.filter(field => mitreBadgeLabel(field.label) === 'MITRE ATT&CK');
    const defenseFields = mitreFields.filter(field => mitreBadgeLabel(field.label) !== 'MITRE ATT&CK');
    const hasNativeMitreField = attackFields.length > 0;
    const techniqueIds = [...new Set(attackFields.flatMap(field => [...field.value.toUpperCase().matchAll(/\bT\d{4}(?:\.\d{3})?\b/g)].map(match => match[0])))];
    const tactics = new Set<string>();
    for (const id of techniqueIds) for (const tactic of techniqueLookup.get(id)?.tactics ?? []) tactics.add(tactic);
    for (const field of attackFields) {
      for (const fragment of field.value.split(/[,;|/]+/)) {
        const tactic = tacticLookup.get(simplify(fragment.replace(/\bT\d{4}(?:\.\d{3})?\b/gi, '')));
        if (tactic) tactics.add(tactic);
      }
    }
    const orderedTactics = data.tacticOrder.filter(tactic => tactics.has(tactic));
    const severityValue = record.fields.find(field => /^(severity|criticality|priority)$/i.test(field.label.trim()))?.value;
    const severity = normalizeSeverity(severityValue);
    const weight = numberField(record, /^weight$/i);
    const score = numberField(record, /^score$/i);
    const title = recordTitleFor(record, titleColumn, fallbackToolId);
    const categories = [...new Set([
      ...record.fields
      .filter(field => /^category$/i.test(field.label.trim()) && field.value)
      .flatMap(field => field.value.split(/[;|]+/))
      .map(value => value.trim())
      .filter(Boolean),
      ...cisBenchmarkCategories(record, fallbackToolId)
    ])];
    const category = categories[0] ?? '';
    const defenses = [...new Set(defenseFields
      .flatMap(field => field.value.split(/[;|]+/))
      .map(value => value.trim())
      .filter(Boolean))];
    const targets = [...new Set(record.fields
      .filter(field => /^(?:target|ziel|betroffenes objekt|asset|host|hostname|domain|resource|system)$/i.test(field.label.trim()) && field.value)
      .flatMap(field => field.value.split(/[;|]+/))
      .map(value => value.trim())
      .filter(Boolean))];

    const mappingOrigin: 'native' | 'derived' | 'unmapped' = orderedTactics.length === 0
      ? 'unmapped'
      : hasNativeMitreField
        ? 'native'
        : 'derived';

    return {
      id: record.id,
      record,
      title,
      category,
      categories,
      defenses,
      targets,
      tactics: orderedTactics,
      primaryTactic: orderedTactics[0] ?? 'Unmapped',
      techniqueIds,
      severity,
      weight,
      score,
      risk: severityRisk[severity] + weight * 10 + score / 10,
      mappingOrigin
    };
  });
}

export function sharedEntityIds(left: string, right: string, entities: MitreEntity[], entityType: string): string[] {
  return entities
    .filter(entity => (entityType === 'all' || entity.type === entityType) && entity.techniques.includes(left) && entity.techniques.includes(right))
    .map(entity => entity.id);
}