# 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 for the vulnerability module."""

import copy
import re

from osv import vulnerability_pb2
from vanir import signature
from vanir import vulnerability

from google.protobuf import timestamp_pb2
from absl.testing import absltest

_TEST_OSV_ID = 'ASB-A-1234'
_TEST_OSV_MODIFIED = '2024-01-01T00:00:00Z'
_TEST_ECOSYSTEM = 'Android'
_TEST_PACKAGE_NAME = 'platform/foo'
_TEST_OSV_AFFECTED = {
    'package': {
        'name': _TEST_PACKAGE_NAME,
        'ecosystem': _TEST_ECOSYSTEM,
    }
}
_TEST_OSV_VULN = {
    'id': _TEST_OSV_ID,
    'modified': _TEST_OSV_MODIFIED,
    'affected': [_TEST_OSV_AFFECTED],
}
_TEST_OSV_SIG = {
    'id': 'ASB-A-1234-sigid1',
    'signature_type': 'Function',
    'signature_version': 'v1',
    'source': 'sourceurl',
    'target': {'file': 'targetfile', 'function': 'func'},
    'deprecated': False,
    'digest': {'function_hash': '1234', 'length': 100},
}
_TEST_OSV_SIG2 = {
    'id': 'ASB-A-1234-sigid2',
    'withdrawn': '2023-10-27T10:30:00Z',
    'signature_type': 'Function',
    'signature_version': 'v1',
    'source': 'sourceurl',
    'target': {'file': 'targetfile', 'function': 'func'},
    'deprecated': False,
    'digest': {'function_hash': '1234', 'length': 100},
}


