From f27e06e6194d91742b49bc8df46ced2f09bd4b90 Mon Sep 17 00:00:00 2001
From: Jeremy Stretch <stretch@packetlife.net>
Date: Tue, 14 Jan 2020 12:11:14 -0500
Subject: [PATCH] Move utility functions for secrets to secrets/utils.py

---
 netbox/secrets/models.py            | 30 ++--------------------------
 netbox/secrets/tests/test_models.py |  3 ++-
 netbox/secrets/utils.py             | 31 +++++++++++++++++++++++++++++
 3 files changed, 35 insertions(+), 29 deletions(-)
 create mode 100644 netbox/secrets/utils.py

diff --git a/netbox/secrets/models.py b/netbox/secrets/models.py
index e01e502c0..2a413b1ca 100644
--- a/netbox/secrets/models.py
+++ b/netbox/secrets/models.py
@@ -1,7 +1,7 @@
 import os
 import sys
 
-from Crypto.Cipher import AES, PKCS1_OAEP
+from Crypto.Cipher import AES
 from Crypto.PublicKey import RSA
 from Crypto.Util import strxor
 from django.conf import settings
@@ -19,6 +19,7 @@ from utilities.models import ChangeLoggedModel
 from .exceptions import InvalidKey
 from .hashers import SecretValidationHasher
 from .querysets import UserKeyQuerySet
+from .utils import encrypt_master_key, decrypt_master_key, generate_random_key
 
 
 __all__ = (
@@ -29,33 +30,6 @@ __all__ = (
 )
 
 
-def generate_random_key(bits=256):
-    """
-    Generate a random encryption key. Sizes is given in bits and must be in increments of 32.
-    """
-    if bits % 32:
-        raise Exception("Invalid key size ({}). Key sizes must be in increments of 32 bits.".format(bits))
-    return os.urandom(int(bits / 8))
-
-
-def encrypt_master_key(master_key, public_key):
-    """
-    Encrypt a secret key with the provided public RSA key.
-    """
-    key = RSA.importKey(public_key)
-    cipher = PKCS1_OAEP.new(key)
-    return cipher.encrypt(master_key)
-
-
-def decrypt_master_key(master_key_cipher, private_key):
-    """
-    Decrypt a secret key with the provided private RSA key.
-    """
-    key = RSA.importKey(private_key)
-    cipher = PKCS1_OAEP.new(key)
-    return cipher.decrypt(master_key_cipher)
-
-
 class UserKey(models.Model):
     """
     A UserKey stores a user's personal RSA (public) encryption key, which is used to generate their unique encrypted
diff --git a/netbox/secrets/tests/test_models.py b/netbox/secrets/tests/test_models.py
index b3ba0cee1..80bc953bc 100644
--- a/netbox/secrets/tests/test_models.py
+++ b/netbox/secrets/tests/test_models.py
@@ -7,7 +7,8 @@ from django.core.exceptions import ValidationError
 from django.test import TestCase
 
 from secrets.hashers import SecretValidationHasher
-from secrets.models import UserKey, Secret, encrypt_master_key, decrypt_master_key, generate_random_key
+from secrets.models import Secret, UserKey
+from secrets.utils import encrypt_master_key, decrypt_master_key, generate_random_key
 
 
 class UserKeyTestCase(TestCase):
diff --git a/netbox/secrets/utils.py b/netbox/secrets/utils.py
new file mode 100644
index 000000000..c43a392b3
--- /dev/null
+++ b/netbox/secrets/utils.py
@@ -0,0 +1,31 @@
+import os
+
+from Crypto.Cipher import PKCS1_OAEP
+from Crypto.PublicKey import RSA
+
+
+def generate_random_key(bits=256):
+    """
+    Generate a random encryption key. Sizes is given in bits and must be in increments of 32.
+    """
+    if bits % 32:
+        raise Exception("Invalid key size ({}). Key sizes must be in increments of 32 bits.".format(bits))
+    return os.urandom(int(bits / 8))
+
+
+def encrypt_master_key(master_key, public_key):
+    """
+    Encrypt a secret key with the provided public RSA key.
+    """
+    key = RSA.importKey(public_key)
+    cipher = PKCS1_OAEP.new(key)
+    return cipher.encrypt(master_key)
+
+
+def decrypt_master_key(master_key_cipher, private_key):
+    """
+    Decrypt a secret key with the provided private RSA key.
+    """
+    key = RSA.importKey(private_key)
+    cipher = PKCS1_OAEP.new(key)
+    return cipher.decrypt(master_key_cipher)