/*
 * This file is part of Dependency-Track.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 * SPDX-License-Identifier: Apache-2.0
 * Copyright (c) OWASP Foundation. All Rights Reserved.
 */
package org.dependencytrack.tasks;

import alpine.common.logging.Logger;
import alpine.event.framework.Event;
import alpine.event.framework.Subscriber;
import alpine.persistence.ScopedCustomization;
import org.dependencytrack.event.ComponentVulnerabilityAnalysisEvent;
import org.dependencytrack.event.InternalAnalysisEvent;
import org.dependencytrack.event.OssIndexAnalysisEvent;
import org.dependencytrack.event.PortfolioVulnerabilityAnalysisEvent;
import org.dependencytrack.event.ProjectMetricsUpdateEvent;
import org.dependencytrack.event.ProjectVulnerabilityAnalysisEvent;
import org.dependencytrack.event.SnykAnalysisEvent;
import org.dependencytrack.event.TrivyAnalysisEvent;
import org.dependencytrack.event.VulnDbAnalysisEvent;
import org.dependencytrack.model.Component;
import org.dependencytrack.model.ComponentAnalysisCache;
import org.dependencytrack.model.FindingAttribution;
import org.dependencytrack.model.Project;
import org.dependencytrack.model.Vulnerability;
import org.dependencytrack.model.VulnerabilityAnalysisLevel;
import org.dependencytrack.persistence.QueryManager;
import org.dependencytrack.policy.PolicyEngine;
import org.dependencytrack.tasks.scanners.AnalyzerIdentity;
import org.dependencytrack.tasks.scanners.CacheableScanTask;
import org.dependencytrack.tasks.scanners.InternalAnalysisTask;
import org.dependencytrack.tasks.scanners.OssIndexAnalysisTask;
import org.dependencytrack.tasks.scanners.ScanTask;
import org.dependencytrack.tasks.scanners.SnykAnalysisTask;
import org.dependencytrack.tasks.scanners.TrivyAnalysisTask;
import org.dependencytrack.tasks.scanners.VulnDbAnalysisTask;
import org.slf4j.MDC;

import javax.jdo.Query;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Date;
import java.util.HashMap;
import java.util.List;
import java.util.UUID;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.locks.ReentrantLock;

import static org.dependencytrack.common.MdcKeys.MDC_EVENT_TOKEN;
import static org.dependencytrack.common.MdcKeys.MDC_PROJECT_NAME;
import static org.dependencytrack.common.MdcKeys.MDC_PROJECT_UUID;
import static org.dependencytrack.common.MdcKeys.MDC_PROJECT_VERSION;
import static org.dependencytrack.common.MdcKeys.MDC_VULN_ANALYSIS_LEVEL;
import static org.dependencytrack.util.LockUtil.getLockForProjectAndNamespace;

public class VulnerabilityAnalysisTask implements Subscriber {

    private static final Logger LOGGER = Logger.getLogger(VulnerabilityAnalysisTask.class);

    /**
     * {@inheritDoc}
     */
    @Override
    public void inform(final Event e) {
        switch (e) {
            case ComponentVulnerabilityAnalysisEvent event -> {
                try (var ignoredMdcEventToken = MDC.putCloseable(MDC_EVENT_TOKEN, event.getChainIdentifier().toString())) {
                    analyzeComponent(event.componentUuid());
                }
            }
            case ProjectVulnerabilityAnalysisEvent event -> {
                try (var ignoredMdcEventToken = MDC.putCloseable(MDC_EVENT_TOKEN, event.getChainIdentifier().toString())) {
                    analyzeProject(event.projectUuid(), event.analysisLevel());
                }
            }
            case PortfolioVulnerabilityAnalysisEvent ignored -> analyzePortfolio();
            default -> throw new IllegalArgumentException("Unexpected event: " + e);
        }
    }

    private void analyzePortfolio() {
        try (final var qm = new QueryManager()) {
            List<Project> projects = fetchNextProjectBatch(qm, null);
            if (projects.isEmpty()) {
                LOGGER.info("Portfolio does not have any active projects; Nothing to analyze");
                return;
            }

            while (!projects.isEmpty()) {
                if (Thread.currentThread().isInterrupted()) {
                    LOGGER.warn("Interrupted before all projects could be analyzed");
                    break;
                }

                LOGGER.info("Analyzing batch of %d projects".formatted(projects.size()));

                for (final Project project : projects) {
                    try (var ignoredMdcProjectUuid = MDC.putCloseable(MDC_PROJECT_UUID, project.getUuid().toString());
                         var ignoredMdcProjectName = MDC.putCloseable(MDC_PROJECT_NAME, project.getName());
                         var ignoredMdcProjectVersion = MDC.putCloseable(MDC_PROJECT_VERSION, project.getVersion());
                         var ignoredMdcAnalysisLevel = MDC.putCloseable(MDC_VULN_ANALYSIS_LEVEL, VulnerabilityAnalysisLevel.PERIODIC_ANALYSIS.name())) {
                        if (Thread.currentThread().isInterrupted()) {
                            LOGGER.warn("Interrupted before project could be analyzed");
                            break;
                        }

                        try {
                            analyzeProject(qm, project, VulnerabilityAnalysisLevel.PERIODIC_ANALYSIS);
                        } catch (RuntimeException e) {
                            LOGGER.error("Failed to analyze project", e);
                        }
                    } finally {
                        qm.getPersistenceManager().evictAll(false, Component.class);
                        qm.getPersistenceManager().evictAll(false, ComponentAnalysisCache.class);
                        qm.getPersistenceManager().evictAll(false, FindingAttribution.class);
                        qm.getPersistenceManager().evictAll(false, Vulnerability.class);
                    }
                }

                qm.getPersistenceManager().evictAll(false, Project.class);
                projects = fetchNextProjectBatch(qm, projects.getLast().getId());
            }
        }
    }

