SyncVulnPolicyBundleActivity.java

/*
 * 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.policy.vulnerability;

import com.fasterxml.jackson.core.JsonParser;
import com.fasterxml.jackson.databind.DeserializationFeature;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.dataformat.yaml.YAMLFactory;
import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule;
import com.networknt.schema.JsonSchema;
import com.networknt.schema.JsonSchemaFactory;
import com.networknt.schema.SpecVersion;
import com.networknt.schema.ValidationMessage;
import com.networknt.schema.serialization.DefaultJsonNodeReader;
import org.apache.commons.codec.digest.DigestUtils;
import org.dependencytrack.dex.api.Activity;
import org.dependencytrack.dex.api.ActivityContext;
import org.dependencytrack.dex.api.ActivitySpec;
import org.dependencytrack.dex.api.failure.TerminalApplicationFailureException;
import org.dependencytrack.persistence.jdbi.VulnerabilityPolicyDao;
import org.dependencytrack.persistence.jdbi.VulnerabilityPolicyDao.VulnPolicyBundleRow;
import org.dependencytrack.persistence.jdbi.VulnerabilityPolicyDao.VulnPolicyIdentityRow;
import org.dependencytrack.policy.cel.CelPolicyScriptHost;
import org.dependencytrack.policy.cel.CelPolicyType;
import org.dependencytrack.proto.internal.workflow.v1.SyncVulnPolicyBundleArg;
import org.dependencytrack.util.HttpUtil;
import org.eclipse.microprofile.config.Config;
import org.jspecify.annotations.NullMarked;
import org.jspecify.annotations.Nullable;
import org.projectnessie.cel.tools.ScriptCreateException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.io.UncheckedIOException;
import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.nio.file.Files;
import java.nio.file.Path;
import java.security.DigestOutputStream;
import java.security.MessageDigest;
import java.time.Duration;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.HexFormat;
import java.util.List;
import java.util.Map;
import java.util.NoSuchElementException;
import java.util.Objects;
import java.util.Set;
import java.util.UUID;
import java.util.function.Function;
import java.util.regex.Pattern;
import java.util.stream.Collectors;
import java.util.zip.ZipEntry;
import java.util.zip.ZipException;
import java.util.zip.ZipFile;

import static org.dependencytrack.persistence.jdbi.JdbiFactory.inJdbiTransaction;
import static org.dependencytrack.persistence.jdbi.JdbiFactory.useJdbiTransaction;

/**
 * @since 5.7.0
 */
@NullMarked
@ActivitySpec(name = "sync-vuln-policy-bundle")
public final class SyncVulnPolicyBundleActivity implements Activity<SyncVulnPolicyBundleArg, Void> {

    private static final Logger LOGGER = LoggerFactory.getLogger(SyncVulnPolicyBundleActivity.class);
    private static final Pattern POLICY_FILE_NAME_PATTERN =
            Pattern.compile("^(?![._])[^/\\\\]+\\.(yaml|yml)$", Pattern.CASE_INSENSITIVE);
    private static final int MAX_BUNDLE_SIZE = 10 * 1024 * 1024;
    private static final int MAX_ZIP_ENTRIES = 1000;
    private static final int MAX_ZIP_ENTRY_SIZE = 1024 * 1024;

    private final Config config;
    private final HttpClient httpClient;
    private final ObjectMapper objectMapper;
    private final JsonSchema schema;

    public SyncVulnPolicyBundleActivity(Config config, HttpClient httpClient) {
        this.config = config;
        this.httpClient = httpClient;
        this.objectMapper = new ObjectMapper(new YAMLFactory())
                .registerModule(new JavaTimeModule())
                .disable(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES)
                .disable(JsonParser.Feature.AUTO_CLOSE_SOURCE);
        this.schema = JsonSchemaFactory
                .builder(JsonSchemaFactory.getInstance(SpecVersion.VersionFlag.V202012))
                .jsonNodeReader(DefaultJsonNodeReader.builder()
                        .yamlMapper(objectMapper)
                        .build())
                .build()
                .getSchema(getClass().getResourceAsStream("/schema/vulnerability-policy-v1.schema.json"));
    }

