vulnerability.py

"""Vulnerability API."""
# Standard Python Libraries
from collections import defaultdict
import csv
from datetime import datetime, timedelta
import io
import logging

# Third-Party Libraries
from django.core.paginator import Paginator
from django.db.models import Count, Prefetch, Q
from django.shortcuts import get_object_or_404
from fastapi import HTTPException
from xfd_mini_dl.models import (  # Add Domain import here when done
    CisaKevCatalog,
    CredentialBreaches,
    CredentialExposures,
    Service,
    ShodanVulns,
    Ticket,
    TicketEvent,
    Vulnerability,
)

from ..auth import get_org_memberships, is_analytics_user, is_global_view_admin
from ..helpers.filter_helpers import apply_vuln_filters, sort_direction
from ..helpers.s3_client import S3Client
from ..helpers.uuid_helpers import is_valid_uuid
from ..models import Domain
from ..schema_models.vulnerability import (
    GetVulnerabilityResponse,
    UserExposure,
    VsTicketEvent,
    VsVulnerabilityResponse,
)
from ..schema_models.vulnerability import (
    VulnerabilityFilters,
    VulnerabilityGroupResponse,
    VulnerabilitySearch,
)
from ..schema_models.vulnerability import Vulnerability as VulnerabilitySchema

LOGGER = logging.getLogger(__name__)


def get_default_dates():
    """Return default earliest and latest dates."""
    # Set latest_date default to today
    latest_date = datetime.now()
    print(f"Setting default search latest_date: {latest_date}")
    # Set earliest_date default to 30 days ago
    earliest_date = latest_date - timedelta(days=30)
    print(f"Setting default search earliest_date: {earliest_date}")
    return earliest_date.isoformat(), latest_date.isoformat()


def get_vulnerability_by_id(vulnerability_id, current_user):
    """Get vulnerability by id."""
    try:
        # Initialize a VulnerabilitySearch with filters.id set to the vulnerability_id
        search = VulnerabilitySearch(
            page=1,
            sort="ASC",
            order="id",
            filters=VulnerabilityFilters(id=str(vulnerability_id)),
        )
        vulnerabilities, count = search_vulnerabilities(search, current_user)

        if count == 0:
            raise HTTPException(status_code=404, detail="Vulnerability not found.")

        enrich_kev_fields(vulnerabilities)

        # Assuming vulnerabilities is a list, get the first (and only) item
        vulnerability = vulnerabilities[0]

        # Serialize using Pydantic model
        return GetVulnerabilityResponse.from_orm(vulnerability)

    except HTTPException as he:
        raise he
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


