فهرست منبع

Move utility functions for secrets to secrets/utils.py

Jeremy Stretch 6 سال پیش
والد
کامیت
f27e06e619
3فایلهای تغییر یافته به همراه35 افزوده شده و 29 حذف شده
  1. 2 28
      netbox/secrets/models.py
  2. 2 1
      netbox/secrets/tests/test_models.py
  3. 31 0
      netbox/secrets/utils.py

+ 2 - 28
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

+ 2 - 1
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):

+ 31 - 0
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)