    private void analyzeProject(
            final QueryManager qm,
            final Project project,
            final VulnerabilityAnalysisLevel analysisLevel) {
        final ReentrantLock projectLock = getLockForProjectAndNamespace(project, getClass().getSimpleName());

        try {
            try {
                final boolean lockAcquired = projectLock.tryLock(5, TimeUnit.MINUTES);
                if (!lockAcquired) {
                    LOGGER.warn("Failed to acquire lock after 5min; Skipping analysis");
                    return;
                }
            } catch (InterruptedException e) {
                LOGGER.warn("Interrupted while waiting for lock; Not performing analysis", e);
                Thread.currentThread().interrupt();
                return;
            }

            // NB: Some analyzers require all components of a project to be analyzed in one go.
            // Trivy for example checks for the existence of OPERATING_SYSTEM components.
            // If we were to analyze components in batches, this logic might not work for large projects.
            final List<Component> components = fetchComponents(qm, project);
            if (components.isEmpty()) {
                LOGGER.info("Project does not have any components; Nothing to analyze");
                return;
            }

            analyzeComponents(qm, components, analysisLevel);

            if (analysisLevel == VulnerabilityAnalysisLevel.PERIODIC_ANALYSIS) {
                try {
                    performPolicyEvaluation(project, components);
                } catch (RuntimeException e) {
                    LOGGER.warn("Policy evaluation against %d components failed".formatted(
                            components.size()), e);
                }
            }

            project.setLastVulnerabilityAnalysis(new Date());
        } finally {
            projectLock.unlock();
        }
    }

    private void analyzeProject(final UUID projectUuid, final VulnerabilityAnalysisLevel analysisLevel) {
        try (final var qm = new QueryManager()) {
            final Project project;
            try (var ignoredPersistenceCustomization = new ScopedCustomization(qm.getPersistenceManager())
                    .withFetchGroup(Project.FetchGroup.PROJECT_VULN_ANALYSIS.name())) {
                project = qm.getObjectByUuid(Project.class, projectUuid);
            }
            if (project == null) {
                LOGGER.warn("Project with UUID %s does not exist; Skipping analysis".formatted(projectUuid));
                return;
            }

            try (var ignoredMdcProjectUuid = MDC.putCloseable(MDC_PROJECT_UUID, project.getUuid().toString());
                 var ignoredMdcProjectName = MDC.putCloseable(MDC_PROJECT_NAME, project.getName());
                 var ignoredMdcProjectVersion = MDC.putCloseable(MDC_PROJECT_VERSION, project.getVersion());
                 var ignoredMdcAnalysisLevel = MDC.putCloseable(MDC_VULN_ANALYSIS_LEVEL, analysisLevel.name())) {
                try {
                    analyzeProject(qm, project, analysisLevel);
                } catch (RuntimeException e) {
                    LOGGER.error("Failed to analyze project", e);
                }
            }
        }
    }

    private void analyzeComponent(final UUID componentUuid) {
        try (final var qm = new QueryManager()) {
            final Component component;
            try (var ignoredPersistenceCustomization = new ScopedCustomization(qm.getPersistenceManager())
                    .withFetchGroup(Component.FetchGroup.COMPONENT_VULN_ANALYSIS.name())) {
                component = qm.getObjectByUuid(Component.class, componentUuid);
            }
            if (component == null) {
                LOGGER.warn("Component with UUID %s does not exist; Skipping analysis".formatted(componentUuid));
                return;
            }

            analyzeComponents(qm, List.of(component), VulnerabilityAnalysisLevel.ON_DEMAND);
        }
    }