def get_vulnerability_by_scan_source_and_id(
    scan_source: str, vulnerability_id: str, current_user
):
    """Get vulnerability by scan_source and id."""
    try:
        # Initialize a VulnerabilitySearch with filters.id set to the vulnerability_id
        # Permissions check
        filters = {}

        if (
            not is_global_view_admin(current_user)
            and not is_analytics_user(current_user)
            and not current_user.user_type == "regionalAdmin"
        ):
            org_ids = get_org_memberships(current_user)
            if not org_ids:
                raise HTTPException(
                    status_code=404, detail="No valid organization linked to user."
                )  # User has no accessible organizations
            filters["organization_id__in"] = org_ids

        # Regional Admins can only view vulnerabilities in their region
        if current_user.user_type == "regionalAdmin" and current_user.region_id:
            filters["organization__region_id"] = current_user.region_id

        if scan_source == "vuln_scanning_tickets":
            try:
                severity_map = {
                    "0.00": "N/A",
                    "1.00": "Low",
                    "2.00": "Medium",
                    "3.00": "High",
                    "4.00": "Critical",
                }
                ticket = (
                    Ticket.objects.select_related("organization")
                    .prefetch_related(
                        Prefetch(
                            "ticket_events",
                            queryset=TicketEvent.objects.select_related(
                                "vuln_scan"
                            ).order_by("-event_timestamp")[:10],
                            to_attr="prefetched_events",
                        )
                    )
                    .get(id=vulnerability_id, **filters)
                )

                ticket_response = {
                    "id": str(ticket.id),
                    "cve_string": ticket.cve_string,
                    # "cve_id":
                    "cvss_base_score": ticket.cvss_base_score,
                    "cvss_version": ticket.cvss_version,
                    "vuln_name": ticket.vuln_name,
                    "cvss_score_source": ticket.cvss_score_source,
                    "cvss_severity": severity_map.get(str(ticket.cvss_severity), None),
                    "vpr_score": ticket.vpr_score,
                    "false_positive": ticket.false_positive,
                    "ip_string": ticket.ip_string,
                    # "ip_id":
                    "updated_timestamp": ticket.updated_timestamp,
                    "location_longitude": ticket.location_longitude,
                    "location_latitude": ticket.location_latitude,
                    "found_in_latest_host_scan": ticket.found_in_latest_host_scan,
                    "organization_name": ticket.organization.name,
                    "vuln_port": ticket.vuln_port,
                    "port_protocol": ticket.port_protocol,
                    # "snapshots_bool":ticket.snapshots_bool
                    "vuln_source": ticket.vuln_source,
                    # "vuln_source_id":ticket.vuln_source
                    "closed_timestamp": ticket.closed_timestamp,
                    "opened_timestamp": ticket.opened_timestamp,
                    "is_open": ticket.is_open,
                    "is_kev": ticket.is_kev,
                    "is_risky": ticket.is_risky,
                    "service_name": ticket.service_name,
                    "operating_system": ticket.operating_system,
                    "ticket_events": [
                        VsTicketEvent(
                            id=str(event.id),
                            action=event.action,
                            reason=event.reason,
                            event_timestamp=event.event_timestamp,
                            delta=event.delta,
                            vuln_scan_description=event.vuln_scan.description
                            if event.vuln_scan
                            else None,
                            vuln_scan_cpe=event.vuln_scan.cpe
                            if event.vuln_scan
                            else None,
                            vuln_scan_solution=event.vuln_scan.solution
                            if event.vuln_scan
                            else None,
                            vuln_scan_synopsis=event.vuln_scan.synopsis
                            if event.vuln_scan
                            else None,
                            # add other fields you need
                        )
                        for event in ticket.prefetched_events
                    ],
                }

                return VsVulnerabilityResponse(**ticket_response)
            except Ticket.DoesNotExist:
                raise HTTPException(status_code=404, detail="VS Ticket not found")
        elif scan_source == "credential_breach":
            try:
                exposure = CredentialExposures.objects.select_related(
                    "credential_breach"
                ).get(credential_exposures_uid=vulnerability_id, **filters)

                related_breach = exposure.credential_breach

                related_exposures = related_breach.exposures.select_related(
                    "organization"
                ).filter(**filters)

                grouped_exposures = defaultdict(list)

                for exposure in related_exposures:
                    org = exposure.organization
                    grouped_exposures[org.name].append(
                        UserExposure(
                            email=exposure.email,
                            name=exposure.name,
                            has_password=bool(exposure.password),
                            uid=str(exposure.credential_exposures_uid),
                            # any other exposure fields you want
                        )
                    )

                response = {
                    "breach_name": related_breach.breach_name,
                    "description": related_breach.description,
                    "breach_date": related_breach.breach_date,
                    "added_date": related_breach.added_date,
                    "modified_date": related_breach.modified_date,
                    "data_classes": related_breach.data_classes,
                    "password_included": related_breach.password_included,
                    "data_source": related_breach.data_source.name,
                    "exposures_by_org": grouped_exposures,
                }
                return response

            except CredentialExposures.DoesNotExist:
                raise HTTPException(
                    status_code=404, detail="Credential exposure not found"
                )
            except CredentialExposures.MultipleObjectsReturned:
                raise HTTPException(
                    status_code=500, detail="Multiple exposures found. Expected one."
                )
            except CredentialBreaches.DoesNotExist:
                raise HTTPException(
                    status_code=404, detail="Associated CredentialBreach not found"
                )

        elif scan_source == "shodan_vulnerability":
            shodan_vuln = ShodanVulns.objects.select_related("organization").get(
                shodan_vuln_uid=vulnerability_id, **filters
            )

            vuln_response = {
                "shodan_identified_organization_name": shodan_vuln.organization_name,
                "ip_string": shodan_vuln.ip_string,
                "port": shodan_vuln.port,
                "protocol": shodan_vuln.protocol,
                "timestamp": shodan_vuln.timestamp,
                "cve": shodan_vuln.cve,
                "severity": shodan_vuln.severity,
                "cvss": shodan_vuln.cvss,
                "summary": shodan_vuln.summary,
                "product": shodan_vuln.product,
                "attack_vector": shodan_vuln.attack_vector,
                "av_description": shodan_vuln.av_description,
                "attack_complexity": shodan_vuln.attack_complexity,
                "ac_description": shodan_vuln.ac_description,
                "confidentiality_impact": shodan_vuln.confidentiality_impact,
                "ci_description": shodan_vuln.ci_description,
                "integrity_impact": shodan_vuln.integrity_impact,
                "ii_description": shodan_vuln.ii_description,
                "availability_impact": shodan_vuln.availability_impact,
                "ai_description": shodan_vuln.ai_description,
                "tags": shodan_vuln.tags,
                "domains": shodan_vuln.domains,
                "hostnames": shodan_vuln.hostnames,
                "isp": shodan_vuln.isp,
                "asn": shodan_vuln.asn,
                "type": shodan_vuln.type,
                "name": shodan_vuln.name,
                "potential_vulns": shodan_vuln.potential_vulns,
                "mitigation": shodan_vuln.mitigation,
                "server": shodan_vuln.server,
                "is_verified": shodan_vuln.is_verified,
                "banner": shodan_vuln.banner,
                "version": shodan_vuln.version,
                "cpe": shodan_vuln.cpe,
                "ip_uid": str(shodan_vuln.ip_id),
            }

            return vuln_response
        else:
            raise HTTPException(
                status_code=404, detail="Provided scan_source not found"
            )

    except HTTPException as he:
        raise he
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


