from utilities.testing import APIViewTestCases
from dcim.models import Device, DeviceRole, Manufacturer, DeviceType, Site
from django.contrib.contenttypes.models import ContentType

from nb_risk.tests.custom import APITestCase
from nb_risk.models import VulnerabilityAssignment, Vulnerability
from nb_risk import choices




from nb_risk.tests.custom import APITestCase


class VulnerabilityAssignmentAPITestCase(
    APITestCase,
    APIViewTestCases.GetObjectViewTestCase,
    APIViewTestCases.ListObjectsViewTestCase,
    APIViewTestCases.CreateObjectViewTestCase,
    APIViewTestCases.UpdateObjectViewTestCase,
    APIViewTestCases.DeleteObjectViewTestCase,
):

    model = VulnerabilityAssignment


    @classmethod
    def setUpTestData(cls):
        site = Site.objects.create(name="Site 1")
        manufacturer = Manufacturer.objects.create(name='Manufacturer 1')
        device_type = DeviceType.objects.create(manufacturer=manufacturer)
        device_role = DeviceRole.objects.create(name='Device Role 1')

        devices = (
            Device(name='Device 1', device_type=device_type, role=device_role, site=site),
            Device(name='Device 2', device_type=device_type, role=device_role, site=site),
            Device(name='Device 3', device_type=device_type, role=device_role, site=site),
        )
        Device.objects.bulk_create(devices)

        vulnerabilities = (
            Vulnerability(
                name="Vulnerability 1",
                cve="CVE-2021-1234",
                description="A test vulnerability",
            ),
            Vulnerability(
                name="Vulnerability 2",
                cve="CVE-2021-1235",
                description="Another test vulnerability",
            ),
            Vulnerability(
                name="Vulnerability 3",
                cve="CVE-2021-1236",
                description="Yet another test vulnerability",
            ),
        )
        Vulnerability.objects.bulk_create(vulnerabilities)


        v_assignments = (
            VulnerabilityAssignment(
                vulnerability=vulnerabilities[0],
                asset=devices[0],
            ),
            VulnerabilityAssignment(
                vulnerability=vulnerabilities[1],
                asset=devices[1],
            ),
            VulnerabilityAssignment(
                vulnerability=vulnerabilities[2],
                asset=devices[2],
            ),
        )
        VulnerabilityAssignment.objects.bulk_create(v_assignments)

        cls.create_data = [
            {
                "vulnerability": vulnerabilities[1].name,
                "asset_object_type": f"{ContentType.objects.get_for_model(devices[0]).app_label}.{ContentType.objects.get_for_model(devices[0]).model}",
                "asset_id": devices[0].pk,
            },
            {
                "vulnerability": vulnerabilities[2].name,
                "asset_object_type": f"{ContentType.objects.get_for_model(devices[1]).app_label}.{ContentType.objects.get_for_model(devices[1]).model}",
                "asset_id": devices[1].pk,
            },
            {
                "vulnerability": vulnerabilities[0].name,
                "asset_object_type": f"{ContentType.objects.get_for_model(devices[2]).app_label}.{ContentType.objects.get_for_model(devices[2]).model}",
                "asset_id": devices[2].pk,
            },
        ]

        cls.validation_excluded_fields = ['vulnerability']

        cls.brief_fields = [ 'display', 'id', 'url'   ]