From 6d30fdb83dbc5f4f7eaccddd95658a5791c9020a Mon Sep 17 00:00:00 2001
From: Jeremy Stretch <jstretch@digitalocean.com>
Date: Tue, 21 Mar 2017 15:30:36 -0400
Subject: [PATCH] Finished work on secrets views; removed path from cookie
 assignment

---
 netbox/secrets/api/views.py |  5 ++++-
 netbox/secrets/views.py     | 40 +++++++++++++++++++++----------------
 netbox/users/views.py       |  2 +-
 3 files changed, 28 insertions(+), 19 deletions(-)

diff --git a/netbox/secrets/api/views.py b/netbox/secrets/api/views.py
index adab5e51d..63476b126 100644
--- a/netbox/secrets/api/views.py
+++ b/netbox/secrets/api/views.py
@@ -169,6 +169,9 @@ class GetSessionKeyViewSet(ViewSet):
         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 not isinstance(encoded_key, str):
+            encoded_key = encoded_key.decode()
 
         # Craft the response
         response = Response({
@@ -177,7 +180,7 @@ class GetSessionKeyViewSet(ViewSet):
 
         # If token authentication is not in use, assign the session key as a cookie
         if request.auth is None:
-            response.set_cookie('session_key', value=encoded_key, path=reverse('secrets-api:secret-list'))
+            response.set_cookie('session_key', value=encoded_key)
 
         return response
 
diff --git a/netbox/secrets/views.py b/netbox/secrets/views.py
index 5a90eca6e..8a1b35a6b 100644
--- a/netbox/secrets/views.py
+++ b/netbox/secrets/views.py
@@ -79,23 +79,30 @@ def secret_add(request, pk):
         form = forms.SecretForm(request.POST, instance=secret)
         if form.is_valid():
 
-            # Retrieve the master key from the current user's UserKey
-            master_key = uk.get_master_key(form.cleaned_data['private_key'])
-            if master_key is None:
-                form.add_error(None, "Invalid private key! Unable to encrypt secret data.")
+            # We need a valid session key in order to create a Secret
+            session_key = base64.b64decode(request.COOKIES.get('session_key', None))
+            if session_key is None:
+                form.add_error(None, "No session key was provided with the request. Unable to encrypt secret data.")
 
             # Create and encrypt the new Secret
             else:
-                secret = form.save(commit=False)
-                secret.plaintext = str(form.cleaned_data['plaintext'])
-                secret.encrypt(master_key)
-                secret.save()
+                master_key = None
+                try:
+                    sk = SessionKey.objects.get(userkey__user=request.user)
+                    master_key = sk.get_master_key(session_key)
+                except SessionKey.DoesNotExist:
+                    form.add_error(None, "No session key found for this user.")
 
-                messages.success(request, u"Added new secret: {}.".format(secret))
-                if '_addanother' in request.POST:
-                    return redirect('dcim:device_addsecret', pk=device.pk)
-                else:
-                    return redirect('secrets:secret', pk=secret.pk)
+                if master_key is not None:
+                    secret = form.save(commit=False)
+                    secret.plaintext = str(form.cleaned_data['plaintext'])
+                    secret.encrypt(master_key)
+                    secret.save()
+                    messages.success(request, u"Added new secret: {}.".format(secret))
+                    if '_addanother' in request.POST:
+                        return redirect('dcim:device_addsecret', pk=device.pk)
+                    else:
+                        return redirect('secrets:secret', pk=secret.pk)
 
     else:
         form = forms.SecretForm(instance=secret)
@@ -118,14 +125,13 @@ def secret_edit(request, pk):
         if form.is_valid():
 
             # Re-encrypt the Secret if a plaintext and session key have been provided.
-            session_key = request.COOKIES.get('session_key', None)
+            session_key = base64.b64decode(request.COOKIES.get('session_key', None))
             if form.cleaned_data['plaintext'] and session_key is not None:
 
                 # Retrieve the master key using the provided session key
-                session_key = base64.b64decode(session_key)
                 master_key = None
                 try:
-                    sk = SessionKey.objects.get(user=request.user)
+                    sk = SessionKey.objects.get(userkey__user=request.user)
                     master_key = sk.get_master_key(session_key)
                 except SessionKey.DoesNotExist:
                     form.add_error(None, "No session key found for this user.")
@@ -186,7 +192,7 @@ def secret_import(request):
             session_key = base64.b64decode(session_key)
             master_key = None
             try:
-                sk = SessionKey.objects.get(user=request.user)
+                sk = SessionKey.objects.get(userkey__user=request.user)
                 master_key = sk.get_master_key(session_key)
             except SessionKey.DoesNotExist:
                 form.add_error(None, "No session key found for this user.")
diff --git a/netbox/users/views.py b/netbox/users/views.py
index 41cecb96a..711c5a9c0 100644
--- a/netbox/users/views.py
+++ b/netbox/users/views.py
@@ -149,7 +149,7 @@ class SessionKeyDeleteView(LoginRequiredMixin, View):
 
             # Delete cookie
             response = redirect('user:userkey')
-            response.delete_cookie('session_key', path=reverse('secrets-api:secret-list'))
+            response.delete_cookie('session_key')
 
             return response