
import json
from io import BytesIO
import gzip
from unittest.mock import patch

from django.urls import reverse
from django.test import TestCase
from django.test import Client as TestClient
from django.utils import timezone
from django.core.management import call_command
import requests
from rest_framework import status

from api.models import Vulnerability


class BaseMockHTTPResponse:
    def raise_for_status(self):
        pass


class TestVulnerability(TestCase):

    def setUp(self):
        self.client = TestClient()


    def test_vi_is_marked_as_being_displayed(self):
        vi = Vulnerability.objects.create(
            cve_identifier="cve-1111", description="foobar")

        self.assertEqual(Vulnerability.objects.count(), 1)
        self.assertIsNone(vi.displayed_at)

        url = reverse("api-get-vulnerability")
        response = self.client.get(url)
        self.assertEqual(response.status_code, status.HTTP_200_OK)

        vi.refresh_from_db()
        self.assertIsNotNone(vi.displayed_at)
        self.assertTrue(vi.displayed_once)


    def test_vi_displayed_at_is_reset_when_all_vi_have_been_displayed(self):
        vi1 = Vulnerability.objects.create(
            cve_identifier="cve-1111", description="foobar",
            displayed_at=timezone.now(), displayed_once=True)
        vi2 = Vulnerability.objects.create(
            cve_identifier="cve-2222", description="foobar",
            displayed_at=timezone.now(), displayed_once=True)

        url = reverse("api-get-vulnerability")
        response = self.client.get(url)
        self.assertEqual(response.status_code, status.HTTP_200_OK)
        self.assertEqual(response.data['cve_identifier'], "cve-1111")


        vi1.refresh_from_db()
        vi2.refresh_from_db()
        self.assertIsNotNone(vi1.displayed_at)
        self.assertIsNone(vi2.displayed_at)
        self.assertTrue(vi1.displayed_once)
        self.assertTrue(vi2.displayed_once)


    def test_vulnerability_items_are_created_when_cronjob_runs(self):
        mock_get = patch.object(requests, "get").start()
        class MockHTTPResponse(BaseMockHTTPResponse):
            @property
            def content(self):
                bytes_data = json.dumps({"CVE_Items":[
                    {
                        'cve':{
                            'CVE_data_meta':{
                                "ID":"cve-11111"
                            },
                            "description":{
                                "description_data":[
                                    {
                                        "value":"foobar"
                                    }
                                ]
                            }
                        }
                    }
                ]}).encode()
                buffer = BytesIO()
                with gzip.GzipFile(fileobj=buffer, mode='w') as f:
                    f.write(bytes_data)
                return buffer.getvalue()

        mock_get.return_value = MockHTTPResponse()

        self.assertEqual(Vulnerability.objects.count(), 0)
        call_command("update_vulnerabilities_records")
        mock_get.assert_called_once()
        self.assertEqual(Vulnerability.objects.count(), 1)

        vi = Vulnerability.objects.first()
        self.assertEqual(vi.cve_identifier, "cve-11111")
        self.assertEqual(vi.description, "foobar")
        self.assertFalse(vi.displayed_once)
        self.assertIsNone(vi.displayed_at)
        mock_get.stop()

    def test_vulnerability_items_are_deleted_when_cronjob_runs_and_item_has_been_displayed_and_no_longer_present_in_feed(self):
        vi1 = Vulnerability.objects.create(
            cve_identifier="cve-1111", description="foobar",
            displayed_at=timezone.now(), displayed_once=True)

        vi1_id = vi1.id

        mock_get = patch.object(requests, "get").start()
        class MockHTTPResponse(BaseMockHTTPResponse):
            @property
            def content(self):
                bytes_data = json.dumps({"CVE_Items":[
                    {
                        'cve':{
                            'CVE_data_meta':{
                                "ID":"cve-2222"
                            },
                            "description":{
                                "description_data":[
                                    {
                                        "value":"foobar"
                                    }
                                ]
                            }
                        }
                    }
                ]}).encode()
                buffer = BytesIO()
                with gzip.GzipFile(fileobj=buffer, mode='w') as f:
                    f.write(bytes_data)
                return buffer.getvalue()
        mock_get.return_value = MockHTTPResponse()

        call_command("update_vulnerabilities_records")
        mock_get.assert_called_once()
        self.assertEqual(Vulnerability.objects.count(), 1)

        self.assertFalse(Vulnerability.objects.filter(id=vi1_id).exists())
        vi = Vulnerability.objects.first()
        self.assertEqual(vi.cve_identifier, "cve-2222")
        self.assertEqual(vi.description, "foobar")
        self.assertFalse(vi.displayed_once)
        self.assertIsNone(vi.displayed_at)

        mock_get.stop()


    def test_vulnerability_items_are_not_deleted_when_cronjob_runs_and_item_has_not_been_displayed_and_no_longer_present_in_feed(self):
        vi1 = Vulnerability.objects.create(
            cve_identifier="cve-1111", description="foobar",
            displayed_at=None, displayed_once=False)

        vi1_id = vi1.id

        mock_get = patch.object(requests, "get").start()
        class MockHTTPResponse(BaseMockHTTPResponse):
            @property
            def content(self):
                bytes_data = json.dumps({"CVE_Items":[
                    {
                        'cve':{
                            'CVE_data_meta':{
                                "ID":"cve-2222"
                            },
                            "description":{
                                "description_data":[
                                    {
                                        "value":"foobar"
                                    }
                                ]
                            }
                        }
                    }
                ]}).encode()
                buffer = BytesIO()
                with gzip.GzipFile(fileobj=buffer, mode='w') as f:
                    f.write(bytes_data)
                return buffer.getvalue()
        mock_get.return_value = MockHTTPResponse()

        call_command("update_vulnerabilities_records")
        mock_get.assert_called_once()
        self.assertEqual(Vulnerability.objects.count(), 2)

        mock_get.stop()


    def test_vulnerability_items_are_not_deleted_when_cronjob_runs_and_item_has_been_displayed_and_is_present_in_feed(self):
        vi1 = Vulnerability.objects.create(
            cve_identifier="cve-1111", description="foobar",
            displayed_at=timezone.now(), displayed_once=True)

        vi1_id = vi1.id

        mock_get = patch.object(requests, "get").start()
        class MockHTTPResponse(BaseMockHTTPResponse):
            @property
            def content(self):
                bytes_data = json.dumps({"CVE_Items":[
                    {
                        'cve':{
                            'CVE_data_meta':{
                                "ID":"cve-1111"
                            },
                            "description":{
                                "description_data":[
                                    {
                                        "value":"foobar"
                                    }
                                ]
                            }
                        }
                    },{
                        'cve':{
                            'CVE_data_meta':{
                                "ID":"cve-2222"
                            },
                            "description":{
                                "description_data":[
                                    {
                                        "value":"foobar"
                                    }
                                ]
                            }
                        }
                    }
                ]}).encode()
                buffer = BytesIO()
                with gzip.GzipFile(fileobj=buffer, mode='w') as f:
                    f.write(bytes_data)
                return buffer.getvalue()
        mock_get.return_value = MockHTTPResponse()

        call_command("update_vulnerabilities_records")
        mock_get.assert_called_once()
        self.assertEqual(Vulnerability.objects.count(), 2)

        mock_get.stop()
