/*
 * 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) Steve Springett. All Rights Reserved.
 */
package org.dependencytrack.resources.v1;

import alpine.auth.PermissionRequired;
import alpine.persistence.PaginatedResult;
import alpine.resources.AlpineResource;
import alpine.util.UuidUtil;
import io.swagger.annotations.Api;
import io.swagger.annotations.ApiOperation;
import io.swagger.annotations.ApiParam;
import io.swagger.annotations.ApiResponse;
import io.swagger.annotations.ApiResponses;
import io.swagger.annotations.Authorization;
import io.swagger.annotations.ResponseHeader;
import org.dependencytrack.auth.Permissions;
import org.dependencytrack.model.Component;
import org.dependencytrack.model.Cwe;
import org.dependencytrack.model.Project;
import org.dependencytrack.model.Vulnerability;
import org.dependencytrack.persistence.QueryManager;
import us.springett.cvss.Cvss;
import us.springett.cvss.Score;
import javax.validation.Validator;
import javax.ws.rs.Consumes;
import javax.ws.rs.DELETE;
import javax.ws.rs.GET;
import javax.ws.rs.POST;
import javax.ws.rs.PUT;
import javax.ws.rs.Path;
import javax.ws.rs.PathParam;
import javax.ws.rs.Produces;
import javax.ws.rs.QueryParam;
import javax.ws.rs.core.MediaType;
import javax.ws.rs.core.Response;
import java.math.BigDecimal;
import java.util.List;

/**
 * JAX-RS resources for processing vulnerabilities.
 *
 * @author Steve Springett
 * @since 3.0.0
 */
@Path("/v1/vulnerability")
@Api(value = "vulnerability", authorizations = @Authorization(value = "X-Api-Key"))
public class VulnerabilityResource extends AlpineResource {

    @GET
    @Path("/component/{ident}")
    @Produces(MediaType.APPLICATION_JSON)
    @ApiOperation(
            value = "Returns a list of all vulnerabilities for a specific component",
            notes = "A valid UUID of the component may be specified, or the MD5 or SHA1 hash of the component",
            response = Vulnerability.class,
            responseContainer = "List",
            responseHeaders = @ResponseHeader(name = TOTAL_COUNT_HEADER, response = Long.class, description = "The total number of vulnerabilities")
    )
    @ApiResponses(value = {
            @ApiResponse(code = 401, message = "Unauthorized"),
            @ApiResponse(code = 404, message = "The component could not be found")
    })
    @PermissionRequired(Permissions.Constants.VIEW_PORTFOLIO)
    public Response getVulnerabilitiesByComponent(@PathParam("ident") String ident,
                                                  @ApiParam(value = "Optionally includes suppressed vulnerabilities")
                                                  @QueryParam("suppressed") boolean suppressed) {
        try (QueryManager qm = new QueryManager(getAlpineRequest())) {
            final Component component;
            if (UuidUtil.isValidUUID(ident)) {
                component = qm.getObjectByUuid(Component.class, ident);
            } else {
                component = qm.getComponentByHash(ident);
            }
            if (component != null) {
                final PaginatedResult result = qm.getVulnerabilities(component, suppressed);
                return Response.ok(result.getObjects()).header(TOTAL_COUNT_HEADER, result.getTotal()).build();
            } else {
                return Response.status(Response.Status.NOT_FOUND).entity("The component could not be found.").build();
            }
        }
    }

    @GET
    @Path("/project/{uuid}")
    @Produces(MediaType.APPLICATION_JSON)
    @ApiOperation(
            value = "Returns a list of all vulnerabilities for a specific project",
            response = Vulnerability.class,
            responseContainer = "List",
            responseHeaders = @ResponseHeader(name = TOTAL_COUNT_HEADER, response = Long.class, description = "The total number of vulnerabilities")
    )
    @ApiResponses(value = {
            @ApiResponse(code = 401, message = "Unauthorized"),
            @ApiResponse(code = 404, message = "The project could not be found")
    })
    @PermissionRequired(Permissions.Constants.VIEW_PORTFOLIO)
    public Response getVulnerabilitiesByProject(@PathParam("uuid") String uuid,
                                                @ApiParam(value = "Optionally includes suppressed vulnerabilities")
                                                @QueryParam("suppressed") boolean suppressed) {
        try (QueryManager qm = new QueryManager(getAlpineRequest())) {
            final Project project = qm.getObjectByUuid(Project.class, uuid);
            if (project != null) {
                final List<Vulnerability> vulnerabilities = qm.getVulnerabilities(project, suppressed);
                return Response.ok(vulnerabilities).header(TOTAL_COUNT_HEADER, vulnerabilities.size()).build();
            } else {
                return Response.status(Response.Status.NOT_FOUND).entity("The project could not be found.").build();
            }
        }
    }

