# Copyright (C) 2021-2022 Greenbone Networks GmbH
#
# SPDX-License-Identifier: AGPL-3.0-or-later
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.

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)