def update_vulnerability(
    vulnerability_id, vulnerability_data: VulnerabilitySchema, current_user
):
    """Update vulnerability by id."""
    try:
        if not is_global_view_admin(current_user):
            raise HTTPException(status_code=401, detail="Unauthorized")

        # Validate the vulnerability ID
        if not is_valid_uuid(vulnerability_id):
            raise HTTPException(status_code=404, detail="Vulnerability not found")

        # Fetch the existing vulnerability
        vulnerability = Vulnerability.objects.get(id=vulnerability_id)

        # Create a mapping of fields to update
        fields_to_update = {
            "title": vulnerability_data.title,
            "cve": vulnerability_data.cve,
            "cwe": vulnerability_data.cwe,
            "cpe": vulnerability_data.cpe,
            "description": vulnerability_data.description,
            "references": vulnerability_data.references,
            "cvss": vulnerability_data.cvss,
            "severity": vulnerability_data.severity,
            "needs_population": vulnerability_data.needs_population,
            "state": vulnerability_data.state,
            "substate": vulnerability_data.substate,
            "source": vulnerability_data.source,
            "notes": vulnerability_data.notes,
            "actions": vulnerability_data.actions,
            "structured_data": vulnerability_data.structured_data,
            "is_kev": vulnerability_data.is_kev,
            "domain": vulnerability_data.domain_id,
            "service": vulnerability_data.service_id,
        }

        # Update fields that are not None
        for field, value in fields_to_update.items():
            if value is not None:
                if field == "domain":
                    # Handle domain ID to fetch the Domain instance
                    domain_instance = get_object_or_404(
                        Domain, id=vulnerability_data.domain_id
                    )
                    vulnerability.domain = domain_instance
                elif field == "service":
                    # Handle service ID to fetch the Service instance
                    service_instance = get_object_or_404(
                        Service, id=vulnerability_data.service_id
                    )
                    vulnerability.service = service_instance
                else:
                    setattr(vulnerability, field, value)
        # Save the updated vulnerability object
        vulnerability.save()
        return vulnerability

    except Vulnerability.DoesNotExist:
        raise HTTPException(status_code=404, detail="Vulnerability not found.")
    except Exception as e:
        raise HTTPException(
            status_code=500, detail="Internal Server Error: {}".format(e)
        )


def search_vulnerabilities(vulnerability_search: VulnerabilitySearch, current_user):
    """List vulnerabilities by search filter."""
    try:
        # Base query with related data
        vulnerabilities = Vulnerability.objects.select_related(
            "domain", "domain__organization"
        ).order_by(
            sort_direction(vulnerability_search.sort, vulnerability_search.order)
        )

        # Permissions check
        if (
            not is_global_view_admin(current_user)
            and not is_analytics_user(current_user)
            and not current_user.user_type == "regionalAdmin"
        ):
            org_ids = get_org_memberships(current_user)
            if not org_ids:
                return [], 0  # User has no accessible organizations
            vulnerabilities = vulnerabilities.filter(
                domain__organization_id__in=org_ids
            )
        # Regional Admins can only view vulnerabilities in their region
        if current_user.user_type == "regionalAdmin" and current_user.region_id:
            vulnerabilities = vulnerabilities.filter(
                domain__organization__region_id=current_user.region_id
            )

        # Apply Open Status only filters by default unless showAll is set to True
        if not vulnerability_search.show_all:
            vulnerabilities = vulnerabilities.filter(Q(state="open"))

        # # Set default dates if not provided
        if vulnerability_search.filters and not getattr(
            vulnerability_search.filters, "earliest_date", None
        ):
            earliest_date, _ = get_default_dates()
            vulnerability_search.filters.earliest_date = earliest_date

        if vulnerability_search.filters and not getattr(
            vulnerability_search.filters, "latest_date", None
        ):
            _, latest_date = get_default_dates()
            vulnerability_search.filters.latest_date = latest_date

        # Apply additional filters
        if vulnerability_search.filters:
            vulnerabilities = apply_vuln_filters(
                vulnerabilities, vulnerability_search.filters
            )

        # Apply grouping if specified
        if vulnerability_search.group_by:
            vulnerabilities = (
                vulnerabilities.values(vulnerability_search.group_by)
                .annotate(cnt=Count("id"))
                .order_by("-cnt")
            )

        # Paginate the results
        paginator = Paginator(vulnerabilities, vulnerability_search.page_size)
        page_obj = paginator.get_page(vulnerability_search.page)

        # If group_by is used, handle raw values
        if vulnerability_search.group_by:
            result = list(page_obj)
            count = paginator.count
            return result, count

        result = list(page_obj.object_list)
        count = paginator.count

        return result, count

    except HTTPException as he:
        raise he
    except Vulnerability.DoesNotExist as e:
        print(e)
        raise HTTPException(status_code=404, detail="Vulnerability not found.")
    except Exception as e:
        print(e)
        raise HTTPException(status_code=500, detail=str(e))