    @GET
    @Path("/{uuid}")
    @Produces(MediaType.APPLICATION_JSON)
    @ApiOperation(
            value = "Returns a specific vulnerability",
            response = Vulnerability.class
    )
    @ApiResponses(value = {
            @ApiResponse(code = 401, message = "Unauthorized"),
            @ApiResponse(code = 404, message = "The vulnerability could not be found")
    })
    @PermissionRequired(Permissions.Constants.VIEW_PORTFOLIO)
    public Response getVulnerabilityByUuid(@ApiParam(value = "The UUID of the vulnerability", required = true)
                                             @PathParam("uuid") String uuid) {
        try (QueryManager qm = new QueryManager()) {
            final Vulnerability vulnerability = qm.getObjectByUuid(Vulnerability.class, uuid);
            if (vulnerability != null) {
                return Response.ok(vulnerability).build();
            } else {
                return Response.status(Response.Status.NOT_FOUND).entity("The vulnerability could not be found.").build();
            }
        }
    }

    @GET
    @Path("/source/{source}/vuln/{vuln}")
    @Produces(MediaType.APPLICATION_JSON)
    @ApiOperation(
            value = "Returns a specific vulnerability",
            response = Vulnerability.class
    )
    @ApiResponses(value = {
            @ApiResponse(code = 401, message = "Unauthorized"),
            @ApiResponse(code = 404, message = "The vulnerability could not be found")
    })
    @PermissionRequired(Permissions.Constants.VIEW_PORTFOLIO)
    public Response getVulnerabilityByVulnId(@PathParam("source") String source,
                                             @PathParam("vuln") String vuln) {
        try (QueryManager qm = new QueryManager()) {
            final Vulnerability vulnerability = qm.getVulnerabilityByVulnId(source, vuln);
            if (vulnerability != null) {
                return Response.ok(vulnerability).build();
            } else {
                return Response.status(Response.Status.NOT_FOUND).entity("The vulnerability could not be found.").build();
            }
        }
    }

    @GET
    @Path("/source/{source}/vuln/{vuln}/projects")
    @Produces(MediaType.APPLICATION_JSON)
    @ApiOperation(
            value = "Returns a list of all projects affected by a specific vulnerability",
            response = Project.class,
            responseContainer = "List",
            responseHeaders = @ResponseHeader(name = TOTAL_COUNT_HEADER, response = Long.class, description = "The total number of projects")
    )
    @ApiResponses(value = {
            @ApiResponse(code = 401, message = "Unauthorized"),
            @ApiResponse(code = 404, message = "The vulnerability could not be found")
    })
    @PermissionRequired(Permissions.Constants.VIEW_PORTFOLIO)
    public Response getAffectedProject(@PathParam("source") String source,
                                                @PathParam("vuln") String vuln) {
        try (QueryManager qm = new QueryManager(getAlpineRequest())) {
            final Vulnerability vulnerability = qm.getVulnerabilityByVulnId(source, vuln);
            if (vulnerability != null) {
                final List<Project> projects = qm.getProjects(vulnerability);
                final long totalCount = projects.size();
                return Response.ok(projects).header(TOTAL_COUNT_HEADER, totalCount).build();
            } else {
                return Response.status(Response.Status.NOT_FOUND).entity("The vulnerability could not be found.").build();
            }
        }
    }

