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

package com.hiveops.aria.service;

import com.hiveops.aria.entity.IocFeedRun;
import com.hiveops.aria.entity.ThreatIoc;
import com.hiveops.aria.feed.MalwareBazaarFeedClient;
import com.hiveops.aria.feed.OtxFeedClient;
import com.hiveops.aria.feed.ThreatFoxFeedClient;
import com.hiveops.aria.repository.IocFeedRunRepository;
import com.hiveops.aria.repository.ThreatIocRepository;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;

import java.time.Instant;
import java.util.List;
import java.util.Map;
import java.util.Optional;

@Service
@RequiredArgsConstructor
@Slf4j
public class IocFeedService {

    private final OtxFeedClient           otxClient;
    private final MalwareBazaarFeedClient  malwareBazaarClient;
    private final ThreatFoxFeedClient      threatFoxClient;
    private final ThreatIocRepository      iocRepository;
    private final IocFeedRunRepository     feedRunRepository;

    private static final Map<String, ThreatIoc.IocSource> FEED_SOURCES = Map.of(
        "OTX",            ThreatIoc.IocSource.OTX,
        "MALWARE_BAZAAR", ThreatIoc.IocSource.MALWARE_BAZAAR,
        "THREAT_FOX",     ThreatIoc.IocSource.THREAT_FOX
    );

    @Scheduled(cron = "${aria.feed.refresh.cron:0 0 3 * * *}")
    public void scheduledRefresh() {
        log.info("IOC feed scheduled refresh starting");
        refreshAll();
    }

    public void refreshAll() {
        for (String feedName : FEED_SOURCES.keySet()) {
            try {
                refreshFeed(feedName);
            } catch (Exception e) {
                log.error("IOC feed {} failed: {}", feedName, e.getMessage(), e);
            }
        }
    }

    @Transactional
    public IocFeedRun refreshFeed(String feedName) {
        if (feedRunRepository.countRunning(feedName) > 0) {
            log.info("IOC feed {} already running, skipping", feedName);
            return feedRunRepository.findTopByFeedNameOrderByStartedAtDesc(feedName).orElseThrow();
        }

        IocFeedRun run = feedRunRepository.save(IocFeedRun.builder()
                .feedName(feedName)
                .startedAt(Instant.now())
                .build());

        try {
            List<ThreatIoc> fetched = fetchFromFeed(feedName);
            int[] counts = upsertIocs(fetched);

            run.setStatus("SUCCESS");
            run.setIocsAdded(counts[0]);
            run.setIocsUpdated(counts[1]);
            run.setCompletedAt(Instant.now());
            log.info("IOC feed {} complete: {} added, {} updated", feedName, counts[0], counts[1]);
        } catch (Exception e) {
            run.setStatus("ERROR");
            run.setErrorMessage(e.getMessage());
            run.setCompletedAt(Instant.now());
            log.error("IOC feed {} error: {}", feedName, e.getMessage(), e);
        }

        return feedRunRepository.save(run);
    }

    private List<ThreatIoc> fetchFromFeed(String feedName) {
        return switch (feedName) {
            case "OTX"            -> otxClient.fetch();
            case "MALWARE_BAZAAR" -> malwareBazaarClient.fetch();
            case "THREAT_FOX"     -> threatFoxClient.fetch();
            default               -> throw new IllegalArgumentException("Unknown feed: " + feedName);
        };
    }

    private int[] upsertIocs(List<ThreatIoc> iocs) {
        int added = 0, updated = 0;
        for (ThreatIoc ioc : iocs) {
            if (ioc.getValue() == null || ioc.getValue().isBlank()) continue;
            Optional<ThreatIoc> existing = iocRepository
                    .findByActiveTrueAndIocTypeAndValueIgnoreCase(ioc.getIocType(), ioc.getValue());
            if (existing.isPresent()) {
                ThreatIoc e = existing.get();
                e.setDescription(ioc.getDescription());
                e.setSourceRef(ioc.getSourceRef());
                e.setConfidence(ioc.getConfidence());
                e.setReportedAt(ioc.getReportedAt());
                e.setExpiresAt(ioc.getExpiresAt());
                iocRepository.save(e);
                updated++;
            } else {
                iocRepository.save(ioc);
                added++;
            }
        }
        return new int[]{ added, updated };
    }

    public List<String> feedNames() {
        return List.of("OTX", "MALWARE_BAZAAR", "THREAT_FOX");
    }

    public boolean isConfigured(String feedName) {
        return switch (feedName) {
            case "OTX"            -> otxClient.isConfigured();
            case "MALWARE_BAZAAR" -> malwareBazaarClient.isConfigured();
            case "THREAT_FOX"     -> threatFoxClient.isConfigured();
            default               -> false;
        };
    }
}