hiveops-aria/src/main/java/com/hiveops/aria/service/ThreatEventClassifier.java

package com.hiveops.aria.service;

import com.hiveops.aria.dto.ThreatEventItem;
import com.hiveops.aria.entity.ThreatEvent;
import com.hiveops.aria.entity.ThreatEvent.AttackPhase;
import com.hiveops.aria.entity.ThreatEvent.ThreatSeverity;
import com.hiveops.aria.entity.ThreatIoc;
import com.hiveops.aria.repository.ThreatIocRepository;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import java.util.Optional;

@Service
@RequiredArgsConstructor
@Slf4j
public class ThreatEventClassifier {

    private final ThreatIocRepository iocRepository;

    public ClassificationResult classify(ThreatEventItem item) {
        String signal = item.getSignalKey();
        ThreatSeverity severity;
        AttackPhase phase = null;
        ThreatIoc matchedIoc = null;

        switch (signal) {
            // Windows Security Event IDs — FBI FLASH-20260219-001
            case "EVENT_1102" -> {
                severity = ThreatSeverity.CRITICAL;
                phase = AttackPhase.CLEANUP;
            }
            case "EVENT_4719" -> {
                severity = ThreatSeverity.HIGH;
                phase = AttackPhase.CLEANUP;
            }
            case "EVENT_6416", "EVENT_2003" -> {
                severity = ThreatSeverity.HIGH;
                phase = AttackPhase.PHYSICAL_ACCESS;
                matchedIoc = matchUsb(item);
            }
            case "EVENT_4663" -> {
                severity = ThreatSeverity.HIGH;
                phase = AttackPhase.MALWARE_STAGING;
                matchedIoc = matchFilePath(item);
            }
            case "EVENT_4688" -> {
                matchedIoc = matchProcessName(item);
                if (matchedIoc != null) {
                    severity = ThreatSeverity.HIGH;
                    phase = AttackPhase.MALWARE_EXECUTION;
                } else {
                    severity = ThreatSeverity.INFO;
                }
            }
            case "EVENT_4697" -> {
                matchedIoc = matchServiceName(item);
                if (matchedIoc != null) {
                    severity = ThreatSeverity.HIGH;
                    phase = AttackPhase.PERSISTENCE;
                } else {
                    severity = ThreatSeverity.MEDIUM;
                    phase = AttackPhase.PERSISTENCE;
                }
            }
            // Physical sensor signals
            case "DOOR_OPEN" -> {
                severity = ThreatSeverity.MEDIUM;
                phase = AttackPhase.PHYSICAL_ACCESS;
            }
            case "CARD_READER_TAMPER" -> {
                severity = ThreatSeverity.HIGH;
                phase = AttackPhase.PHYSICAL_ACCESS;
            }
            case "USB_INSERTION" -> {
                severity = ThreatSeverity.HIGH;
                phase = AttackPhase.PHYSICAL_ACCESS;
            }
            case "POWER_CYCLE_UNEXPECTED", "OUT_OF_SERVICE" -> {
                severity = ThreatSeverity.MEDIUM;
            }
            case "VIBRATION_ANOMALY" -> {
                severity = ThreatSeverity.MEDIUM;
            }
            case "CASH_LOW_UNEXPECTED" -> {
                severity = ThreatSeverity.HIGH;
            }
            // File system signals
            case "HASH_DRIFT" -> {
                severity = ThreatSeverity.HIGH;
                phase = AttackPhase.MALWARE_STAGING;
            }
            case "IOC_PROCESS_MATCH", "IOC_FILENAME_MATCH" -> {
                matchedIoc = matchProcessName(item);
                severity = ThreatSeverity.HIGH;
                phase = AttackPhase.MALWARE_EXECUTION;
            }
            case "REGISTRY_PERSISTENCE_DETECTED" -> {
                matchedIoc = matchRegistryKey(item);
                severity = ThreatSeverity.HIGH;
                phase = AttackPhase.PERSISTENCE;
            }
            case "SUSPICIOUS_DIRECTORY" -> {
                matchedIoc = matchDirectory(item);
                severity = ThreatSeverity.HIGH;
                phase = AttackPhase.MALWARE_STAGING;
            }
            case "REMOTE_TOOL_DETECTED" -> {
                matchedIoc = matchProcessName(item);
                severity = ThreatSeverity.MEDIUM;
            }
            default -> {
                severity = ThreatSeverity.INFO;
                log.debug("Unclassified signal key: {}", signal);
            }
        }

        return new ClassificationResult(severity, phase, matchedIoc);
    }

    private ThreatIoc matchProcessName(ThreatEventItem item) {
        String name = extractString(item, "processName", "imageName", "fileName");
        if (name == null) return null;
        return iocRepository.findByActiveTrueAndIocTypeAndValueIgnoreCase(ThreatIoc.IocType.FILENAME, name).orElse(null);
    }

    private ThreatIoc matchServiceName(ThreatEventItem item) {
        String name = extractString(item, "serviceName", "serviceDisplayName");
        if (name == null) return null;
        return iocRepository.findByActiveTrueAndIocTypeAndValueIgnoreCase(ThreatIoc.IocType.SERVICE_NAME, name).orElse(null);
    }

    private ThreatIoc matchFilePath(ThreatEventItem item) {
        String path = extractString(item, "objectName", "filePath");
        if (path == null) return null;
        Optional<ThreatIoc> dir = iocRepository.findByActiveTrueAndIocTypeAndValueIgnoreCase(ThreatIoc.IocType.DIRECTORY, path);
        return dir.orElseGet(() ->
                iocRepository.findByActiveTrueAndIocTypeAndValueIgnoreCase(ThreatIoc.IocType.FILE_PATH, path).orElse(null));
    }

    private ThreatIoc matchRegistryKey(ThreatEventItem item) {
        String key = extractString(item, "registryKey", "keyPath");
        if (key == null) return null;
        return iocRepository.findByActiveTrueAndIocTypeAndValueIgnoreCase(ThreatIoc.IocType.REGISTRY_KEY, key).orElse(null);
    }

    private ThreatIoc matchDirectory(ThreatEventItem item) {
        String path = extractString(item, "directory", "path");
        if (path == null) return null;
        return iocRepository.findByActiveTrueAndIocTypeAndValueIgnoreCase(ThreatIoc.IocType.DIRECTORY, path).orElse(null);
    }

    private ThreatIoc matchUsb(ThreatEventItem item) {
        String deviceId = extractString(item, "deviceId", "deviceDescription");
        if (deviceId == null) return null;
        return iocRepository.findByActiveTrueAndIocTypeAndValueIgnoreCase(ThreatIoc.IocType.IP, deviceId).orElse(null);
    }

    private String extractString(ThreatEventItem item, String... keys) {
        if (item.getRawPayload() == null) return null;
        for (String key : keys) {
            Object val = item.getRawPayload().get(key);
            if (val instanceof String s && !s.isBlank()) return s;
        }
        return null;
    }

    public record ClassificationResult(ThreatSeverity severity, AttackPhase attackPhase, ThreatIoc matchedIoc) {}
}