def export_vulnerabilities(vulnerability_search, current_user):
    """Export vulnerabilities into a CSV and upload to S3."""
    try:
        # Retrieve vulnerabilities based on search filters
        vulnerability_search.page_size = 1000  # Set to max 1000 results
        vulnerabilities, count = search_vulnerabilities(
            vulnerability_search, current_user
        )

        if vulnerability_search.group_by:
            # Handle grouped data
            result = [
                VulnerabilityGroupResponse.model_validate(v) for v in vulnerabilities
            ]
            # Define CSV fields for grouped data
            group_field = vulnerability_search.group_by
            csv_fields = [group_field, "cnt"]
            # Create CSV content
            output = io.StringIO()
            writer = csv.DictWriter(output, fieldnames=csv_fields)
            writer.writeheader()
            for v in result:
                row = {group_field: getattr(v, group_field, ""), "cnt": v.cnt}
                writer.writerow(row)
            csv_content = output.getvalue()
        else:
            # Serialize ORM instances to Pydantic models
            serialized_vulnerabilities = [
                GetVulnerabilityResponse.model_validate(v) for v in vulnerabilities
            ]
            # Convert to list of dictionaries
            data = [v.dict() for v in serialized_vulnerabilities]
            # Map organization and domain names
            for entry in data:
                entry["organization"] = (
                    entry["domain"]["organization"]["name"]
                    if entry.get("domain") and entry["domain"].get("organization")
                    else ""
                )
                entry["domain"] = entry["domain"]["name"] if entry.get("domain") else ""
            # Define CSV fields
            csv_fields = [
                "organization",
                "domain",
                "title",
                "description",
                "cve",
                "is_kev",
                "cwe",
                "cpe",
                "cvss",
                "severity",
                "state",
                "substate",
                "last_seen",
                "created_at",
                "id",
            ]
            # Create CSV content
            output = io.StringIO()
            writer = csv.DictWriter(output, fieldnames=csv_fields)
            writer.writeheader()
            for entry in data:
                row = {field: entry.get(field, "") for field in csv_fields}
                writer.writerow(row)
            csv_content = output.getvalue()

        # Initialize S3 client
        client = S3Client()

        # Save CSV to S3
        url = client.save_csv(csv_content, "vulnerabilities")

        return {"url": url}

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


def is_publicid_in_kev(publicid: str) -> bool:
    """Check if the given public ID matches a known CVE ID in the KEV catalog."""
    if not publicid:
        return False
    return CisaKevCatalog.objects.filter(cve_id=publicid).exists()


def enrich_kev_fields(vulnerabilities):
    """Update vulnerabilities with is_kev and kev_results using KEV catalog lookup."""
    for v in vulnerabilities:
        try:
            if v.cve:
                if v.is_kev is None:
                    v.is_kev = is_publicid_in_kev(v.cve)

                if not v.kev_results:
                    kev_entry = CisaKevCatalog.objects.filter(cve_id=v.cve).first()
                    if kev_entry:
                        v.kev_results = {
                            "vendor": kev_entry.vendor_project,
                            "product": kev_entry.product,
                            "vuln_name": kev_entry.vulnerability_name,
                            "cwe": kev_entry.cwe_id,
                            "date_added": str(kev_entry.date_added),
                            "ransomware": kev_entry.known_ransomware_campaign_use,
                        }

            # Set fallback defaults if still None
            if v.kev_results is None:
                v.kev_results = {}

            if v.is_kev is None:
                v.is_kev = False

        except Exception as e:
            LOGGER.warning(f"Skipping vulnerability {getattr(v, 'id', 'unknown')}: {e}")