# Copyright 2023 Google LLC
#
# Use of this source code is governed by a BSD-style
# license that can be found in the LICENSE file or at
# https://developers.google.com/open-source/licenses/bsd

"""Manage OSV Vulnerabilities and their related data structures."""

from collections.abc import MutableMapping
import copy
import dataclasses
import enum
import functools
import json
import re
from typing import Any, Dict, List, Mapping, Optional, Sequence, Union
from google.protobuf import json_format
from osv import vulnerability_pb2

from vanir import signature

OSV_AFFECTED = 'affected'
OSV_PACKAGE = 'package'
OSV_PACKAGE_NAME = 'name'
OSV_ECOSYSTEM = 'ecosystem'
OSV_VERSIONS = 'versions'
OSV_ECOSYSTEM_SPECIFIC = 'ecosystem_specific'
OSV_DATABASE_SPECIFIC = 'database_specific'
OSV_VANIR_SIGNATURES = 'vanir_signatures'
OSV_SIGNATURE_ID = 'id'


@enum.unique
class MetaPackage(enum.Enum):
  """Special set of OSV packages.

  An actual package name of a metapacakge can contain postfix strings.
  E.g., ":linux_kernel:", ":linux_kernel:Qualcomm", ":linux_kernel:Mediatek".
  """
  UNKNOWN = ':unknown:'
  ANDROID_KERNEL = ':linux_kernel:'
  ANDROID_MODEM = ':modem:'

  @functools.cached_property
  def package_pattern(self) -> re.Pattern[str]:
    """Returns regex pattern of packages belonging to the meta-package."""
    pattern = re.escape(self.value) + '.*'
    return re.compile(pattern)


class AffectedEntry:
  """Class for representing an affected package entry in OSV format."""

  def __init__(
      self,
      osv_affected: Mapping[str, Any],
      store_signatures_in_legacy_location: bool = False,
  ):
    self._signatures_location = (
        OSV_ECOSYSTEM_SPECIFIC if store_signatures_in_legacy_location
        else OSV_DATABASE_SPECIFIC
    )
    package = osv_affected.get(OSV_PACKAGE)
    if package and (
        not package.get(OSV_ECOSYSTEM)
        or not package.get(OSV_PACKAGE_NAME)
    ):
      raise ValueError(f'Invalid package info: {osv_affected}')

    self._osv_affected: Dict[str, Any] = dict(copy.deepcopy(osv_affected))
    if OSV_ECOSYSTEM_SPECIFIC not in self._osv_affected:
      self._osv_affected[OSV_ECOSYSTEM_SPECIFIC] = {}
    if OSV_DATABASE_SPECIFIC not in self._osv_affected:
      self._osv_affected[OSV_DATABASE_SPECIFIC] = {}
    if (
        OSV_VANIR_SIGNATURES in self.ecosystem_specific
        and OSV_VANIR_SIGNATURES in self.database_specific
    ):
      raise ValueError(
          'Signatures exist in both ecosystem_specific and database_specific: '
          f'{osv_affected}'
      )
    # Convert all signatures from dictionaries/JSON to Signature objects.
    if OSV_VANIR_SIGNATURES in self.database_specific:
      original_signatures = self.database_specific[OSV_VANIR_SIGNATURES]
    elif OSV_VANIR_SIGNATURES in self.ecosystem_specific:
      original_signatures = self.ecosystem_specific[OSV_VANIR_SIGNATURES]
    else:
      original_signatures = []
    self.vanir_signatures = [
        sig if isinstance(sig, signature.Signature) else
        signature.Signature.from_osv_dict(sig)
        for sig in original_signatures
    ]

  @property
  def ecosystem(self) -> str:
    """Returns the ecosystem of the affected package."""
    return self._osv_affected.get(OSV_PACKAGE, {}).get(OSV_ECOSYSTEM, '')

  @ecosystem.setter
  def ecosystem(self, value: str):
    if OSV_PACKAGE not in self._osv_affected:
      self._osv_affected[OSV_PACKAGE] = {}
    self._osv_affected[OSV_PACKAGE][OSV_ECOSYSTEM] = value

  @property
  def osv_package_name(self) -> str:
    """Returns the raw OSV package name of the affected package."""
    return self._osv_affected.get(OSV_PACKAGE, {}).get(OSV_PACKAGE_NAME, '')

  @property
  def package_name(self) -> str:
    """Returns the normalized package name of the affected package.

    Usually, an OSV package name is the package name we want to use. However,
    there are special packages, e.g. ":linux_kernel:Qualcomm" or
    ":linux_kernel:Mediatek", which we normalize to a MetaPackage
    e.g. ":linux_kernel:".
    """
    for meta_package in MetaPackage:
      if self.osv_package_name.startswith(meta_package.value):
        return meta_package.value
    return self.osv_package_name

  @osv_package_name.setter
  def osv_package_name(self, value: str):
    if OSV_PACKAGE not in self._osv_affected:
      self._osv_affected[OSV_PACKAGE] = {}
    self._osv_affected[OSV_PACKAGE][OSV_PACKAGE_NAME] = value

  @property
  def versions(self) -> Sequence[str]:
    """Returns the list of affected versions of this affected entry."""
    return self._osv_affected.get(OSV_VERSIONS, [])

  @property
  def ecosystem_specific(self) -> MutableMapping[str, Any]:
    return self._osv_affected[OSV_ECOSYSTEM_SPECIFIC]

  @property
  def database_specific(self) -> MutableMapping[str, Any]:
    return self._osv_affected[OSV_DATABASE_SPECIFIC]

  @property
  def vanir_signatures(self) -> List[signature.Signature]:
    """Returns the list of Vanir signatures of this affected entry."""
    return self._osv_affected[self._signatures_location][OSV_VANIR_SIGNATURES]

  @vanir_signatures.setter
  def vanir_signatures(self, value: Sequence[signature.Signature]):
    self._osv_affected[self._signatures_location][OSV_VANIR_SIGNATURES] = value

  def sort_vanir_signatures(self):
    """Sorts the Vanir signatures according to their signature ID."""
    self.vanir_signatures.sort(
        key=lambda sig: (
            sig.signature_id if isinstance(sig, signature.Signature) else
            sig[OSV_SIGNATURE_ID],
        )
    )

  def to_osv_dict(self, use_string_hashes: bool = True) -> Mapping[str, Any]:
    """Returns the affected package as a dict in OSV format."""
    osv_affected = copy.deepcopy(self._osv_affected)
    osv_sigs = [
        sig.to_osv_dict(use_string_hashes)
        if isinstance(sig, signature.Signature) else sig
        for sig in self.vanir_signatures
    ]
    # Replace signature objects with their OSV dict representation.
    if OSV_VANIR_SIGNATURES in osv_affected[OSV_ECOSYSTEM_SPECIFIC]:
      del osv_affected[OSV_ECOSYSTEM_SPECIFIC][OSV_VANIR_SIGNATURES]
    if OSV_VANIR_SIGNATURES in osv_affected[OSV_DATABASE_SPECIFIC]:
      del osv_affected[OSV_DATABASE_SPECIFIC][OSV_VANIR_SIGNATURES]
    osv_affected[self._signatures_location][OSV_VANIR_SIGNATURES] = (
        osv_sigs
    )
    if not osv_sigs:
      del osv_affected[self._signatures_location][OSV_VANIR_SIGNATURES]
    # Remove empty ecosystem_specific and database_specific fields.
    if not osv_affected[OSV_ECOSYSTEM_SPECIFIC]:
      del osv_affected[OSV_ECOSYSTEM_SPECIFIC]
    if not osv_affected[OSV_DATABASE_SPECIFIC]:
      del osv_affected[OSV_DATABASE_SPECIFIC]
    return osv_affected

  def to_proto(self) -> vulnerability_pb2.Affected:
    """Converts this AffectedEntry object to OSV affected proto."""
    return json_format.Parse(
        json.dumps(self, default=_to_proto_dict),
        vulnerability_pb2.Affected(),
    )