    @Override
    public @Nullable Void execute(ActivityContext ctx, @Nullable SyncVulnPolicyBundleArg arg) throws Exception {
        if (arg == null || arg.getBundleUuid().isBlank()) {
            throw new TerminalApplicationFailureException("No bundle UUID provided");
        }

        final var bundleUuid = UUID.fromString(arg.getBundleUuid());

        // TODO: Eventually each bundle record will have its own URL.
        //  For the moment we only support a single default bundle.
        final URI bundleUri;
        try {
            bundleUri = config.getValue("dt.vulnerability.policy.bundle.url", URI.class);
        } catch (NoSuchElementException e) {
            throw new TerminalApplicationFailureException("No bundle URL configured", e);
        }

        final Credentials credentials = getCredentials();

        final BundleFile bundleFile = downloadBundle(bundleUri, credentials);
        if (bundleFile == null) {
            throw new TerminalApplicationFailureException("Bundle file not found");
        }

        try {
            final VulnPolicyBundleRow bundle = getOrCreateBundle(bundleUuid, bundleUri);
            if (isBundleUnchanged(bundle, bundleFile.digest())) {
                LOGGER.info("Bundle {} is unchanged", bundleUuid);
                return null;
            }

            final List<VulnerabilityPolicy> policies = parseAndValidatePolicies(bundleFile);
            useJdbiTransaction(handle -> {
                final var dao = handle.attach(VulnerabilityPolicyDao.class);
                reconcilePolicies(dao, bundle.id(), policies);
                dao.updateBundleAfterSync(bundleUuid, bundleUri.toString(), bundleFile.digest());
            });
        } finally {
            Files.deleteIfExists(bundleFile.path());
        }

        return null;
    }

    private VulnPolicyBundleRow getOrCreateBundle(UUID bundleUuid, URI bundleUri) {
        return inJdbiTransaction(handle -> {
            final var dao = handle.attach(VulnerabilityPolicyDao.class);

            final VulnPolicyBundleRow bundle = dao.getBundleByUuid(bundleUuid);
            if (bundle != null) {
                return bundle;
            }

            return dao.createBundle(bundleUuid, bundleUri);
        });
    }

    private static boolean isBundleUnchanged(VulnPolicyBundleRow bundle, String bundleDigest) {
        return Objects.equals(bundle.hash(), bundleDigest);
    }

    private record BundleFile(URI uri, String digest, Path path) {
    }

    private @Nullable BundleFile downloadBundle(
            URI bundleUri,
            @Nullable Credentials credentials) throws InterruptedException {
        final var requestBuilder = HttpRequest.newBuilder()
                .uri(bundleUri)
                .timeout(Duration.ofSeconds(10))
                .GET();
        if (credentials != null) {
            requestBuilder.header("Authorization", credentials.authHeaderValue());
        }

        final HttpResponse<InputStream> response;
        try {
            response = httpClient.send(
                    requestBuilder.build(),
                    HttpResponse.BodyHandlers.ofInputStream());
        } catch (IOException e) {
            throw new UncheckedIOException(e);
        }

        try (final InputStream bodyStream = response.body()) {
            if (response.statusCode() == 200) {
                final Path filePath = Files.createTempFile(null, null);
                try {
                    final MessageDigest fileDigest = DigestUtils.getSha256Digest();
                    try (final var fos = Files.newOutputStream(filePath);
                         final var dos = new DigestOutputStream(fos, fileDigest)) {
                        long totalRead = 0;
                        final var buffer = new byte[8192];
                        int bytesRead;
                        while ((bytesRead = bodyStream.read(buffer)) != -1) {
                            totalRead += bytesRead;
                            if (totalRead > MAX_BUNDLE_SIZE) {
                                throw new TerminalApplicationFailureException(
                                        "Bundle exceeds maximum size of %d bytes".formatted(MAX_BUNDLE_SIZE));
                            }
                            dos.write(buffer, 0, bytesRead);
                        }
                    }

                    return new BundleFile(
                            bundleUri,
                            HexFormat.of().formatHex(fileDigest.digest()),
                            filePath);
                } catch (Exception e) {
                    try {
                        Files.deleteIfExists(filePath);
                    } catch (IOException ex) {
                        e.addSuppressed(ex);
                    }
                    throw e;
                }
            }

            bodyStream.transferTo(OutputStream.nullOutputStream());

            if (response.statusCode() == 404) {
                return null;
            }

            throw new IllegalStateException("Unexpected response code: " + response.statusCode());
        } catch (IOException e) {
            throw new UncheckedIOException(e);
        }
    }

