netbox/netbox/secrets/tests/test_views.py

import base64

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 StandardTestCases
from .constants import PRIVATE_KEY, PUBLIC_KEY


class SecretRoleTestCase(StandardTestCases.Views):
    model = SecretRole

    # Disable inapplicable tests
    test_get_object = None
    test_delete_object = None

    @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'),
        ])

        cls.form_data = {
            'name': 'Secret Role X',
            'slug': 'secret-role-x',
            'description': 'A secret role',
            'users': [],
            'groups': [],
        }

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


class SecretTestCase(StandardTestCases.Views):
    model = Secret

    # Disable inapplicable tests
    test_create_object = None

    # TODO: Check permissions enforcement on secrets.views.secret_edit
    test_edit_object = None

    @classmethod
    def setUpTestData(cls):

        site = Site.objects.create(name='Site 1', slug='site-1')
        manufacturer = Manufacturer.objects.create(name='Manufacturer 1', slug='manufacturer-1')
        devicetype = DeviceType.objects.create(manufacturer=manufacturer, model='Device Type 1')
        devicerole = DeviceRole.objects.create(name='Device Role 1', slug='device-role-1')
        device = Device.objects.create(name='Device 1', site=site, device_type=devicetype, device_role=devicerole)
        secretrole = SecretRole.objects.create(name='Secret Role 1', slug='secret-role-1')

        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'),
        ])

        cls.form_data = {
            'device': device.pk,
            'role': secretrole.pk,
            'name': 'Secret X',
        }

    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_import_objects(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.assertHttpStatus(response, 200)
        self.assertEqual(Secret.objects.count(), 6)