#!/usr/bin/env python3

import unittest
import time
import os

from pyvulnerabilitylookup import PyVulnerabilityLookup


class TestPublic(unittest.TestCase):

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

    #  Test default

    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:
        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:
        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(self) -> 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='admin')
        self.assertTrue(len(comments['data']) >= 1)
        for comment in comments['data']:
            self.assertEqual(comment['author']['login'], 'admin')

        comments = self.client.get_comments(uuid='a309d024-2714-4a81-a425-60f83f6d5740',
                                            vuln_id='CVE-2024-20401',
                                            author='admin')
        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')

    # TODO: POST / Delete Comment
    # TODO: POST / Get user

    # Test bundles

    def test_get_bundles(self) -> 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='admin')
        self.assertTrue(len(bundles['data']) >= 1)
        for bundle in bundles['data']:
            self.assertEqual(bundle['author']['login'], 'admin')

        bundles = self.client.get_bundles(uuid='a23cbcad-e890-4df8-8736-9332ed4c3d47',
                                          vuln_id='CVE-2024-39573',
                                          author='admin')
        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')