# SPDX-FileCopyrightText: 2021-2024 Greenbone AG
#
# SPDX-License-Identifier: AGPL-3.0-or-later

from unittest import TestCase

from notus.scanner.models.packages.package import PackageAdvisory
from notus.scanner.models.packages.rpm import RPMPackage
from notus.scanner.models.vulnerability import Vulnerabilities, Vulnerability


class VulnerabilityTestCase(TestCase):
    def test_vulnerability(self):
        vul = Vulnerability()

        package = RPMPackage.from_full_name(
            "foo-1.2.3-4.x86_64",
        )
        if not package:
            self.fail("Cannot build package")
        fixed_package = RPMPackage.from_full_name(
            "foo-1.2.4-4.x86_64",
        )
        if not fixed_package:
            self.fail("Cannot build package")

        fixed = PackageAdvisory(
            fixed_package, "1.2.3.4.5", ">=", lambda _: False
        )

        vul.add(package, fixed)

        for key, val in vul.get().items():
            self.assertEqual(package, key)
            self.assertListEqual([fixed], val)

        fixed_package = RPMPackage.from_full_name(
            "foo-1.2.2-4.x86_64",
        )
        if not fixed_package:
            self.fail("Cannot build package")

        fixed = PackageAdvisory(
            fixed_package, "1.2.3.4.5", "<", lambda _: False
        )

        vul.add(package, fixed)

        self.assertEqual(len(vul.get()), 1)
        self.assertEqual(len(vul.get()[package]), 2)

        vul2 = Vulnerability()

        package = RPMPackage.from_full_name(
            "bar-1.2.3-4.x86_64",
        )
        if not package:
            self.fail("Cannot build package")

        fixed_package = RPMPackage.from_full_name(
            "bar-1.2.4-4.x86_64",
        )
        if not fixed_package:
            self.fail("Cannot build package")

        fixed = PackageAdvisory(
            fixed_package, "1.2.3.4.5", ">=", lambda _: False
        )

        vul2.add(package, fixed)
        vul.combine(vul2)

        self.assertEqual(len(vul.get()), 2)


class VulnerabilitiesTestCase(TestCase):
    def test_vulnerabilities(self):
        vuls = Vulnerabilities()
        vul = Vulnerability()

        package = RPMPackage.from_full_name(
            "foo-1.2.3-4.x86_64",
        )
        if not package:
            self.fail("Cannot build package")

        fixed_package = RPMPackage.from_full_name(
            "foo-1.2.4-4.x86_64",
        )
        if not fixed_package:
            self.fail("Cannot build package")

        fixed = PackageAdvisory(
            fixed_package, "1.2.3.4.5", ">=", lambda _: False
        )

        vul.add(package, fixed)
        vuls.add("1.2.3.4.5", vul)

        self.assertEqual(len(vuls), 1)
        self.assertTrue("1.2.3.4.5" in vuls.get())
        self.assertEqual(vuls.get()["1.2.3.4.5"], vul)

        vul = Vulnerability()

        package = RPMPackage.from_full_name(
            "bar-1.2.3-4.x86_64",
        )
        if not package:
            self.fail("Cannot build package")

        fixed_package = RPMPackage.from_full_name(
            "bar-1.2.4-4.x86_64",
        )
        if not fixed_package:
            self.fail("Cannot build package")

        fixed = PackageAdvisory(
            fixed_package, "1.2.3.4.6", ">=", lambda _: False
        )

        vul.add(package, fixed)
        vuls.add("1.2.3.4.6", vul)

        self.assertEqual(len(vuls), 2)
        self.assertTrue("1.2.3.4.6" in vuls.get())
        self.assertEqual(vuls.get()["1.2.3.4.6"], vul)