    @GET
    @Produces(MediaType.APPLICATION_JSON)
    @ApiOperation(
            value = "Returns a list of all vulnerabilities",
            response = Vulnerability.class,
            responseContainer = "List",
            responseHeaders = @ResponseHeader(name = TOTAL_COUNT_HEADER, response = Long.class, description = "The total number of vulnerabilities")
    )
    @ApiResponses(value = {
            @ApiResponse(code = 401, message = "Unauthorized")
    })
    @PermissionRequired(Permissions.Constants.VIEW_PORTFOLIO)
    public Response getAllVulnerabilities() {
        try (QueryManager qm = new QueryManager(getAlpineRequest())) {
            final PaginatedResult result = qm.getVulnerabilities();
            return Response.ok(result.getObjects()).header(TOTAL_COUNT_HEADER, result.getTotal()).build();
        }
    }

    @PUT
    @Consumes(MediaType.APPLICATION_JSON)
    @Produces(MediaType.APPLICATION_JSON)
    @ApiOperation(
            value = "Creates a new vulnerability",
            response = Vulnerability.class,
            code = 201
    )
    @ApiResponses(value = {
            @ApiResponse(code = 401, message = "Unauthorized"),
            @ApiResponse(code = 409, message = "A vulnerability with the specified vulnId already exists")
    })
    @PermissionRequired(Permissions.Constants.PORTFOLIO_MANAGEMENT)
    public Response createVulnerability(Vulnerability jsonVulnerability) {
        final Validator validator = super.getValidator();
        failOnValidationError(
                validator.validateProperty(jsonVulnerability, "vulnId"),
                validator.validateProperty(jsonVulnerability, "title"),
                validator.validateProperty(jsonVulnerability, "subTitle"),
                validator.validateProperty(jsonVulnerability, "description"),
                validator.validateProperty(jsonVulnerability, "recommendation"),
                validator.validateProperty(jsonVulnerability, "references"),
                validator.validateProperty(jsonVulnerability, "credits"),
                validator.validateProperty(jsonVulnerability, "created"),
                validator.validateProperty(jsonVulnerability, "published"),
                validator.validateProperty(jsonVulnerability, "updated"),
                validator.validateProperty(jsonVulnerability, "cwe"),
                validator.validateProperty(jsonVulnerability, "cvssV2Vector"),
                validator.validateProperty(jsonVulnerability, "cvssV3Vector"),
                validator.validateProperty(jsonVulnerability, "vulnerableVersions"),
                validator.validateProperty(jsonVulnerability, "patchedVersions")
        );

        try (QueryManager qm = new QueryManager()) {
            Vulnerability vulnerability = qm.getVulnerabilityByVulnId(
                    Vulnerability.Source.INTERNAL, jsonVulnerability.getVulnId().trim());
            if (vulnerability == null) {
                Cwe cwe = null;
                if (jsonVulnerability.getCwe() != null) {
                    cwe = qm.getCweById(jsonVulnerability.getCwe().getCweId());
                }
                jsonVulnerability.setCwe(cwe);
                recalculateScoresFromVector(jsonVulnerability);
                jsonVulnerability.setSource(Vulnerability.Source.INTERNAL);
                vulnerability = qm.createVulnerability(jsonVulnerability, true);
                return Response.status(Response.Status.CREATED).entity(vulnerability).build();
            } else {
                return Response.status(Response.Status.CONFLICT).entity("A vulnerability with the specified vulnId already exists.").build();
            }
        }
    }