    private List<VulnerabilityPolicy> parseAndValidatePolicies(BundleFile bundleFile) {
        final var parsedPolicies = new ArrayList<VulnerabilityPolicy>();
        int numTotalPolicies = 0;
        int numFailures = 0;

        try (final var zipFile = new ZipFile(bundleFile.path().toFile())) {
            if (zipFile.size() > MAX_ZIP_ENTRIES) {
                throw new TerminalApplicationFailureException(
                        "Bundle contains more than %d entries".formatted(MAX_ZIP_ENTRIES));
            }

            final var entries = zipFile.entries();
            while (entries.hasMoreElements()) {
                final ZipEntry zipEntry = entries.nextElement();
                if (zipEntry.isDirectory()) {
                    continue;
                }

                final String fileName = Path.of(zipEntry.getName()).getFileName().toString();
                if (!POLICY_FILE_NAME_PATTERN.matcher(fileName).matches()) {
                    LOGGER.debug("Skipping file '{}': Unexpected file name pattern", zipEntry.getName());
                    continue;
                }

                if (zipEntry.getSize() > MAX_ZIP_ENTRY_SIZE) {
                    throw new TerminalApplicationFailureException(
                            "Entry '%s' exceeds maximum size of %d bytes".formatted(
                                    zipEntry.getName(), MAX_ZIP_ENTRY_SIZE));
                }

                final byte[] entryBytes;
                try (final InputStream entryStream = zipFile.getInputStream(zipEntry)) {
                    entryBytes = entryStream.readNBytes(MAX_ZIP_ENTRY_SIZE + 1);
                    if (entryBytes.length > MAX_ZIP_ENTRY_SIZE) {
                        throw new TerminalApplicationFailureException(
                                "Entry '%s' exceeds maximum size of %d bytes".formatted(
                                        zipEntry.getName(), MAX_ZIP_ENTRY_SIZE));
                    }
                }

                final JsonNode policyNode = objectMapper.readTree(entryBytes);
                if (!"Vulnerability Policy".equals(policyNode.path("type").asText(null))) {
                    LOGGER.debug("Skipping file '{}': Not a vulnerability policy", zipEntry.getName());
                    continue;
                }

                numTotalPolicies++;

                final Set<ValidationMessage> validationMessages = schema.validate(policyNode);
                if (!validationMessages.isEmpty()) {
                    LOGGER.warn(
                            "Schema validation of '{}' failed: {}",
                            zipEntry.getName(),
                            validationMessages.stream()
                                    .map(ValidationMessage::getMessage)
                                    .collect(Collectors.joining(", ")));
                    numFailures++;
                    continue;
                }

                final var vulnPolicy = objectMapper.treeToValue(policyNode, VulnerabilityPolicy.class);
                boolean conditionValid = true;
                try {
                    CelPolicyScriptHost
                            .getInstance(CelPolicyType.VULNERABILITY)
                            .compile(vulnPolicy.getCondition(), CelPolicyScriptHost.CacheMode.NO_CACHE);
                } catch (ScriptCreateException e) {
                    LOGGER.warn("Failed to compile condition of policy '{}'", vulnPolicy.getName(), e);
                    conditionValid = false;
                }

                if (conditionValid) {
                    parsedPolicies.add(vulnPolicy);
                } else {
                    numFailures++;
                }
            }
        } catch (ZipException e) {
            throw new TerminalApplicationFailureException("Bundle is not a valid ZIP archive", e);
        } catch (IOException e) {
            throw new UncheckedIOException(e);
        }

        if (numFailures > 0) {
            throw new TerminalApplicationFailureException("""
                    Will not update policies in datastore because %d/%d policy definitions failed to validate. \
                    Policy bundles can only be applied when all policy definitions within them are valid.\
                    """.formatted(numFailures, numTotalPolicies));
        }

        return parsedPolicies;
    }

