From 5d022a575a1f17efe192203870a4c011790ea154 Mon Sep 17 00:00:00 2001
From: Jeremy Stretch <jstretch@digitalocean.com>
Date: Tue, 28 Mar 2017 11:13:13 -0400
Subject: [PATCH] Closes #985: Added preserve_key to get-session-key endpoint

---
 netbox/secrets/api/views.py  | 34 ++++++++++++++++++++++++----------
 netbox/secrets/exceptions.py |  4 ++--
 netbox/secrets/models.py     | 15 +++++++++++++--
 3 files changed, 39 insertions(+), 14 deletions(-)

diff --git a/netbox/secrets/api/views.py b/netbox/secrets/api/views.py
index 4a44776c3..4e20ee63c 100644
--- a/netbox/secrets/api/views.py
+++ b/netbox/secrets/api/views.py
@@ -1,7 +1,6 @@
 import base64
 from Crypto.PublicKey import RSA
 
-from django.core.urlresolvers import reverse
 from django.http import HttpResponseBadRequest
 
 from rest_framework.exceptions import ValidationError
@@ -9,7 +8,7 @@ from rest_framework.permissions import IsAuthenticated
 from rest_framework.response import Response
 from rest_framework.viewsets import ModelViewSet, ViewSet
 
-from secrets.exceptions import InvalidSessionKey
+from secrets.exceptions import InvalidKey
 from secrets.filters import SecretFilter
 from secrets.models import Secret, SecretRole, SessionKey, UserKey
 from utilities.api import WritableSerializerMixin
@@ -84,7 +83,7 @@ class SecretViewSet(WritableSerializerMixin, ModelViewSet):
                 try:
                     sk = SessionKey.objects.get(userkey__user=request.user)
                     self.master_key = sk.get_master_key(session_key)
-                except (SessionKey.DoesNotExist, InvalidSessionKey):
+                except (SessionKey.DoesNotExist, InvalidKey):
                     raise ValidationError("Invalid session key.")
 
     def retrieve(self, request, *args, **kwargs):
@@ -140,6 +139,9 @@ class GetSessionKeyViewSet(ViewSet):
         {
             "session_key": "+8t4SI6XikgVmB5+/urhozx9O5qCQANyOk1MNe6taRf="
         }
+
+    This endpoint accepts one optional parameter: `preserve_key`. If True and a session key exists, the existing session
+    key will be returned instead of a new one.
     """
     permission_classes = [IsAuthenticated]
 
@@ -163,14 +165,26 @@ class GetSessionKeyViewSet(ViewSet):
         if master_key is None:
             return HttpResponseBadRequest(ERR_PRIVKEY_INVALID)
 
-        # Delete the existing SessionKey for this user if one exists
-        SessionKey.objects.filter(userkey__user=request.user).delete()
+        try:
+            current_session_key = SessionKey.objects.get(userkey__user_id=request.user.pk)
+        except SessionKey.DoesNotExist:
+            current_session_key = None
 
-        # Create a new SessionKey
-        sk = SessionKey(userkey=user_key)
-        sk.save(master_key=master_key)
-        encoded_key = base64.b64encode(sk.key)
-        # b64decode() returns a bytestring under Python 3
+        if current_session_key and request.GET.get('preserve_key', False):
+
+            # Retrieve the existing session key
+            key = current_session_key.get_session_key(master_key)
+
+        else:
+
+            # Create a new SessionKey
+            SessionKey.objects.filter(userkey__user=request.user).delete()
+            sk = SessionKey(userkey=user_key)
+            sk.save(master_key=master_key)
+            key = sk.key
+
+        # Encode the key using base64. (b64decode() returns a bytestring under Python 3.)
+        encoded_key = base64.b64encode(key)
         if not isinstance(encoded_key, str):
             encoded_key = encoded_key.decode()
 
diff --git a/netbox/secrets/exceptions.py b/netbox/secrets/exceptions.py
index 417d135de..c163e5907 100644
--- a/netbox/secrets/exceptions.py
+++ b/netbox/secrets/exceptions.py
@@ -1,5 +1,5 @@
-class InvalidSessionKey(Exception):
+class InvalidKey(Exception):
     """
-    Raised when the a provided session key is invalid.
+    Raised when a provided key is invalid.
     """
     pass
diff --git a/netbox/secrets/models.py b/netbox/secrets/models.py
index 7d7b1cd50..9afe2af84 100644
--- a/netbox/secrets/models.py
+++ b/netbox/secrets/models.py
@@ -13,7 +13,7 @@ from django.utils.encoding import force_bytes, python_2_unicode_compatible
 from dcim.models import Device
 from utilities.models import CreatedUpdatedModel
 
-from .exceptions import InvalidSessionKey
+from .exceptions import InvalidKey
 from .hashers import SecretValidationHasher
 
 
@@ -221,13 +221,24 @@ class SessionKey(models.Model):
 
         # Validate the provided session key
         if not check_password(session_key, self.hash):
-            raise InvalidSessionKey()
+            raise InvalidKey("Invalid session key")
 
         # Decrypt master key using provided session key
         master_key = xor_keys(session_key, bytes(self.cipher))
 
         return master_key
 
+    def get_session_key(self, master_key):
+
+        # Recover session key using the master key
+        session_key = xor_keys(master_key, bytes(self.cipher))
+
+        # Validate the recovered session key
+        if not check_password(session_key, self.hash):
+            raise InvalidKey("Invalid master key")
+
+        return session_key
+
 
 @python_2_unicode_compatible
 class SecretRole(models.Model):