    @POST
    @Consumes(MediaType.APPLICATION_JSON)
    @Produces(MediaType.APPLICATION_JSON)
    @ApiOperation(
            value = "Updates an internal vulnerability",
            response = Project.class
    )
    @ApiResponses(value = {
            @ApiResponse(code = 401, message = "Unauthorized"),
            @ApiResponse(code = 404, message = "The vulnerability could not be found"),
            @ApiResponse(code = 406, message = "The vulnId may not be changed")
    })
    @PermissionRequired(Permissions.Constants.PORTFOLIO_MANAGEMENT)
    public Response updateVulnerability(Vulnerability jsonVuln) {
        final Validator validator = super.getValidator();
        failOnValidationError(
                validator.validateProperty(jsonVuln, "title"),
                validator.validateProperty(jsonVuln, "subTitle"),
                validator.validateProperty(jsonVuln, "description"),
                validator.validateProperty(jsonVuln, "recommendation"),
                validator.validateProperty(jsonVuln, "references"),
                validator.validateProperty(jsonVuln, "credits"),
                validator.validateProperty(jsonVuln, "created"),
                validator.validateProperty(jsonVuln, "published"),
                validator.validateProperty(jsonVuln, "updated"),
                validator.validateProperty(jsonVuln, "cwe"),
                validator.validateProperty(jsonVuln, "cvssV2Vector"),
                validator.validateProperty(jsonVuln, "cvssV3Vector"),
                validator.validateProperty(jsonVuln, "vulnerableVersions"),
                validator.validateProperty(jsonVuln, "patchedVersions")
        );
        try (QueryManager qm = new QueryManager()) {
            Vulnerability vulnerability = qm.getObjectByUuid(Vulnerability.class, jsonVuln.getUuid());
            if (vulnerability != null && Vulnerability.Source.INTERNAL.name().equals(vulnerability.getSource())) {
                if (!vulnerability.getVulnId().equals(jsonVuln.getVulnId())) {
                    return Response.status(Response.Status.NOT_ACCEPTABLE).entity("The vulnId may not be changed.").build();
                }
                Cwe cwe = null;
                if (jsonVuln.getCwe() != null) {
                    cwe = qm.getCweById(jsonVuln.getCwe().getCweId());
                }
                jsonVuln.setCwe(cwe);
                recalculateScoresFromVector(jsonVuln);
                vulnerability = qm.updateVulnerability(jsonVuln, true);
                return Response.ok(vulnerability).build();
            } else {
                return Response.status(Response.Status.NOT_FOUND).entity("The vulnerability could not be found.").build();
            }
        }
    }

    private void recalculateScoresFromVector(Vulnerability vuln) {
        // Recalculate V2 score based on vector passed to resource and normalize vector
        final Cvss v2 = Cvss.fromVector(vuln.getCvssV2Vector());
        if (v2 != null) {
            final Score score = v2.calculateScore();
            vuln.setCvssV2BaseScore(BigDecimal.valueOf(score.getBaseScore()));
            vuln.setCvssV2ImpactSubScore(BigDecimal.valueOf(score.getImpactSubScore()));
            vuln.setCvssV2ExploitabilitySubScore(BigDecimal.valueOf(score.getExploitabilitySubScore()));
            vuln.setCvssV2Vector(v2.getVector());
        }

        // Recalculate V3 score based on vector passed to resource and normalize vector
        final Cvss v3 = Cvss.fromVector(vuln.getCvssV3Vector());
        if (v3 != null) {
            final Score score = v3.calculateScore();
            vuln.setCvssV3BaseScore(BigDecimal.valueOf(score.getBaseScore()));
            vuln.setCvssV3ImpactSubScore(BigDecimal.valueOf(score.getImpactSubScore()));
            vuln.setCvssV3ExploitabilitySubScore(BigDecimal.valueOf(score.getExploitabilitySubScore()));
            vuln.setCvssV3Vector(v3.getVector());
        }
    }

    @POST
    @Path("/source/{source}/vuln/{vulnId}/component/{component}")
    @Consumes(MediaType.APPLICATION_JSON)
    @Produces(MediaType.APPLICATION_JSON)
    @ApiOperation(
            value = "Assigns a vulnerability to a component"
    )
    @ApiResponses(value = {
            @ApiResponse(code = 401, message = "Unauthorized"),
            @ApiResponse(code = 404, message = "The vulnerability or component could not be found")
    })
    @PermissionRequired(Permissions.Constants.PORTFOLIO_MANAGEMENT)
    public Response assignVulnerability(@ApiParam(value = "The vulnerability source", required = true)
                                        @PathParam("source") String source,
                                        @ApiParam(value = "The vulnId", required = true)
                                        @PathParam("vulnId") String vulnId,
                                        @ApiParam(value = "The UUID of the component", required = true)
                                        @PathParam("component") String componentUuid) {
        try (QueryManager qm = new QueryManager()) {
            Vulnerability vulnerability = qm.getVulnerabilityByVulnId(source, vulnId);
            if (vulnerability == null) {
                return Response.status(Response.Status.NOT_FOUND).entity("The vulnerability could not be found.").build();
            }
            Component component = qm.getObjectByUuid(Component.class, componentUuid);
            if (component == null) {
                return Response.status(Response.Status.NOT_FOUND).entity("The component could not be found.").build();
            }
            qm.addVulnerability(vulnerability, component);
            return Response.ok().build();
        }
    }