    private void reconcilePolicies(VulnerabilityPolicyDao dao, long bundleId, List<VulnerabilityPolicy> policies) {
        final Map<String, VulnPolicyIdentityRow> existingPolicyByName =
                dao.getAllByBundleId(bundleId).stream()
                        .collect(Collectors.toMap(VulnPolicyIdentityRow::name, Function.identity()));
        final var policiesToCreate = new ArrayList<VulnerabilityPolicy>();
        final var policiesToUpdate = new ArrayList<VulnerabilityPolicy>();
        final var policyNamesSeen = new HashSet<String>();
        final var duplicateNames = new ArrayList<String>();

        for (final VulnerabilityPolicy policy : policies) {
            if (!policyNamesSeen.add(policy.getName())) {
                duplicateNames.add(policy.getName());
                continue;
            }

            if (existingPolicyByName.containsKey(policy.getName())) {
                policiesToUpdate.add(policy);
            } else {
                policiesToCreate.add(policy);
            }
        }

        if (!duplicateNames.isEmpty()) {
            throw new TerminalApplicationFailureException(
                    "Bundle contains duplicate policy names: " + duplicateNames);
        }

        // Detect name conflicts between bundle policies and user-managed policies.
        // Names are globally unique, so if a user-managed policy has the same name as
        // a bundle policy, the sync cannot proceed.
        if (!policiesToCreate.isEmpty()) {
            final Set<String> newNames = policiesToCreate.stream()
                    .map(VulnerabilityPolicy::getName)
                    .collect(Collectors.toSet());
            final List<String> conflicting = dao.getUserManagedPolicyNames().stream()
                    .filter(newNames::contains)
                    .toList();
            if (!conflicting.isEmpty()) {
                throw new TerminalApplicationFailureException("""
                        Bundle sync aborted: the following policy names conflict with \
                        existing user-managed policies and must be renamed or deleted \
                        before the bundle can be applied:\s""" + conflicting);
            }
        }

        final var policyNamesToDelete = new HashSet<>(existingPolicyByName.keySet());
        policyNamesToDelete.removeAll(policyNamesSeen);

        for (final String policyName : policyNamesToDelete) {
            LOGGER.info("Deleting policy '{}'", policyName);
            dao.unassignAndDeleteByNameAndBundleId(policyName, bundleId);
        }

        if (!policiesToCreate.isEmpty()) {
            LOGGER.info("Creating {} policies", policiesToCreate.size());
            dao.createAll(bundleId, policiesToCreate);
        }

        if (!policiesToUpdate.isEmpty()) {
            LOGGER.info("Updating {} policies", policiesToUpdate.size());
            dao.updateAll(policiesToUpdate);
        }
    }

    private sealed interface Credentials {

        String authHeaderValue();

        record BasicAuth(String username, String password) implements Credentials {

            @Override
            public String authHeaderValue() {
                return HttpUtil.basicAuthHeaderValue(username, password);
            }

        }

        record BearerToken(String token) implements Credentials {

            @Override
            public String authHeaderValue() {
                return "Bearer " + token;
            }

        }

    }

    private @Nullable Credentials getCredentials() {
        final var username = config
                .getOptionalValue("dt.vulnerability.policy.bundle.auth.username", String.class)
                .orElse(null);
        final var password = config
                .getOptionalValue("dt.vulnerability.policy.bundle.auth.password", String.class)
                .orElse(null);
        if (username != null && password != null) {
            return new Credentials.BasicAuth(username, password);
        }

        return config
                .getOptionalValue("dt.vulnerability.policy.bundle.auth.bearer.token", String.class)
                .map(Credentials.BearerToken::new)
                .orElse(null);
    }

}