@dataclasses.dataclass
class Vulnerability:
  """Class for representing a vulnerability entry in OSV format.

  Fields definitions: https://ossf.github.io/osv-schema/
  """

  schema_version: Optional[str] = None
  id: str = ''
  modified: str = ''
  published: Optional[str] = None
  withdrawn: Optional[str] = None
  aliases: Optional[Sequence[str]] = None
  related: Optional[Sequence[str]] = None
  summary: Optional[str] = None
  details: Optional[str] = None
  references: Optional[Sequence[Mapping[str, str]]] = None
  severity: Optional[Mapping[str, str]] = None
  affected: Sequence[AffectedEntry] = ()
  credits: Optional[Sequence[Mapping[str, str]]] = None
  database_specific: Optional[Mapping[str, Any]] = None

  def __init__(
      self,
      osv_vuln: Mapping[str, Any],
      store_signatures_in_legacy_location: bool = False,
  ):
    if not isinstance(osv_vuln, Mapping):
      raise ValueError(f'Not a Mapping: {osv_vuln}')
    for field in dataclasses.fields(self):
      if field.name in osv_vuln and field.name != OSV_AFFECTED:
        setattr(self, field.name, copy.deepcopy(osv_vuln[field.name]))

    # Convert all affected entries to AffectedEntry objects.
    self.affected = []
    for osv_affected in osv_vuln.get(OSV_AFFECTED, []):
      self.affected.append(
          AffectedEntry(osv_affected, store_signatures_in_legacy_location)
      )

    if not self.id or not self.modified or not self.affected:
      raise ValueError('Missing required fields: id, modified, or affected')

  def to_osv_dict(self, use_string_hashes: bool = True) -> Mapping[str, Any]:
    """Returns the vulnerability as a dict in OSV format."""
    osv_vuln = dataclasses.asdict(self)
    if self.affected:
      affected_entries = []
      for affected in osv_vuln[OSV_AFFECTED]:
        # At this point, affected field should be a list of AffectedEntry due
        # to __post_init__(). This assert is needed to make pytype happy.
        assert isinstance(affected, AffectedEntry)
        affected_entries.append(affected.to_osv_dict(use_string_hashes))
      osv_vuln[OSV_AFFECTED] = affected_entries
    empty_entries = {key for key in osv_vuln if not osv_vuln[key]}
    for key in empty_entries:
      del osv_vuln[key]
    return osv_vuln

  def to_proto(self) -> vulnerability_pb2.Vulnerability:
    """Converts this Vulnerability object to OSV vulnerability proto."""
    return json_format.Parse(
        json.dumps(self, default=_to_proto_dict),
        vulnerability_pb2.Vulnerability(),
    )


def _to_proto_dict(
    item: Union[Vulnerability, AffectedEntry, signature.Signature],
) -> Mapping[str, Any]:
  return item.to_osv_dict(use_string_hashes=True)