    @POST
    @Path("/{uuid}/component/{component}")
    @Consumes(MediaType.APPLICATION_JSON)
    @Produces(MediaType.APPLICATION_JSON)
    @ApiOperation(
            value = "Assigns a vulnerability to a component"
    )
    @ApiResponses(value = {
            @ApiResponse(code = 401, message = "Unauthorized"),
            @ApiResponse(code = 404, message = "The vulnerability or component could not be found")
    })
    @PermissionRequired(Permissions.Constants.PORTFOLIO_MANAGEMENT)
    public Response assignVulnerability(@ApiParam(value = "The UUID of the vulnerability", required = true)
                                        @PathParam("uuid") String uuid,
                                        @ApiParam(value = "The UUID of the component", required = true)
                                        @PathParam("component") String componentUuid) {
        try (QueryManager qm = new QueryManager()) {
            Vulnerability vulnerability = qm.getObjectByUuid(Vulnerability.class, uuid);
            if (vulnerability == null) {
                return Response.status(Response.Status.NOT_FOUND).entity("The vulnerability could not be found.").build();
            }
            Component component = qm.getObjectByUuid(Component.class, componentUuid);
            if (component == null) {
                return Response.status(Response.Status.NOT_FOUND).entity("The component could not be found.").build();
            }
            qm.addVulnerability(vulnerability, component);
            return Response.ok().build();
        }
    }

    @DELETE
    @Path("/source/{source}/vuln/{vulnId}/component/{component}")
    @Consumes(MediaType.APPLICATION_JSON)
    @Produces(MediaType.APPLICATION_JSON)
    @ApiOperation(
            value = "Removes assignment of a vulnerability from a component"
    )
    @ApiResponses(value = {
            @ApiResponse(code = 401, message = "Unauthorized"),
            @ApiResponse(code = 404, message = "The vulnerability or component could not be found")
    })
    @PermissionRequired(Permissions.Constants.PORTFOLIO_MANAGEMENT)
    public Response unassignVulnerability(@ApiParam(value = "The vulnerability source", required = true)
                                          @PathParam("source") String source,
                                          @ApiParam(value = "The vulnId", required = true)
                                          @PathParam("vulnId") String vulnId,
                                          @ApiParam(value = "The UUID of the component", required = true)
                                          @PathParam("component") String componentUuid) {
        try (QueryManager qm = new QueryManager()) {
            Vulnerability vulnerability = qm.getVulnerabilityByVulnId(source, vulnId);
            if (vulnerability == null) {
                return Response.status(Response.Status.NOT_FOUND).entity("The vulnerability could not be found.").build();
            }
            Component component = qm.getObjectByUuid(Component.class, componentUuid);
            if (component == null) {
                return Response.status(Response.Status.NOT_FOUND).entity("The component could not be found.").build();
            }
            qm.removeVulnerability(vulnerability, component);
            return Response.ok().build();
        }
    }

    @DELETE
    @Path("/{uuid}/component/{component}")
    @Consumes(MediaType.APPLICATION_JSON)
    @Produces(MediaType.APPLICATION_JSON)
    @ApiOperation(
            value = "Removes assignment of a vulnerability from a component"
    )
    @ApiResponses(value = {
            @ApiResponse(code = 401, message = "Unauthorized"),
            @ApiResponse(code = 404, message = "The vulnerability or component could not be found")
    })
    @PermissionRequired(Permissions.Constants.PORTFOLIO_MANAGEMENT)
    public Response unassignVulnerability(@ApiParam(value = "The UUID of the vulnerability", required = true)
                                          @PathParam("uuid") String uuid,
                                          @ApiParam(value = "The UUID of the component", required = true)
                                          @PathParam("component") String componentUuid) {
        try (QueryManager qm = new QueryManager()) {
            Vulnerability vulnerability = qm.getObjectByUuid(Vulnerability.class, uuid);
            if (vulnerability == null) {
                return Response.status(Response.Status.NOT_FOUND).entity("The vulnerability could not be found.").build();
            }
            Component component = qm.getObjectByUuid(Component.class, componentUuid);
            if (component == null) {
                return Response.status(Response.Status.NOT_FOUND).entity("The component could not be found.").build();
            }
            qm.removeVulnerability(vulnerability, component);
            return Response.ok().build();
        }
    }
}