PyVulnerabilityLookup/tests/test_web.py

#!/usr/bin/env python3

import unittest
import uuid
import time
import os

from datetime import datetime, timezone, timedelta

from pyvulnerabilitylookup import PyVulnerabilityLookup

# NOTE:
# * to run the tests with a pre-configured admin key:
#   API_KEY="<APIKEY>" pytest tests/test_web.py
# * to run the test against a local instance (fallback to public):
#   INSTANCE_URL="http://0.0.0.0:10001" pytest tests/test_web.py


class TestPublic(unittest.TestCase):

    def setUp(self) -> None:
        self.admin_token = os.getenv("API_KEY", '').strip()
        instance_url = os.getenv("INSTANCE_URL", 'https://vulnerability.circl.lu').strip()
        self.public_test = instance_url == 'https://vulnerability.circl.lu'
        self.client = PyVulnerabilityLookup(root_url=instance_url, token=self.admin_token)

    def test_up(self) -> None:
        self.assertTrue(self.client.is_up)
        self.assertTrue(self.client.redis_up())

    def test_get_vulnerability(self) -> None:
        while True:
            if vuln := self.client.get_vulnerability('PYSEC-2024-4'):
                self.assertEqual(vuln['id'], 'PYSEC-2024-4')
                break
            print('waiting for pysec to be imported')
            time.sleep(1)

    def test_get_info(self) -> None:
        info = self.client.get_info()
        self.assertTrue(info['last_updates'])
        self.assertTrue(info['db_sizes'])

    def test_get_last(self) -> None:
        last = self.client.get_last()
        self.assertTrue(last)
        self.assertTrue(isinstance(last, list))
        last = self.client.get_last(number=1)
        self.assertTrue(isinstance(last, list))
        self.assertEqual(len(last), 1)
        last = self.client.get_last(source='pysec')
        for vuln in last:
            self.assertTrue(vuln['id'].startswith('PYSEC'))
        last = self.client.get_last(source='pysec', number=1)
        self.assertEqual(len(last), 1)
        self.assertTrue(last[-1]['id'].startswith('PYSEC'))

    # TODO: POST Vulnerability / Delete vulnerability
    # Test API

    def test_get_vendors(self) -> None:
        vendors = self.client.get_vendors()
        self.assertTrue(isinstance(vendors, list))

    def test_get_vendor_products(self) -> None:
        if not self.public_test:
            return None
        products = self.client.get_vendor_products('misp')
        self.assertTrue(isinstance(products, list))
        self.assertTrue('misp' in products)

    def test_get_vendor_product_vulnerabilities(self) -> None:
        if not self.public_test:
            return None
        vulns = self.client.get_vendor_product_vulnerabilities('misp', 'misp')
        self.assertTrue(isinstance(vulns, dict))
        self.assertTrue('cvelistv5' in vulns)

    # Test comments

    def test_get_comments_public(self) -> None:
        if not self.public_test:
            return None
        comments = self.client.get_comments()
        self.assertTrue('metadata' in comments)
        self.assertTrue('data' in comments)
        self.assertTrue(len(comments['data']) > 0)

        comments = self.client.get_comments(uuid='a309d024-2714-4a81-a425-60f83f6d5740')
        self.assertTrue(len(comments['data']) == 1)
        self.assertEqual(comments['data'][0]['uuid'], 'a309d024-2714-4a81-a425-60f83f6d5740')

        comments = self.client.get_comments(vuln_id='CVE-2024-20401')
        self.assertTrue(len(comments['data']) >= 1)
        for comment in comments['data']:
            self.assertEqual(comment['vulnerability'], 'CVE-2024-20401')

        comments = self.client.get_comments(author='adulau')
        self.assertTrue(len(comments['data']) >= 1)
        for comment in comments['data']:
            self.assertEqual(comment['author']['login'], 'adulau')

        comments = self.client.get_comments(uuid='a309d024-2714-4a81-a425-60f83f6d5740',
                                            vuln_id='CVE-2024-20401',
                                            author='adulau')
        self.assertTrue(len(comments['data']) == 1)
        self.assertEqual(comments['data'][0]['uuid'], 'a309d024-2714-4a81-a425-60f83f6d5740')
        self.assertEqual(comments['data'][0]['vulnerability'], 'CVE-2024-20401')
        self.assertEqual(comments['data'][0]['author']['login'], 'adulau')

    def test_comments_local(self) -> None:
        if not self.admin_token:
            # this test is only working if the admin token is set
            return None
        # Makes sure the userkey is set to the right one
        self.client.set_apikey(self.admin_token)
        comments = self.client.get_comments()
        self.assertTrue("metadata" in comments)
        self.assertTrue("data" in comments)
        self.assertTrue(len(comments['data']) == 0)
        self.assertEqual(comments['metadata']['count'], 0)

        comment = {
            "uuid": "a309d024-2714-4a81-a425-60f83f6d5740",
            "title": "Comment",
            "description": "Comment",
            "description_format": "markdown",
            "vulnerability": "CVE-2024-20401",
            "related_vulnerabilities": ["ghsa-4rcj-fmjg-q9fv"],
        }
        comments = self.client.create_comment(comment=comment)
        self.assertTrue(len(comments["data"]) == 1)
        self.assertEqual(
            comments["data"][0]["uuid"], "a309d024-2714-4a81-a425-60f83f6d5740"
        )

        comment = self.client.get_comment("a309d024-2714-4a81-a425-60f83f6d5740")
        self.assertEqual(
            comment["uuid"], "a309d024-2714-4a81-a425-60f83f6d5740"
        )
        comments = self.client.get_comments(uuid="a309d024-2714-4a81-a425-60f83f6d5740")
        self.assertTrue(len(comments["data"]) == 1)
        self.assertEqual(
            comments["data"][0]["uuid"], "a309d024-2714-4a81-a425-60f83f6d5740"
        )

        comments = self.client.get_comments(vuln_id="CVE-2024-20401")
        self.assertTrue(len(comments["data"]) >= 1)
        for comment in comments["data"]:
            self.assertEqual(comment["vulnerability"], "CVE-2024-20401")

        # comments = self.client.get_comments(author='admin')
        # self.assertTrue(len(comments['data']) >= 1)
        # for comment in comments['data']:
        #     self.assertEqual(comment['author']['login'], 'admin') # type: ignore[call-overload]

        comments = self.client.get_comments(
            uuid="a309d024-2714-4a81-a425-60f83f6d5740", vuln_id="CVE-2024-20401"
        )
        self.assertTrue(len(comments["data"]) == 1)
        self.assertEqual(
            comments["data"][0]["uuid"], "a309d024-2714-4a81-a425-60f83f6d5740"
        )
        self.assertEqual(comments["data"][0]["vulnerability"], "CVE-2024-20401")
        # self.assertEqual(comments['data'][0]['author']['login'], 'admin')

        status_code = self.client.delete_comment("a309d024-2714-4a81-a425-60f83f6d5740")
        self.assertTrue(status_code == 204)
        comments = self.client.get_comments(uuid="a309d024-2714-4a81-a425-60f83f6d5740")
        self.assertTrue(len(comments["data"]) == 0)

    # Test bundles

    def test_get_bundles_public(self) -> None:
        if not self.public_test:
            return None
        bundles = self.client.get_bundles()
        self.assertTrue('metadata' in bundles)
        self.assertTrue('data' in bundles)
        self.assertTrue(len(bundles['data']) > 0)

        bundles = self.client.get_bundles(uuid='a23cbcad-e890-4df8-8736-9332ed4c3d47')
        self.assertTrue(len(bundles['data']) == 1)
        self.assertEqual(bundles['data'][0]['uuid'], 'a23cbcad-e890-4df8-8736-9332ed4c3d47')

        bundles = self.client.get_bundles(vuln_id='CVE-2024-39573')
        self.assertTrue(len(bundles['data']) >= 1)
        for bundle in bundles['data']:
            self.assertTrue('CVE-2024-39573' in bundle['related_vulnerabilities'])

        bundles = self.client.get_bundles(author='adulau')
        self.assertTrue(len(bundles['data']) >= 1)
        for bundle in bundles['data']:
            self.assertEqual(bundle['author']['login'], 'adulau')

        bundles = self.client.get_bundles(uuid='a23cbcad-e890-4df8-8736-9332ed4c3d47',
                                          vuln_id='CVE-2024-39573',
                                          author='adulau')
        self.assertTrue(len(bundles['data']) == 1)
        self.assertEqual(bundles['data'][0]['uuid'], 'a23cbcad-e890-4df8-8736-9332ed4c3d47')
        self.assertTrue('CVE-2024-39573' in bundles['data'][0]['related_vulnerabilities'])
        self.assertEqual(bundles['data'][0]['author']['login'], 'adulau')

    def test_bundles_local(self) -> None:
        if not self.admin_token:
            # this test is only working if the admin token is set
            return None
        # Makes sure the userkey is set to the right one
        self.client.set_apikey(self.admin_token)
        bundles = self.client.get_bundles()
        self.assertTrue("metadata" in bundles)
        self.assertTrue("data" in bundles)
        self.assertTrue(len(bundles['data']) == 0)
        self.assertEqual(bundles['metadata']['count'], 0)

        bundle = {
            "uuid": "a23cbcad-e890-4df8-8736-9332ed4c3d47",
            "name": "Bundle",
            "description": "Bundle",
            "description_format": "markdown",
            "related_vulnerabilities": ["ghsa-4rcj-fmjg-q9fv", "CVE-2024-39573"],
        }
        bundles = self.client.create_bundle(bundle=bundle)
        self.assertTrue(len(bundles["data"]) == 1)
        self.assertEqual(
            bundles["data"][0]["uuid"], "a23cbcad-e890-4df8-8736-9332ed4c3d47"
        )

        bundles = self.client.get_bundles(uuid="a23cbcad-e890-4df8-8736-9332ed4c3d47")
        self.assertTrue(len(bundles["data"]) == 1)
        self.assertEqual(
            bundles["data"][0]["uuid"], "a23cbcad-e890-4df8-8736-9332ed4c3d47"
        )

        bundles = self.client.get_bundles(vuln_id="CVE-2024-39573")
        self.assertTrue(len(bundles["data"]) >= 1)
        for bundle in bundles["data"]:
            self.assertTrue("CVE-2024-39573" in bundle["related_vulnerabilities"])

        # bundles = self.client.get_bundles(author='admin')
        # self.assertTrue(len(bundles['data']) >= 1)
        # for bundle in bundles['data']:
        #     self.assertEqual(bundle['author']['login'], 'admin') # type: ignore[call-overload]

        bundles = self.client.get_bundles(
            uuid="a23cbcad-e890-4df8-8736-9332ed4c3d47", vuln_id="CVE-2024-39573"
        )
        self.assertTrue(len(bundles["data"]) == 1)
        self.assertEqual(
            bundles["data"][0]["uuid"], "a23cbcad-e890-4df8-8736-9332ed4c3d47"
        )
        self.assertTrue(
            "CVE-2024-39573" in bundles["data"][0]["related_vulnerabilities"]
        )
        # self.assertEqual(bundles['data'][0]['author']['login'], 'admin')

        status_code = self.client.delete_bundle("a23cbcad-e890-4df8-8736-9332ed4c3d47")
        self.assertTrue(status_code == 204)
        comments = self.client.get_comments(uuid="a23cbcad-e890-4df8-8736-9332ed4c3d47")
        self.assertTrue(len(comments["data"]) == 0)

    # Test User
    def test_users_info(self) -> None:
        if not self.admin_token:
            # this test is only working if the admin token is set
            return None
        # Makes sure the userkey is set to the right one
        self.client.set_apikey(self.admin_token)
        user_info = self.client.get_user_information()
        self.assertTrue(isinstance(user_info, dict))
        self.assertTrue('login' in user_info)
        self.assertTrue('apikey' in user_info)
        self.assertEqual(user_info['apikey'], self.admin_token)

    def test_list_users(self) -> None:
        if not self.admin_token:
            # this test is only working if the admin token is set
            return None
        # Makes sure the userkey is set to the right one
        self.client.set_apikey(self.admin_token)
        users = self.client.list_users()
        self.assertTrue(isinstance(users, dict))
        self.assertTrue(len(users['data']) > 0)
        got_admin = False
        for user in users['data']:
            self.assertTrue('login' in user)
            self.assertTrue('apikey' in user)
            if user['apikey'] == self.admin_token:
                self.assertTrue(user['is_admin'])
                got_admin = True
        self.assertTrue(got_admin)

    def test_create_user_comment(self) -> None:
        if self.public_test:
            # Do not run that test against the public instance, it would create users.
            return None
        instance_config = self.client.get_config_info()
        if not instance_config.get('registration'):
            return None

        user = self.client.create_user(name='test Name', login='alan',
                                       organisation='test Organization', email='test@testorg.lu')
        self.assertTrue(user)
        self.assertTrue('id' in user, user)
        uid = user['id']
        self.assertTrue('login' in user, user)
        self.assertTrue('apikey' in user, user)
        self.assertTrue('is_commenter' in user, user)
        self.assertTrue(user['is_commenter'])
        self.client.set_apikey(user['apikey'])
        comment = {'title': 'test', 'description': 'test',
                   'vulnerability': 'CVE-2024-20401',
                   'related_vulnerabilities': ['CVE-2024-20402']}
        created_comment = self.client.create_comment(comment=comment)
        new_comment_uuid = created_comment['data'][0]['uuid']
        comments = self.client.get_comments(uuid=new_comment_uuid)
        self.assertTrue(len(comments['data']) == 1, comments)
        deleted_comment = self.client.delete_comment(new_comment_uuid)
        self.assertTrue(deleted_comment < 300)

        self.client.set_apikey(self.admin_token)
        deleted_user = self.client.delete_user(str(uid))
        self.assertTrue(deleted_user < 300)

    # test Sightings
    def test_sightings_public(self) -> None:
        if not self.public_test:
            return None
        sighting_cve = self.client.get_sighting('6febe45d-d8de-4df7-b3ba-6cf7acd2e2b5')
        self.assertTrue(sighting_cve)
        self.assertTrue('uuid' in sighting_cve)
        self.assertEqual(sighting_cve['uuid'], '6febe45d-d8de-4df7-b3ba-6cf7acd2e2b5')

        sighting_cve_list = self.client.get_sightings(sighting_uuid='6febe45d-d8de-4df7-b3ba-6cf7acd2e2b5')
        self.assertTrue(sighting_cve_list)
        self.assertTrue('data' in sighting_cve_list)
        self.assertTrue(len(sighting_cve_list['data']) == 1)
        self.assertEqual(sighting_cve_list['data'][0], sighting_cve)

        sighting_cve_list = self.client.get_sightings(vuln_id='CVE-2020-3532')
        self.assertTrue(sighting_cve_list)
        self.assertTrue('data' in sighting_cve_list)
        self.assertTrue(len(sighting_cve_list['data']) > 0)

        sighting_cve_list = self.client.get_sightings(author='automation')
        self.assertTrue(sighting_cve_list)
        self.assertTrue('data' in sighting_cve_list)
        self.assertTrue(len(sighting_cve_list['data']) > 0)

        sighting_cve_list = self.client.get_sightings(author='automation', vuln_id='CVE-2020-3532')
        self.assertTrue(sighting_cve_list)
        self.assertTrue('data' in sighting_cve_list)
        self.assertTrue(len(sighting_cve_list['data']) > 0)

        sighting_cve_list = self.client.get_sightings(author='automation', date_from=datetime.now(tz=timezone.utc) - timedelta(days=1), date_to=datetime.now(tz=timezone.utc))
        self.assertTrue(sighting_cve_list)
        self.assertTrue('data' in sighting_cve_list)
        self.assertTrue(len(sighting_cve_list['data']) > 0)

    def test_sightings_local(self) -> None:
        if not self.admin_token:
            # this test is only working if the admin token is set
            return None

        u1 = str(uuid.uuid4())
        sighting = self.client.create_sighting(
            sighting={
                "vulnerability": "CVE-2024-20401",
                "source": u1,
                "type": "seen"
            }
        )
        self.assertTrue(sighting)
        print(sighting)
        s = self.client.get_sighting(sighting['data'][0]['uuid'])
        self.assertTrue('uuid' in s)
        self.assertTrue('vulnerability' in s)
        self.assertTrue('source' in s)
        self.assertTrue('type' in s)
        self.assertEqual(s['source'], u1)

        u2 = str(uuid.uuid4())
        sighting = self.client.create_sighting(source=u2, sighting_type='seen', vulnerability='CVE-2024-20401')
        s = self.client.get_sighting(sighting['data'][0]['uuid'])
        self.assertEqual(s['source'], u2)