vulnerability_manager_test.py

# 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

"""Tests Vulnerability Manager."""

from collections.abc import Mapping
import copy
import dataclasses
import json
import re
from typing import Any
from unittest import mock

from absl import logging
from vanir import osv_client
from vanir import refiner
from vanir import sign_generator
from vanir import signature
from vanir import vulnerability
from vanir import vulnerability_manager
from vanir.code_extractors import code_extractor
from vanir.code_extractors import code_extractor_base

from absl.testing import absltest


def _get_osv_signs(osv_vuln: Mapping[str, Any]) -> Any:
  """Utility function to get the first package's sigs in the vulnerability."""
  return osv_vuln['affected'][0]['ecosystem_specific']['vanir_signatures']


class VulnerabilityManagerTest(absltest.TestCase):

  def setUp(self):
    super().setUp()
    self._test_sign_factory = signature.SignatureFactory('ASB-A-1234')
    self._test_osv_sign = {
        'id': 'ASB-A-1234-b4e3e62e8c',
        'signature_type': 'Function',
        'signature_version': 'v1234',
        'source': 'http://android.googlesource.com/test_source',
        'target': {
            'file': 'test_file.c',
            'function': 'test_func1'
        },
        'deprecated': False,
        'digest': {
            'function_hash': '1234',
            'length': 100
        }
    }
    self._test_signature = signature.FunctionSignature(
        signature_id=self._test_osv_sign['id'],
        signature_version=self._test_osv_sign['signature_version'],
        source=self._test_osv_sign['source'],
        target_file=self._test_osv_sign['target']['file'],
        deprecated=self._test_osv_sign['deprecated'],
        exact_target_file_match_only=self._test_osv_sign.get(
            'exact_target_file_match_only', False
        ),
        match_only_versions=self._test_osv_sign.get(
            'match_only_versions', None
        ),
        truncated_path_level=self._test_osv_sign['target'].get(
            'truncated_path_level'
        ),
        function_hash=self._test_osv_sign['digest']['function_hash'],
        length=self._test_osv_sign['digest']['length'],
        target_function=self._test_osv_sign['target']['function'],
    )

    self._test_osv_vuln = {
        'id':
            'ASB-A-1234',
        'modified':
            '2000-11-11T21:26:24.496435096Z',
        'affected': [{
            'package': {
                'name': 'test_package',
                'ecosystem': 'Android'
            },
            'versions': ['11'],
            'ecosystem_specific': {
                'vanir_signatures': [self._test_osv_sign]
            }
        }]
    }
    self._test_vul = vulnerability.Vulnerability(self._test_osv_vuln)
    self._test_osv_vuln_without_signatures = {
        'id':
            'ASB-A-1234',
        'modified':
            '2000-11-11T21:26:24.496435096Z',
        'affected': [{
            'package': {
                'name': 'test_package',
                'ecosystem': 'Android'
            },
            'versions': ['11']
        }]
    }
    self._test_qc_kernel_osv_vul = {
        'id':
            'ASB-A-2345',
        'modified':
            '2000-12-15T21:26:24Z',
        'affected': [{
            'package': {
                'name': ':linux_kernel:Qualcomm',
                'ecosystem': 'Android'
            },
            'versions': ['Kernel'],
            'ecosystem_specific': {
                'vanir_signatures': []
            }
        }]
    }
    self._mock_commit = mock.create_autospec(
        code_extractor_base.Commit, instance=True
    )
    self._fake_failed_commit_url = code_extractor_base.FailedCommitUrl(
        'https://bad.url.org', ValueError('bad url.')
    )
    self._mock_extract = self.enter_context(
        mock.patch.object(
            code_extractor,
            'extract_for_affected_entry',
            return_value=([self._mock_commit], [self._fake_failed_commit_url]),
            autospec=True,
        )
    )
    self._mock_extract_tip_of_unaffected_versions = self.enter_context(
        mock.patch.object(
            code_extractor,
            'extract_files_at_tip_of_unaffected_versions',
            return_value=([], []),
            autospec=True,
        )
    )
    self._mock_sign_generator = self.enter_context(
        mock.patch.object(
            sign_generator.SignGenerator,
            'generate_signatures_for_commit',
            return_value=[self._test_signature],
            autospec=True,
        )
    )
    # By default, treat all signatures as good (i.e. the refiner does nothing).
    self._mock_refiner = self.enter_context(
        mock.patch.object(
            refiner.Refiner,
            'refine_against_patch_series',
            side_effect=lambda _, sigs, commits, action: sigs,
            autospec=True,
        )
    )

  def test_osv_id_filter(self):
    test_vul_list = [
        vulnerability.Vulnerability(copy.deepcopy(self._test_osv_vuln))
    ]
    vulnerability_manager.OsvIdFilter(['ASB-A-9876']).filter(test_vul_list)
    self.assertLen(test_vul_list, 1)
    vulnerability_manager.OsvIdFilter(['ASB-A-1234']).filter(test_vul_list)
    self.assertEmpty(test_vul_list)

  def test_osv_id_prefix_filter(self):
    test_vul_list = [
        vulnerability.Vulnerability(copy.deepcopy(self._test_osv_vuln))
    ]
    vulnerability_manager.OsvIdDeniedPrefixFilter(['foo-', 'bar-']).filter(
        test_vul_list
    )
    self.assertLen(test_vul_list, 1)
    vulnerability_manager.OsvIdDeniedPrefixFilter(['foo-', 'ASB-A-']).filter(
        test_vul_list
    )
    self.assertEmpty(test_vul_list)

    # Test reverse matching
    test_vul_list = [
        vulnerability.Vulnerability(copy.deepcopy(self._test_osv_vuln)),
        vulnerability.Vulnerability(
            copy.deepcopy(self._test_qc_kernel_osv_vul)
        ),
    ]
    vulnerability_manager.OsvIdAllowedPrefixFilter(['foo-', 'ASB-A-']).filter(
        test_vul_list
    )
    self.assertLen(test_vul_list, 2)
    vulnerability_manager.OsvIdAllowedPrefixFilter(['foo-', 'ASB-A-12']).filter(
        test_vul_list
    )
    self.assertLen(test_vul_list, 1)
    self.assertEqual(test_vul_list[0].id, 'ASB-A-1234')

    test_vul_list = [
        vulnerability.Vulnerability(copy.deepcopy(self._test_osv_vuln)),
        vulnerability.Vulnerability(
            copy.deepcopy(self._test_qc_kernel_osv_vul)
        ),
    ]
    vulnerability_manager.OsvIdAllowedPrefixFilter(['foo-', 'bar-']).filter(
        test_vul_list
    )
    self.assertEmpty(test_vul_list)

  def test_cve_id_filter(self):
    self._test_osv_vuln['aliases'] = [
        'CVE-12345-12343', 'CVE-11111-11111', 'GSD-1234-12345'
    ]
    test_vul_list = [vulnerability.Vulnerability(self._test_osv_vuln)]
    vulnerability_manager.CveIdFilter(['CVE-00000-00000']).filter(test_vul_list)
    self.assertLen(test_vul_list, 1)
    vulnerability_manager.CveIdFilter(['CVE-12345-12343']).filter(test_vul_list)
    self.assertEmpty(test_vul_list)

  def test_android_severity_filter(self):
    test_vul_list = []
    vul1 = copy.deepcopy(self._test_vul)
    vul1.affected[0].ecosystem_specific['severity'] = 'Low'
    test_vul_list.append(vul1)
    vul2 = copy.deepcopy(self._test_vul)
    vul2.affected[0].ecosystem_specific['severity'] = 'Moderate'
    test_vul_list.append(vul2)
    vul3 = copy.deepcopy(self._test_vul)
    vul3.affected[0].ecosystem_specific['severity'] = 'High'
    test_vul_list.append(vul3)
    vul4 = copy.deepcopy(self._test_vul)
    vul4.affected[0].ecosystem_specific['severity'] = 'Critical'
    test_vul_list.append(vul4)
    vulnerability_manager.AndroidSeverityFilter(
        vulnerability_manager.AndroidSeverityLevel.HIGH).filter(test_vul_list)
    self.assertEmpty(vul1.affected)
    self.assertEmpty(vul2.affected)
    self.assertLen(vul3.affected, 1)
    self.assertLen(vul4.affected, 1)

  def test_android_severity_filter_do_nothing_for_unknown_ecosystem(self):
    self._test_vul.affected[0].ecosystem = 'Windows'
    test_vul_list = [self._test_vul]
    vulnerability_manager.AndroidSeverityFilter(
        vulnerability_manager.AndroidSeverityLevel.CRITICAL).filter(
            test_vul_list)
    self.assertLen(self._test_vul.affected, 1)

    del self._test_osv_vuln['affected'][0]['ecosystem_specific']
    test_vul_list = [vulnerability.Vulnerability(self._test_osv_vuln)]
    vulnerability_manager.AndroidSeverityFilter(
        vulnerability_manager.AndroidSeverityLevel.CRITICAL).filter(
            test_vul_list)
    self.assertLen(self._test_vul.affected, 1)

  def test_android_severity_filter_do_nothing_if_no_severity_exists(self):
    test_vul_list = [self._test_vul]
    vulnerability_manager.AndroidSeverityFilter(
        vulnerability_manager.AndroidSeverityLevel.CRITICAL).filter(
            test_vul_list)
    self.assertLen(self._test_vul.affected, 1)

  def test_android_severity_filter_do_nothing_for_unknown_severity(self):
    test_vul_list = [self._test_vul]
    self._test_vul.affected[0].ecosystem_specific[
        'severity'] = 'Catastrophic'
    with self.assertLogs(level=logging.ERROR) as logs:
      vulnerability_manager.AndroidSeverityFilter(
          vulnerability_manager.AndroidSeverityLevel.CRITICAL).filter(
              test_vul_list)
    self.assertLen(self._test_vul.affected, 1)
    self.assertIn(
        'Unknown severity found: Catastrophic (ID: %s). Skipping' %
        self._test_vul.id, logs.output[0])

  def test_android_spl_filter(self):
    test_vul_list = []
    vul1 = copy.deepcopy(self._test_vul)
    vul1.affected[0].ecosystem_specific['spl'] = '2020-01-10'
    test_vul_list.append(vul1)
    vul2 = copy.deepcopy(self._test_vul)
    vul2.affected[0].ecosystem_specific['spl'] = '2020-01-01'
    test_vul_list.append(vul2)
    vulnerability_manager.AndroidSplFilter('2020-01-05').filter(test_vul_list)
    self.assertEmpty(vul1.affected)
    self.assertLen(vul2.affected, 1)

  def test_android_spl_filter_do_nothing_for_unknown_ecosystem(self):
    self._test_vul.affected[0].ecosystem = 'Windows'
    test_vul_list = [self._test_vul]
    vulnerability_manager.AndroidSplFilter('2020-01-05').filter(test_vul_list)
    self.assertLen(self._test_vul.affected, 1)

    del self._test_osv_vuln['affected'][0]['ecosystem_specific']
    test_vul_list = [vulnerability.Vulnerability(self._test_osv_vuln)]
    vulnerability_manager.AndroidSplFilter('2020-01-05').filter(test_vul_list)
    self.assertLen(self._test_vul.affected, 1)

  def test_android_spl_filter_do_nothing_if_no_severity_exists(self):
    test_vul_list = [self._test_vul]
    vulnerability_manager.AndroidSplFilter('2020-01-05').filter(test_vul_list)
    self.assertLen(self._test_vul.affected, 1)

  def test_android_spl_filter_do_nothing_for_unknown_severity(self):
    test_vul_list = [self._test_vul]
    self._test_vul.affected[0].ecosystem_specific['spl'] = '12-25-2022'
    with self.assertLogs(level=logging.ERROR) as logs:
      vulnerability_manager.AndroidSplFilter('2020-01-05').filter(test_vul_list)
    self.assertLen(self._test_vul.affected, 1)
    self.assertIn(
        'Unknown SPL format: 12-25-2022 (ID: %s). Skipping' %
        self._test_vul.id, logs.output[0])

  def test_target_path_filter(self):
    osv_vul1 = copy.deepcopy(self._test_osv_vuln)
    osv_sign1 = _get_osv_signs(osv_vul1)[0]
    osv_sign1['target']['file'] = 'drivers/nvme/test123.c'
    sign = signature.Signature.from_osv_dict(osv_sign1)
    osv_vul1['affected'][0]['ecosystem_specific']['vanir_signatures'] = [sign]
    vul1 = vulnerability.Vulnerability(osv_vul1)

    osv_vul2 = copy.deepcopy(self._test_osv_vuln)
    osv_sign2 = _get_osv_signs(osv_vul2)[0]
    osv_sign2['id'] = 'ASB-A-9999-9999999999'
    osv_sign2['target']['file'] = 'drivers/nvme_ex/test123.c'
    sign = signature.Signature.from_osv_dict(osv_sign2)
    osv_vul2['affected'][0]['ecosystem_specific']['vanir_signatures'] = [sign]
    vul2 = vulnerability.Vulnerability(osv_vul2)

    test_vul_list = [vul1, vul2]
    test_pattern = re.compile('drivers/nvme/.*')
    vulnerability_manager.TargetPathFilter(test_pattern).filter(test_vul_list)
    self.assertEmpty(vul1.affected[0].vanir_signatures)
    self.assertLen(vul2.affected[0].vanir_signatures, 1)

  def test_architecture_filter(self):
    test_arches = ['powerpc', 'x86', 'arm', 'arm64', 'sparc']
    test_vul_dict = {}
    for arch in test_arches:
      osv_vul = copy.deepcopy(self._test_osv_vuln)
      osv_sign = _get_osv_signs(osv_vul)[0]
      osv_sign['id'] = 'ASB-A-1234-12345%s' % arch
      osv_sign['target']['file'] = 'arch/%s/mm/mem.c' % arch
      sign = signature.Signature.from_osv_dict(osv_sign)
      osv_vul['affected'][0]['ecosystem_specific']['vanir_signatures'] = [sign]
      test_vul_dict[arch] = vulnerability.Vulnerability(osv_vul)

    allowed_arches = ['arm', 'arm64']
    vulnerability_manager.ArchitectureFilter(allowed_arches).filter(
        test_vul_dict.values())

    for arch, vul in test_vul_dict.items():
      if arch in allowed_arches:
        self.assertLen(vul.affected[0].vanir_signatures, 1)
      else:
        self.assertEmpty(vul.affected[0].vanir_signatures)

  def test_affected_ecosystem_filter(self):
    test_vuln_list = [copy.deepcopy(self._test_vul) for _ in range(2)]
    test_vuln_list[0].affected[0].ecosystem = 'Android'
    test_vuln_list[1].affected[0].ecosystem = 'Chromium'
    vulnerability_manager.AffectedEcosystemFilter('Chromium').filter(
        test_vuln_list
    )
    self.assertLen(test_vuln_list, 1)
    self.assertLen(test_vuln_list[0].affected, 1)
    self.assertEqual(test_vuln_list[0].affected[0].ecosystem, 'Chromium')

  def test_affected_package_name_filter(self):
    test_vuln_list = [copy.deepcopy(self._test_vul) for _ in range(2)]
    test_vuln_list[1].affected[0].osv_package_name = 'platform/system/bt'
    vulnerability_manager.AffectedPackageNameFilter(
        'platform/system/.*').filter(test_vuln_list)
    self.assertLen(test_vuln_list, 1)
    self.assertLen(test_vuln_list[0].affected, 1)
    self.assertEqual(test_vuln_list[0].affected[0].osv_package_name,
                     'platform/system/bt')

  def test_affected_package_name_filter_with_regex_pattern(self):
    test_vuln_list = [copy.deepcopy(self._test_vul) for _ in range(2)]
    test_vuln_list[1].affected[0].osv_package_name = 'platform/system/bt'
    regex_pattern = re.compile('platform/system/.*')
    vulnerability_manager.AffectedPackageNameFilter(regex_pattern).filter(
        test_vuln_list
    )
    self.assertLen(test_vuln_list, 1)
    self.assertLen(test_vuln_list[0].affected, 1)
    self.assertEqual(
        test_vuln_list[0].affected[0].osv_package_name,
        'platform/system/bt',
    )

  def test_affected_package_name_filter_with_inverse_match(self):
    test_vuln_list = [copy.deepcopy(self._test_vul) for _ in range(2)]
    test_vuln_list[1].affected[0].osv_package_name = 'platform/system/bt'
    vulnerability_manager.AffectedPackageNameFilter(
        'platform/system/.*', inverse_match=True
    ).filter(test_vuln_list)
    self.assertLen(test_vuln_list, 1)
    self.assertLen(test_vuln_list[0].affected, 1)
    self.assertEqual(
        test_vuln_list[0].affected[0].osv_package_name,
        'test_package',
    )

  def test_vulnerability_manager_init(self):
    manager = vulnerability_manager.VulnerabilityManager([self._test_osv_vuln])
    self.assertLen(manager.vulnerabilities, 1)
    self.assertEqual(
        manager.vulnerabilities[0].to_osv_dict(), self._test_osv_vuln
    )

  def test_vulnerability_manager_init_fails_if_vuls_duplicate(self):
    with self.assertRaisesRegex(ValueError, 'Vulnerability .* already exists'):
      vulnerability_manager.VulnerabilityManager([self._test_osv_vuln] * 2)

  def test_vulnerability_manager_init_fails_if_signs_duplicate(self):
    additional_cve = copy.deepcopy(self._test_osv_vuln)
    additional_cve['id'] = 'ASB-A-2345'
    with self.assertRaisesRegex(ValueError, 'Signature.*already exists.*'):
      vulnerability_manager.VulnerabilityManager(
          [self._test_osv_vuln, additional_cve]
      )

  def test_vulnerability_manager_has_correct_signatures(self):
    vulnerabilities = []
    expected_signs = []
    expected_sign_id_to_osv_id = {}
    for i in range(5):
      vul = copy.deepcopy(self._test_osv_vuln)
      vul['id'] = 'ASB-A-TEST-%d' % i
      sign = _get_osv_signs(vul)[0]
      sign['id'] += str(i)
      vulnerabilities.append(vul)
      expected_signs.append(sign)
      expected_sign_id_to_osv_id[sign['id']] = vul['id']

    vul_with_no_sign = copy.deepcopy(self._test_osv_vuln)
    del vul_with_no_sign['affected'][0]['ecosystem_specific']
    vulnerabilities.append(vul_with_no_sign)

    manager = vulnerability_manager.VulnerabilityManager(vulnerabilities)

    signatures = manager.signatures
    self.assertLen(signatures, 5)
    self.assertCountEqual(expected_signs,
                          [sign.to_osv_dict() for sign in signatures])
    for sign in signatures:
      sign_id = sign.signature_id
      self.assertEqual(expected_sign_id_to_osv_id[sign_id],
                       manager.sign_id_to_osv_id(sign_id))

  def test_vulnerability_manager_returns_correct_cve_ids(self):
    self._test_osv_vuln['aliases'] = [
        'CVE-12345-12343', 'CVE-11111-11111', 'GSD-1234-12345'
    ]
    manager = vulnerability_manager.VulnerabilityManager([self._test_osv_vuln])
    self.assertSetEqual(
        {'CVE-11111-11111', 'CVE-12345-12343'},
        set(manager.sign_id_to_cve_ids(self._test_osv_sign['id']))
    )

  def test_vulnerability_manager_returns_correct_severities(self):
    vul1 = copy.deepcopy(self._test_osv_vuln)
    vul1['affected'][0]['ecosystem_specific']['severity'] = 'Low'
    # Add a second affected with different severity.
    vul1['affected'].append(copy.deepcopy(vul1['affected'][0]))
    vul1['affected'][1]['ecosystem_specific']['vanir_signatures'][0][
        'id'
    ] = 'ASB-A-1235-b4e3e62e8c'
    vul1['affected'][1]['ecosystem_specific']['severity'] = 'High'
    manager = vulnerability_manager.VulnerabilityManager([vul1])
    self.assertEqual(
        {'Low', 'High'}, manager.get_osv_severities(self._test_vul.id)
    )

  def test_vulnerability_manager_overwrites_duplicates(self):
    earlier_vul = copy.deepcopy(self._test_osv_vuln)
    earlier_vul['modified'] = '1830-11-11T21:26:24Z'
    earlier_vul['summary'] = 'This is an extremely old vulnerability.'
    later_vul = copy.deepcopy(self._test_osv_vuln)
    later_vul['modified'] = '2030-11-11T21:26:24Z'
    later_vul['summary'] = 'This is a very dangerious vulnerability.'
    manager = vulnerability_manager.VulnerabilityManager(
        [earlier_vul, later_vul, self._test_osv_vuln],
        overwrite_older_duplicate=True)
    self.assertEqual([later_vul], json.loads(manager.to_json()))

  def test_vulnerability_manager_fails_with_duplicates(self):
    earlier_vul = copy.deepcopy(self._test_osv_vuln)
    earlier_vul['modified'] = '1830-11-11T21:26:24Z'
    earlier_vul['summary'] = 'This is an extremely old vulnerability.'
    with self.assertRaisesRegex(ValueError, 'Vulnerability .* already exists.'):
      vulnerability_manager.VulnerabilityManager(
          [earlier_vul, self._test_osv_vuln],
          overwrite_older_duplicate=False)

  def test_vulnerability_manager_generate_signatures(self):
    manager = vulnerability_manager.VulnerabilityManager(
        [self._test_osv_vuln_without_signatures],
    )
    manager.generate_signatures()
    self.assertLen(manager.vulnerabilities, 1)
    self.assertEqual(
        manager.vulnerabilities[0].to_osv_dict(),
        self._test_osv_vuln
    )

  def test_vulnerability_manager_generate_signature_fails_when_sig_exists(self):
    manager = vulnerability_manager.VulnerabilityManager([self._test_osv_vuln])
    with self.assertRaisesRegex(ValueError, '.*already exists.*'):
      manager.generate_signatures()
    self.assertLen(manager.vulnerabilities, 1)
    self.assertEqual(
        manager.vulnerabilities[0].to_osv_dict(),
        self._test_osv_vuln
    )

  def test_vulnerability_manager_generate_signatures_with_deprecated_sigs(
      self
  ):
    manager = vulnerability_manager.VulnerabilityManager(
        [self._test_osv_vuln_without_signatures],
    )
    manager.generate_signatures(
        deprecated_signatures=[self._test_osv_sign['id']],
    )
    self.assertLen(manager.vulnerabilities, 1)
    expected_osv_vuln = copy.deepcopy(self._test_osv_vuln)
    _get_osv_signs(expected_osv_vuln)[0]['deprecated'] = True
    self.assertEqual(
        manager.vulnerabilities[0].to_osv_dict(), expected_osv_vuln
    )

  def test_vulnerability_manager_generate_signatures_with_deprecated_vulns(
      self
  ):
    manager = vulnerability_manager.VulnerabilityManager(
        [self._test_osv_vuln_without_signatures],
    )
    manager.generate_signatures(
        deprecated_vulns=[self._test_osv_vuln['id']],
    )
    self.assertLen(manager.vulnerabilities, 1)
    expected_osv_vuln = copy.deepcopy(self._test_osv_vuln)
    _get_osv_signs(expected_osv_vuln)[0]['deprecated'] = True
    self.assertEqual(
        manager.vulnerabilities[0].to_osv_dict(), expected_osv_vuln
    )

  def test_vulnerability_manager_generate_signatures_with_deprecated_patch(
      self
  ):
    manager = vulnerability_manager.VulnerabilityManager(
        [self._test_osv_vuln_without_signatures],
    )
    self._mock_commit.get_url.return_value = self._test_osv_sign['source']
    manager.generate_signatures(
        deprecated_patch_urls=[self._test_osv_sign['source']],
    )
    self.assertLen(manager.vulnerabilities, 1)
    expected_osv_vuln = copy.deepcopy(self._test_osv_vuln)
    _get_osv_signs(expected_osv_vuln)[0]['deprecated'] = True
    self.assertEqual(
        manager.vulnerabilities[0].to_osv_dict(), expected_osv_vuln
    )

  def test_vulnerability_manager_generate_signatures_with_exact_match_sigs(
      self
  ):
    manager = vulnerability_manager.VulnerabilityManager(
        [self._test_osv_vuln_without_signatures],
    )
    manager.generate_signatures(
        exact_match_only_signatures=[self._test_osv_sign['id']],
    )
    self.assertLen(manager.vulnerabilities, 1)
    expected_osv_vuln = copy.deepcopy(self._test_osv_vuln)
    _get_osv_signs(expected_osv_vuln)[0]['exact_target_file_match_only'] = True
    self.assertEqual(
        manager.vulnerabilities[0].to_osv_dict(), expected_osv_vuln
    )

  def test_vulnerability_manager_generate_signatures_with_exact_match_patch(
      self
  ):
    manager = vulnerability_manager.VulnerabilityManager(
        [self._test_osv_vuln_without_signatures],
    )
    self._mock_commit.get_url.return_value = self._test_osv_sign['source']
    manager.generate_signatures(
        exact_match_only_patch_urls=[self._test_osv_sign['source']],
    )
    self.assertLen(manager.vulnerabilities, 1)
    expected_osv_vuln = copy.deepcopy(self._test_osv_vuln)
    _get_osv_signs(expected_osv_vuln)[0]['exact_target_file_match_only'] = True
    self.assertEqual(
        manager.vulnerabilities[0].to_osv_dict(), expected_osv_vuln
    )

  def test_vulnerability_manager_generate_signatures_skips_failed_extraction(
      self
  ):
    manager = vulnerability_manager.VulnerabilityManager(
        [self._test_osv_vuln_without_signatures],
    )
    self._mock_extract.side_effect = ValueError('mock error.')
    with self.assertLogs(level=logging.ERROR) as logs:
      manager.generate_signatures()
      self.assertLen(manager.vulnerabilities, 1)
      self.assertEmpty(manager.vulnerabilities[0].affected[0].vanir_signatures)
      self.assertIn('Code extraction failed for', ''.join(logs.output))

  def test_vulnerability_manager_generate_signatures_refine_affected_no_false_positive(
      self
  ):
    manager = vulnerability_manager.VulnerabilityManager(
        [self._test_osv_vuln_without_signatures],
    )
    manager.generate_signatures()
    self.assertLen(manager.vulnerabilities[0].affected[0].vanir_signatures, 1)

  def test_vulnerability_manager_generate_signatures_refine_affected_false_positive_filtered(
      self
  ):
    manager = vulnerability_manager.VulnerabilityManager(
        [self._test_osv_vuln_without_signatures],
    )
    # simulate refinement removing all signatures during first pass
    self._mock_refiner.side_effect = lambda _, sigs, commits, action: (
        [] if isinstance(action, refiner.RemoveBadSignature) else sigs
    )
    manager.generate_signatures()
    self.assertEmpty(manager.vulnerabilities[0].affected[0].vanir_signatures)

  def test_vulnerability_manager_generate_signatures_cross_version_refine(self):
    commit1 = mock.create_autospec(code_extractor_base.Commit, instance=True)
    commit2 = mock.create_autospec(code_extractor_base.Commit, instance=True)
    commit1.get_url.return_value = 'commit_from_ver1'
    commit2.get_url.return_value = 'commit_from_ver2'
    self._mock_extract.side_effect = [([commit1], []), ([commit2], [])]

    signature2 = dataclasses.replace(
        self._test_signature, signature_id='ASB-A-1234-signaturehash2',
    )
    self._mock_sign_generator.side_effect = [
        [self._test_signature],
        [signature2],
    ]

    # simulate signature2 matching against patched files in commit1 in 2nd pass
    # this should cause refiner to mark signature2 as specific to version 2 only
    def mock_matcher(_, signatures, commits, action):
      if (
          signatures[0].signature_id == signature2.signature_id
          and commits[0].get_url() == 'commit_from_ver1'
          and isinstance(action, refiner.MarkAsSpecificToVersions)
      ):
        vers = action._versions
        return [
            (
                sig if sig != signature2
                else dataclasses.replace(sig, match_only_versions=vers)
            )
            for sig in signatures
        ]
      return signatures
    self._mock_refiner.side_effect = mock_matcher

    vuln_with_2_versions = {
        'id': 'ASB-A-1234',
        'modified': '2000-11-11T21:26:24Z',
        'affected': [
            {
                'package': {'name': 'test_package', 'ecosystem': 'Android'},
                'versions': ['1'],
            },
            {
                'package': {'name': 'test_package', 'ecosystem': 'Android'},
                'versions': ['2'],
            },
        ]
    }
    manager = vulnerability_manager.VulnerabilityManager([vuln_with_2_versions])
    manager.generate_signatures()

    # refiner should be called 8 times: once per commit, once for each affected
    # version with all commits, and once for each version against all versions
    # (2x2 = 4)
    self.assertLen(self._mock_refiner.call_args_list, 8)
    self.assertLen(manager.vulnerabilities[0].affected[0].vanir_signatures, 1)
    self.assertLen(manager.vulnerabilities[0].affected[1].vanir_signatures, 1)
    first_sig = manager.vulnerabilities[0].affected[0].vanir_signatures[0]
    second_sig = manager.vulnerabilities[0].affected[1].vanir_signatures[0]
    self.assertEqual(first_sig.signature_id, self._test_signature.signature_id)
    self.assertIsNone(first_sig.match_only_versions)
    self.assertEqual(second_sig.signature_id, signature2.signature_id)
    # signature2 should be marked as version-specific
    self.assertEqual(set(second_sig.match_only_versions), {'2'})

  def test_vulnerability_manager_generate_signatures_cross_version_refine_different_packages(
      self
  ):
    commit1 = mock.create_autospec(code_extractor_base.Commit, instance=True)
    commit2 = mock.create_autospec(code_extractor_base.Commit, instance=True)
    commit1.get_url.return_value = 'commit_from_ver1'
    commit2.get_url.return_value = 'commit_from_ver2'
    self._mock_extract.side_effect = [([commit1], []), ([commit2], [])]

    signature2 = dataclasses.replace(
        self._test_signature, signature_id='ASB-A-1234-signaturehash2',
    )
    self._mock_sign_generator.side_effect = [
        [self._test_signature],
        [signature2],
    ]

    # simulate signature2 matching against patched files in commit1 in 2nd pass
    def mock_matcher(_, signatures, commits, action):
      if (
          signatures[0].signature_id == signature2.signature_id
          and commits[0].get_url() == 'commit_from_ver1'
          and isinstance(action, refiner.MarkAsSpecificToVersions)
      ):
        vers = action._versions
        return [
            (
                sig if sig != signature2
                else dataclasses.replace(sig, match_only_versions=vers)
            )
            for sig in signatures
        ]
      return signatures
    self._mock_refiner.side_effect = mock_matcher

    vuln_with_2_versions = {
        'id': 'ASB-A-1234',
        'modified': '2000-11-11T21:26:24Z',
        'affected': [
            {
                'package': {'name': 'test_package1', 'ecosystem': 'Android'},
                'versions': ['1'],
            },
            {
                'package': {'name': 'test_package2', 'ecosystem': 'Android'},
                'versions': ['2'],
            },
        ]
    }
    manager = vulnerability_manager.VulnerabilityManager([vuln_with_2_versions])
    manager.generate_signatures()

    # refiner should be called 6 times: once for each commit, once for each
    # affected version, and once for each version against all versions of same
    # package (2x1 = 2)
    self.assertLen(self._mock_refiner.call_args_list, 6)
    self.assertLen(manager.vulnerabilities[0].affected[0].vanir_signatures, 1)
    self.assertLen(manager.vulnerabilities[0].affected[1].vanir_signatures, 1)
    first_sig = manager.vulnerabilities[0].affected[0].vanir_signatures[0]
    second_sig = manager.vulnerabilities[0].affected[1].vanir_signatures[0]
    self.assertEqual(first_sig.signature_id, self._test_signature.signature_id)
    self.assertIsNone(first_sig.match_only_versions)
    self.assertEqual(second_sig.signature_id, signature2.signature_id)
    self.assertIsNone(second_sig.match_only_versions)

  def test_vulnerability_manager_generate_signatures_cross_version_refine_with_missing_versions(
      self
  ):
    commit1 = mock.create_autospec(code_extractor_base.Commit, instance=True)
    commit1.get_url.return_value = 'commit_from_ver1'
    tip_commit = mock.create_autospec(code_extractor_base.Commit, instance=True)
    tip_commit.get_url.return_value = 'commit_from_tip'
    self._mock_extract.return_value = ([commit1], [])
    self._mock_sign_generator.side_effect = [[self._test_signature]]

    # simulate _test_signature matching against patched files in tip only
    def mock_refiner(_, sigs, commits, action):
      if (
          list(sigs)[0].signature_id == self._test_signature.signature_id
          and commits[0].get_url() == tip_commit.get_url()
          and isinstance(action, refiner.MarkAsSpecificToVersions)
      ):
        vers = action._versions
        return {
            (
                sig if sig != self._test_signature
                else dataclasses.replace(sig, match_only_versions=vers)
            )
            for sig in sigs
        }
      return sigs
    self._mock_refiner.side_effect = mock_refiner

    self._mock_extract_tip_of_unaffected_versions.return_value = (
        [tip_commit], [],
    )

    manager = vulnerability_manager.VulnerabilityManager(
        [self._test_osv_vuln_without_signatures],
    )
    manager.generate_signatures()

    self.assertLen(manager.vulnerabilities[0].affected[0].vanir_signatures, 1)
    sig = list(manager.vulnerabilities[0].affected[0].vanir_signatures)[0]
    self.assertEqual(sig.signature_id, self._test_signature.signature_id)
    self.assertEqual(
        set(sig.match_only_versions),
        set(self._test_osv_vuln_without_signatures['affected'][0]['versions']),
    )

  def test_vulnerability_manager_generate_signatures_skip_empty_affected(self):
    self._mock_extract.return_value = ([], [self._fake_failed_commit_url])
    manager = vulnerability_manager.VulnerabilityManager(
        [self._test_osv_vuln_without_signatures],
    )
    manager.generate_signatures()
    self.assertEmpty(manager.vulnerabilities[0].affected[0].vanir_signatures)
    self._mock_sign_generator.assert_not_called()
    self._mock_refiner.assert_not_called()

  def test_vulnerability_manager_with_filters(self):
    filters = []
    filters.append(vulnerability_manager.OsvIdFilter(['ASB-A-1234-1']))
    filters.append(vulnerability_manager.CveIdFilter(['CVE-9999-9999']))
    filters.append(
        vulnerability_manager.AndroidSeverityFilter(
            vulnerability_manager.AndroidSeverityLevel.CRITICAL))
    filters.append(vulnerability_manager.AndroidSplFilter('2020-12-05'))
    filters.append(
        vulnerability_manager.TargetPathFilter(re.compile('hello/world/.*')))
    filters.append(
        vulnerability_manager.ArchitectureFilter(
            allowed_arches=[vulnerability_manager.Architecture.RISCV]))

    vul1 = copy.deepcopy(self._test_osv_vuln)
    vul1['id'] = 'ASB-A-1234-1'
    _get_osv_signs(vul1)[0]['id'] = 'ASB-A-1234-b4e3-1'

    vul2 = copy.deepcopy(self._test_osv_vuln)
    vul2['id'] = 'ASB-A-1234-2'
    _get_osv_signs(vul2)[0]['id'] = 'ASB-A-1234-b4e3-2'
    vul2['aliases'] = ['CVE-9999-9999']

    vul3 = copy.deepcopy(self._test_osv_vuln)
    vul3['id'] = 'ASB-A-1234-3'
    _get_osv_signs(vul3)[0]['id'] = 'ASB-A-1234-b4e3-3'
    vul3['affected'][0]['ecosystem_specific']['severity'] = 'Low'

    vul4 = copy.deepcopy(self._test_osv_vuln)
    vul4['id'] = 'ASB-A-1234-4'
    _get_osv_signs(vul4)[0]['id'] = 'ASB-A-1234-b4e3-4'
    vul4['affected'][0]['ecosystem_specific']['spl'] = '2024-01-01'

    vul5 = copy.deepcopy(self._test_osv_vuln)
    vul5['id'] = 'ASB-A-1234-5'
    _get_osv_signs(vul5)[0]['id'] = 'ASB-A-1234-b4e3-5'
    _get_osv_signs(vul5)[0]['target']['file'] = 'hello/world/my/target/file.c'

    vul6 = copy.deepcopy(self._test_osv_vuln)
    vul6['id'] = 'ASB-A-1234-6'
    _get_osv_signs(vul6)[0]['id'] = 'ASB-A-1234-b4e3-6'
    _get_osv_signs(vul6)[0]['target']['file'] = 'arch/x86/mm/mm.c'

    vul7 = copy.deepcopy(self._test_osv_vuln)
    vul7['id'] = 'ASB-A-1234-7'
    _get_osv_signs(vul7)[0]['id'] = 'ASB-A-1234-b4e3-7'
    vul7['aliases'] = ['CVE-0000-0000']
    vul7['affected'][0]['ecosystem_specific']['severity'] = 'Critical'
    vul7['affected'][0]['ecosystem_specific']['spl'] = '2020-01-01'
    _get_osv_signs(vul7)[0]['target']['file'] = 'arch/riscv/mm/mm.c'

    test_vul_list = [vul1, vul2, vul3, vul4, vul5, vul6, vul7]

    manager = vulnerability_manager.VulnerabilityManager(
        test_vul_list, vulnerability_filters=filters)

    self.assertEqual(
        [vuln.id for vuln in manager.vulnerabilities],
        ['ASB-A-1234-5', 'ASB-A-1234-6', 'ASB-A-1234-7'])
    self.assertEqual(
        [sig.signature_id for sig in manager.signatures],
        ['ASB-A-1234-b4e3-7'])

    self.assertLen(manager.get_vulnerabilities(ignore_filters=True), 7)
    self.assertLen(manager.get_signatures(ignore_filters=True), 7)

  def test_get_affected_package_names(self):
    manager = vulnerability_manager.VulnerabilityManager(
        [self._test_osv_vuln, self._test_qc_kernel_osv_vul],
    )
    self.assertEqual(
        manager.affected_package_names,
        {'test_package', vulnerability.MetaPackage.ANDROID_KERNEL.value},
    )

  def test_vulnerability_manager_with_deprecated_signature_filter(self):
    vul1 = copy.deepcopy(self._test_osv_vuln)
    vul1['id'] = 'ASB-A-1234-1'
    _get_osv_signs(vul1)[0]['id'] = 'ASB-A-1234-b4e3-1'

    vul2 = copy.deepcopy(self._test_osv_vuln)
    vul2['id'] = 'ASB-A-1234-2'
    _get_osv_signs(vul2)[0]['id'] = 'ASB-A-1234-b4e3-2'
    _get_osv_signs(vul2)[0]['deprecated'] = True

    filters = [vulnerability_manager.DeprecatedSignatureFilter()]
    manager = vulnerability_manager.VulnerabilityManager(
        [vul1, vul2], vulnerability_filters=filters)

    self.assertEqual(
        [vuln.id for vuln in manager.vulnerabilities],
        ['ASB-A-1234-1', 'ASB-A-1234-2'])
    self.assertEqual(
        [sig.signature_id for sig in manager.signatures],
        ['ASB-A-1234-b4e3-1'])

  def test_generate_from_managers(self):
    vul1 = copy.deepcopy(self._test_osv_vuln)
    manager1 = vulnerability_manager.generate_from_json_string(
        json.dumps([vul1]),
    )
    vul2 = copy.deepcopy(self._test_osv_vuln)
    vul2['id'] = 'ASB-A-0999-TEST-99999'
    _get_osv_signs(vul2)[0]['id'] = 'ASB-A-0999-b4e3e62e8c'
    manager2 = vulnerability_manager.generate_from_json_string(
        json.dumps([vul2]),
    )
    uber_manager = vulnerability_manager.generate_from_managers(
        [manager1, manager2])
    self.assertSequenceEqual([vul2, vul1], json.loads(uber_manager.to_json()))
    vul3 = copy.deepcopy(self._test_osv_vuln)
    vul3['id'] = 'ASB-A-0011-TEST-00011'
    with self.assertRaisesRegex(ValueError, 'Signature.*already exists.*'):
      uber_manager.add_vulnerability(vulnerability.Vulnerability(vul3))

    # Test with empty managers.
    empty_manager = vulnerability_manager.generate_from_managers([])
    self.assertEmpty(empty_manager.vulnerabilities)
    self.assertEmpty(
        empty_manager.get_signatures_for_package('Android', 'test_package')
    )
    self.assertEmpty(empty_manager.get_signatures())
    empty_manager.add_vulnerability(self._test_vul)
    self.assertLen(empty_manager.vulnerabilities, 1)
    self.assertLen(
        empty_manager.get_signatures_for_package('Android', 'test_package'), 1
    )
    self.assertLen(empty_manager.get_signatures(), 1)

  def test_generate_from_managers_with_filters(self):
    vul1 = copy.deepcopy(self._test_osv_vuln)
    filter1 = vulnerability_manager.OsvIdFilter(['ASB-A-1234'])
    manager1 = vulnerability_manager.generate_from_json_string(
        json.dumps([vul1]),
        vulnerability_filters=[filter1])
    vul2 = copy.deepcopy(self._test_osv_vuln)
    vul2['id'] = 'ASB-A-9999-TEST-99999'
    _get_osv_signs(vul2)[0]['id'] = 'ASB-A-9999-b4e3e62e8c'
    manager2 = vulnerability_manager.generate_from_json_string(
        json.dumps([vul2]),
    )
    extra_filter = vulnerability_manager.OsvIdFilter(['ASB-A-9999-TEST-99999'])
    uber_manager = vulnerability_manager.generate_from_managers(
        [manager1, manager2], vulnerability_filters=[extra_filter])
    self.assertEmpty(uber_manager.vulnerabilities)

  def test_generate_from_json_string(self):
    test_vul_string = json.dumps([self._test_osv_vuln])
    manager = vulnerability_manager.generate_from_json_string(test_vul_string)
    self.assertEqual([self._test_osv_vuln], json.loads(manager.to_json()))

  def test_generate_from_file(self):
    test_vul_file = self.create_tempfile(
        content=json.dumps([self._test_osv_vuln]), mode='wt')
    manager = vulnerability_manager.generate_from_file(test_vul_file)
    self.assertLen(manager.vulnerabilities, 1)
    self.assertEqual(
        manager.vulnerabilities[0].id, self._test_osv_vuln['id']
    )

  def test_generate_from_file_fails_when_file_not_exist(self):
    with self.assertRaisesRegex(ValueError,
                                'Failed to find vulnerability file at'):
      vulnerability_manager.generate_from_file('/tmp/non/existing/file.json')

  def test_generate_from_osv_with_packages(self):
    self.enter_context(
        mock.patch.object(
            osv_client.OsvClient,
            'get_vulns_for_packages',
            return_value=[self._test_osv_vuln],
            autospec=True))
    manager = vulnerability_manager.generate_from_osv(
        'Android', vulnerability.MetaPackage.ANDROID_KERNEL)
    self.assertEqual([self._test_osv_vuln], json.loads(manager.to_json()))

  def test_generate_from_osv_all(self):
    self.enter_context(
        mock.patch.object(
            osv_client.OsvClient,
            'get_vulns_for_ecosystem',
            return_value=[self._test_osv_vuln],
            autospec=True))
    manager = vulnerability_manager.generate_from_osv('Android')
    self.assertEqual([self._test_osv_vuln], json.loads(manager.to_json()))


if __name__ == '__main__':
  absltest.main()