    private void analyzeComponents(
            final QueryManager qm,
            final Collection<Component> components,
            final VulnerabilityAnalysisLevel analysisLevel) {
        final var analyzers = List.of(
                new InternalAnalysisTask(),
                new OssIndexAnalysisTask(),
                new SnykAnalysisTask(),
                new TrivyAnalysisTask(),
                new VulnDbAnalysisTask());

        final var candidateComponentsByAnalyzerIdentity = new HashMap<AnalyzerIdentity, List<Component>>();
        for (final Component component : components) {
            for (final ScanTask analyzer : analyzers) {
                if (analyzer.isCapable(component)) {
                    if (analyzer instanceof final CacheableScanTask cacheableScanTask
                        && !cacheableScanTask.shouldAnalyze(component.getPurl())) {
                        cacheableScanTask.applyAnalysisFromCache(component);
                        continue;
                    }

                    candidateComponentsByAnalyzerIdentity
                            .computeIfAbsent(analyzer.getAnalyzerIdentity(), ignored -> new ArrayList<>())
                            .add(component);
                }
            }
        }

        // NB: Detaching components from this persistence manager so they can be used
        // by PMs of analyzers. Components MUST NOT be made transient (QueryManager#makeTransientAll) here,
        // as that would also dis-associate them from DataNucleus' state manager, preventing lazy-loading
        // of other fields, should analyzers need them (e.g. component properties).
        for (final Component component : components) {
            qm.getPersistenceManager().detachCopy(component);
        }

        for (final ScanTask analyzer : analyzers) {
            final List<Component> candidates = candidateComponentsByAnalyzerIdentity.get(analyzer.getAnalyzerIdentity());
            if (candidates == null || candidates.isEmpty()) {
                LOGGER.debug("No analysis candidates for %s; Not invoking analyzer".formatted(
                        analyzer.getAnalyzerIdentity()));
                continue;
            }

            // TODO: It would be better to invoke ScanTask#analyze directly rather than
            //  going through the indirection of Subscriber#inform. Then, we could have analyzers
            //  return the vulnerabilities they identified, which we need to determine if we
            //  can auto-suppress vulns that are no longer reported by ANY analyzer.
            final Event event = switch (analyzer.getAnalyzerIdentity()) {
                case INTERNAL_ANALYZER -> new InternalAnalysisEvent(candidates, analysisLevel);
                case OSSINDEX_ANALYZER -> new OssIndexAnalysisEvent(candidates, analysisLevel);
                case VULNDB_ANALYZER -> new VulnDbAnalysisEvent(candidates, analysisLevel);
                case SNYK_ANALYZER -> new SnykAnalysisEvent(candidates, analysisLevel);
                case TRIVY_ANALYZER -> new TrivyAnalysisEvent(candidates, analysisLevel);
                case NPM_AUDIT_ANALYZER, NONE -> throw new IllegalStateException(
                        "Unsupported analyzer: " + analyzer.getAnalyzerIdentity());
            };

            try {
                LOGGER.debug("Invoking %s with %d components".formatted(
                        analyzer.getAnalyzerIdentity(), candidates.size()));

                final var task = (Subscriber) analyzer;
                task.inform(event);
            } finally {
                // Clear the transient cache result for each component.
                // Each analyzer will have its own result. Therefore, we do not want to mix them.
                for (final Component component : candidates) {
                    component.setCacheResult(null);
                }
            }
        }
    }

    private void performPolicyEvaluation(Project project, List<Component> components) {
        // Evaluate the components against applicable policies via the PolicyEngine.
        final PolicyEngine pe = new PolicyEngine();
        pe.evaluate(components);
        if (project != null) {
            Event.dispatch(new ProjectMetricsUpdateEvent(project.getUuid()));
        }
    }

    private List<Project> fetchNextProjectBatch(final QueryManager qm, final Long lastId) {
        final var filterParts = new ArrayList<String>(2);
        filterParts.add("active");

        final var filterParams = new HashMap<String, Object>();

        if (lastId != null) {
            filterParts.add("id > :lastId");
            filterParams.put("lastId", lastId);
        }

        final Query<Project> query = qm.getPersistenceManager().newQuery(Project.class);
        query.setFilter(String.join(" && ", filterParts));
        query.setNamedParameters(filterParams);
        query.setOrdering("id asc");
        query.setRange(0, 100);

        try (var ignoredPersistenceCustomization = new ScopedCustomization(qm.getPersistenceManager())
                .withFetchGroup(Project.FetchGroup.PROJECT_VULN_ANALYSIS.name())) {
            return List.copyOf(query.executeList());
        } finally {
            query.closeAll();
        }
    }

    private List<Component> fetchComponents(final QueryManager qm, final Project project) {
        final Query<Component> query = qm.getPersistenceManager().newQuery(Component.class);
        query.setFilter("project.id == :projectId");
        query.setParameters(project.getId());

        try (var ignoredPersistenceCustomization = new ScopedCustomization(qm.getPersistenceManager())
                .withFetchGroup(Component.FetchGroup.COMPONENT_VULN_ANALYSIS.name())) {
            return List.copyOf(query.executeList());
        } finally {
            query.closeAll();
        }
    }

}
