netbox/netbox/secrets/tests/test_views.py

import base64
import urllib.parse

from django.urls import reverse

from dcim.models import Device, DeviceRole, DeviceType, Manufacturer, Site
from secrets.models import Secret, SecretRole, SessionKey, UserKey
from utilities.testing import TestCase
from .constants import PRIVATE_KEY, PUBLIC_KEY


class SecretRoleTestCase(TestCase):
    user_permissions = (
        'secrets.view_secretrole',
    )

    @classmethod
    def setUpTestData(cls):

        SecretRole.objects.bulk_create([
            SecretRole(name='Secret Role 1', slug='secret-role-1'),
            SecretRole(name='Secret Role 2', slug='secret-role-2'),
            SecretRole(name='Secret Role 3', slug='secret-role-3'),
        ])

    def test_secretrole_list(self):

        url = reverse('secrets:secretrole_list')

        response = self.client.get(url, follow=True)
        self.assertEqual(response.status_code, 200)

    def test_secretrole_import(self):
        self.add_permissions('secrets.add_secretrole')

        csv_data = (
            "name,slug",
            "Secret Role 4,secret-role-4",
            "Secret Role 5,secret-role-5",
            "Secret Role 6,secret-role-6",
        )

        response = self.client.post(reverse('secrets:secretrole_import'), {'csv': '\n'.join(csv_data)})

        self.assertEqual(response.status_code, 200)
        self.assertEqual(SecretRole.objects.count(), 6)


class SecretTestCase(TestCase):
    user_permissions = (
        'secrets.view_secret',
    )

    @classmethod
    def setUpTestData(cls):

        site = Site(name='Site 1', slug='site-1')
        site.save()

        manufacturer = Manufacturer(name='Manufacturer 1', slug='manufacturer-1')
        manufacturer.save()

        devicetype = DeviceType(manufacturer=manufacturer, model='Device Type 1')
        devicetype.save()

        devicerole = DeviceRole(name='Device Role 1', slug='device-role-1')
        devicerole.save()

        device = Device(name='Device 1', site=site, device_type=devicetype, device_role=devicerole)
        device.save()

        secretrole = SecretRole(name='Secret Role 1', slug='secret-role-1')
        secretrole.save()

        Secret.objects.bulk_create([
            Secret(device=device, role=secretrole, name='Secret 1', ciphertext=b'1234567890'),
            Secret(device=device, role=secretrole, name='Secret 2', ciphertext=b'1234567890'),
            Secret(device=device, role=secretrole, name='Secret 3', ciphertext=b'1234567890'),
        ])

    def setUp(self):

        super().setUp()

        # Set up a master key for the test user
        userkey = UserKey(user=self.user, public_key=PUBLIC_KEY)
        userkey.save()
        master_key = userkey.get_master_key(PRIVATE_KEY)
        self.session_key = SessionKey(userkey=userkey)
        self.session_key.save(master_key)

    def test_secret_list(self):

        url = reverse('secrets:secret_list')
        params = {
            "role": SecretRole.objects.first().slug,
        }

        response = self.client.get('{}?{}'.format(url, urllib.parse.urlencode(params)), follow=True)
        self.assertEqual(response.status_code, 200)

    def test_secret(self):

        secret = Secret.objects.first()
        response = self.client.get(secret.get_absolute_url(), follow=True)
        self.assertEqual(response.status_code, 200)

    def test_secret_import(self):
        self.add_permissions('secrets.add_secret')

        csv_data = (
            "device,role,name,plaintext",
            "Device 1,Secret Role 1,Secret 4,abcdefghij",
            "Device 1,Secret Role 1,Secret 5,abcdefghij",
            "Device 1,Secret Role 1,Secret 6,abcdefghij",
        )

        # Set the session_key cookie on the request
        session_key = base64.b64encode(self.session_key.key).decode('utf-8')
        self.client.cookies['session_key'] = session_key

        response = self.client.post(reverse('secrets:secret_import'), {'csv': '\n'.join(csv_data)})

        self.assertEqual(response.status_code, 200)
        self.assertEqual(Secret.objects.count(), 6)