class VulnerabilityTest(absltest.TestCase):

  def test_osv_vuln_import(self):
    vuln = vulnerability.Vulnerability(_TEST_OSV_VULN)
    self.assertEqual(vuln.id, _TEST_OSV_ID)
    self.assertEqual(vuln.modified, _TEST_OSV_MODIFIED)
    affected = list(vuln.affected)
    self.assertLen(affected, 1)
    self.assertEqual(affected[0].ecosystem, _TEST_ECOSYSTEM)
    self.assertEqual(affected[0].osv_package_name, _TEST_PACKAGE_NAME)

  def test_setting_affected_packages(self):
    vuln = vulnerability.Vulnerability(_TEST_OSV_VULN)
    vuln.affected = [
        vulnerability.AffectedEntry(
            {'package': {'name': 'test1', 'ecosystem': 'bar'}}
        ),
        vulnerability.AffectedEntry(
            {'package': {'name': 'test2', 'ecosystem': 'baz'}}
        ),
    ]
    self.assertLen(vuln.affected, 2)
    self.assertEqual(vuln.affected[0].osv_package_name, 'test1')
    self.assertEqual(vuln.affected[1].osv_package_name, 'test2')

  def test_osv_vuln_import_with_no_id(self):
    with self.assertRaisesRegex(ValueError, '.*Missing.*'):
      vulnerability.Vulnerability(
          {'modified': _TEST_OSV_MODIFIED, 'affected': [_TEST_OSV_AFFECTED]}
      )
    with self.assertRaisesRegex(ValueError, '.*Missing.*'):
      vulnerability.Vulnerability(
          {
              'id': '',
              'modified': _TEST_OSV_MODIFIED,
              'affected': [_TEST_OSV_AFFECTED]
          }
      )

  def test_osv_vuln_import_with_no_modified(self):
    with self.assertRaisesRegex(ValueError, '.*modified.*'):
      vulnerability.Vulnerability(
          {'id': _TEST_OSV_ID, 'affected': [_TEST_OSV_AFFECTED]}
      )
    with self.assertRaisesRegex(ValueError, '.*modified.*'):
      vulnerability.Vulnerability(
          {
              'id': _TEST_OSV_ID,
              'modified': None,
              'affected': [_TEST_OSV_AFFECTED],
          }
      )

  def test_osv_vuln_import_with_invalid_affected(self):
    with self.assertRaisesRegex(ValueError, '.*affected*'):
      vulnerability.Vulnerability(
          {'id': _TEST_OSV_ID, 'modified': _TEST_OSV_MODIFIED}
      )
    with self.assertRaisesRegex(ValueError, '.*affected.*'):
      vulnerability.Vulnerability(
          {'id': _TEST_OSV_ID, 'modified': _TEST_OSV_MODIFIED, 'affected': []}
      )

  def test_osv_vuln_import_with_both_old_and_new_signatures(self):
    osv_vuln = copy.deepcopy(_TEST_OSV_VULN)
    osv_vuln['affected'][0]['database_specific'] = {
        'vanir_signatures': [_TEST_OSV_SIG]
    }
    osv_vuln['affected'][0]['ecosystem_specific'] = {
        'vanir_signatures': [_TEST_OSV_SIG]
    }
    with self.assertRaisesRegex(ValueError, '.*signatures.*'):
      vulnerability.Vulnerability(osv_vuln)

  def test_osv_vuln_import_with_no_package_name(self):
    vuln_with_no_package_name = copy.deepcopy(_TEST_OSV_VULN)
    vuln_with_no_package_name['affected'][0]['package'].pop('name')
    with self.assertRaisesRegex(ValueError, '.*package.*'):
      vulnerability.Vulnerability(vuln_with_no_package_name)

  def test_osv_vuln_import_with_no_ecosystem(self):
    vuln_with_no_ecosystem = copy.deepcopy(_TEST_OSV_VULN)
    vuln_with_no_ecosystem['affected'][0]['package'].pop('ecosystem')
    with self.assertRaisesRegex(ValueError, '.*package.*'):
      vulnerability.Vulnerability(vuln_with_no_ecosystem)

  def test_osv_vuln_import_with_no_package_info(self):
    vuln_with_no_package = copy.deepcopy(_TEST_OSV_VULN)
    vuln_with_no_package['affected'][0].pop('package')
    vuln = vulnerability.Vulnerability(vuln_with_no_package)
    self.assertEqual(vuln.affected[0].ecosystem, '')
    self.assertEqual(vuln.affected[0].osv_package_name, '')
    self.assertEqual(vuln.affected[0].package_name, '')

  def test_to_osv_dict(self):
    vuln = vulnerability.Vulnerability(_TEST_OSV_VULN)
    self.assertEqual(vuln.to_osv_dict(), _TEST_OSV_VULN)

  def test_to_osv_dict_after_adding_affected_entry(self):
    vuln = vulnerability.Vulnerability(_TEST_OSV_VULN)
    second_affected = {
        'package': {'name': 'test2', 'ecosystem': 'baz'},
        'database_specific': {'vanir_signatures': [_TEST_OSV_SIG]},
    }
    vuln.affected.append(vulnerability.AffectedEntry(second_affected))
    vuln.summary = 'test summary'
    vuln.aliases = ['CVE-1234-1234', 'CVE-5678-5678']

    expected_osv_dict = {
        'id': _TEST_OSV_ID,
        'modified': _TEST_OSV_MODIFIED,
        'summary': 'test summary',
        'aliases': ['CVE-1234-1234', 'CVE-5678-5678'],
        'affected': [_TEST_OSV_AFFECTED, second_affected],
    }
    self.assertEqual(vuln.to_osv_dict(), expected_osv_dict)

  def test_to_osv_dict_saving_signature_to_new_location(self):
    test_legacy_osv_vuln = copy.deepcopy(_TEST_OSV_VULN)
    test_legacy_osv_vuln['affected'][0]['ecosystem_specific'] = {
        'vanir_signatures': [_TEST_OSV_SIG]
    }
    vuln = vulnerability.Vulnerability(test_legacy_osv_vuln)
    expected_osv_dict = copy.deepcopy(_TEST_OSV_VULN)
    expected_osv_dict['affected'][0]['database_specific'] = {
        'vanir_signatures': [_TEST_OSV_SIG]
    }
    self.assertEqual(vuln.to_osv_dict(), expected_osv_dict)

  def test_to_osv_dict_after_adding_signatures(self):
    vuln = vulnerability.Vulnerability(_TEST_OSV_VULN)
    vuln.affected[0].vanir_signatures = [_TEST_OSV_SIG]
    expected_osv_dict = copy.deepcopy(_TEST_OSV_VULN)
    expected_osv_dict['affected'][0]['database_specific'] = {
        'vanir_signatures': [_TEST_OSV_SIG]
    }
    self.assertEqual(vuln.to_osv_dict(), expected_osv_dict)

  def test_to_osv_dict_after_adding_signatures_legacy(self):
    vuln = vulnerability.Vulnerability(
        _TEST_OSV_VULN, store_signatures_in_legacy_location=True,
    )
    vuln.affected[0].vanir_signatures = [_TEST_OSV_SIG]
    expected_osv_dict = copy.deepcopy(_TEST_OSV_VULN)
    expected_osv_dict['affected'][0]['ecosystem_specific'] = {
        'vanir_signatures': [_TEST_OSV_SIG]
    }
    self.assertEqual(vuln.to_osv_dict(), expected_osv_dict)

  def test_init_with_mix_of_signature_objects_and_dicts(self):
    osv_vuln = copy.deepcopy(_TEST_OSV_VULN)
    osv_vuln['affected'][0]['database_specific'] = {
        'vanir_signatures': [
            _TEST_OSV_SIG,
            signature.Signature.from_osv_dict(_TEST_OSV_SIG2),
        ]
    }
    vuln = vulnerability.Vulnerability(osv_vuln)
    self.assertLen(vuln.affected[0].vanir_signatures, 2)
    self.assertIsInstance(
        vuln.affected[0].vanir_signatures[0],
        signature.Signature,
    )
    self.assertIsInstance(
        vuln.affected[0].vanir_signatures[1],
        signature.Signature,
    )
    self.assertEqual(
        vuln.affected[0].vanir_signatures[0].signature_id,
        _TEST_OSV_SIG['id']
    )
    self.assertEqual(
        vuln.affected[0].vanir_signatures[1].signature_id,
        _TEST_OSV_SIG2['id']
    )

  def test_sort_vanir_signatures(self):
    affected = vulnerability.AffectedEntry({
        'package': {'name': 'test2', 'ecosystem': 'baz'},
        'ecosystem_specific': {
            'vanir_signatures': [_TEST_OSV_SIG2, _TEST_OSV_SIG],
        },
    })
    affected.sort_vanir_signatures()
    self.assertLen(affected.vanir_signatures, 2)
    self.assertEqual(
        affected.vanir_signatures[0].signature_id,
        _TEST_OSV_SIG['id']
    )
    self.assertEqual(
        affected.vanir_signatures[1].signature_id,
        _TEST_OSV_SIG2['id']
    )

  def test_meta_package_pattern(self):
    self.assertEqual(
        vulnerability.MetaPackage.ANDROID_KERNEL.package_pattern,
        re.compile(':linux_kernel:.*'),
    )
    self.assertEqual(
        vulnerability.MetaPackage.ANDROID_MODEM.package_pattern,
        re.compile(':modem:.*'),
    )
    self.assertEqual(
        vulnerability.MetaPackage.UNKNOWN.package_pattern,
        re.compile(':unknown:.*'),
    )

  def test_affected_package_normalize_package_name(self):
    affected1 = vulnerability.AffectedEntry({
        'package': {
            'name': 'just/foo/bar/package',
            'ecosystem': _TEST_ECOSYSTEM,
        }
    })
    self.assertEqual(affected1.osv_package_name, 'just/foo/bar/package')
    self.assertEqual(affected1.package_name, 'just/foo/bar/package')

    affected2 = vulnerability.AffectedEntry({
        'package': {
            'name': ':linux_kernel:foo_bar',
            'ecosystem': _TEST_ECOSYSTEM,
        }
    })
    self.assertEqual(affected2.osv_package_name, ':linux_kernel:foo_bar')
    self.assertEqual(affected2.package_name, ':linux_kernel:')

    affected3 = vulnerability.AffectedEntry({
        'package': {
            'name': ':modem:',
            'ecosystem': _TEST_ECOSYSTEM,
        }
    })
    self.assertEqual(affected3.package_name, ':modem:')

  def test_affected_get_versions(self):
    affected1 = vulnerability.AffectedEntry({
        'package': {'name': 'test2', 'ecosystem': 'baz'},
        'versions': ['1.0.0', '1.0.1'],
    })
    self.assertEqual(affected1.versions, ['1.0.0', '1.0.1'])

    affected2 = vulnerability.AffectedEntry({
        'package': {'name': 'test2', 'ecosystem': 'baz'},
    })
    self.assertEqual(affected2.versions, [])

  def test_vulnerability_to_proto(self):
    vuln = vulnerability.Vulnerability(_TEST_OSV_VULN)
    proto = vuln.to_proto()
    expected_modified = timestamp_pb2.Timestamp()
    expected_modified.FromJsonString(_TEST_OSV_MODIFIED)
    expected_proto = vulnerability_pb2.Vulnerability(
        id=_TEST_OSV_ID,
        modified=expected_modified,
        affected=[
            vulnerability_pb2.Affected(
                package=vulnerability_pb2.Package(
                    name=_TEST_PACKAGE_NAME,
                    ecosystem=_TEST_ECOSYSTEM,
                ),
            )
        ],
    )
    self.assertEqual(proto, expected_proto)

  def test_affected_to_proto(self):
    affected = vulnerability.AffectedEntry(_TEST_OSV_AFFECTED)
    proto = affected.to_proto()
    expected_proto = vulnerability_pb2.Affected(
        package=vulnerability_pb2.Package(
            name=_TEST_PACKAGE_NAME,
            ecosystem=_TEST_ECOSYSTEM,
        ),
    )
    self.assertEqual(proto, expected_proto)


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