From 5609d33abbff4495f659b2b8550cd08c8a3570f4 Mon Sep 17 00:00:00 2001
From: Jakob Langdal <jakob.langdal@alexandra.dk>
Date: Tue, 20 Apr 2021 22:16:07 +0200
Subject: [PATCH] Encrypt pickled result

---
 optimizerapi/optimizer.py             |  6 +++---
 optimizerapi/securepickle/__init__.py |  3 ++-
 optimizerapi/securepickle/pickler.py  | 14 ++++----------
 optimizerapi/securepickle/secure.py   | 21 +++++++++-----------
 tests/test_securepickle.py            | 28 ++++++---------------------
 5 files changed, 24 insertions(+), 48 deletions(-)

diff --git a/optimizerapi/optimizer.py b/optimizerapi/optimizer.py
index 248f325..6903c61 100644
--- a/optimizerapi/optimizer.py
+++ b/optimizerapi/optimizer.py
@@ -8,7 +8,7 @@
 import base64
 import io 
 from numbers import Number
-from .securepickle import pickleToString, unpickleFromString
+from securepickle import pickleToString, unpickleFromString, get_crypto
 
 import numpy
 numpy.random.seed(42)
@@ -112,8 +112,8 @@ def processResult(result, optimizer, dimensions, cfg, data, space):
         plot_objective(result, dimensions=dimensions, usepartialdependence=False)
         addPlot(response["plots"], "objective", debug=True)
     
-    print(str(response))
-    response["pickle"] = pickleToString(result)
+    prettyResult["pickled"] = pickleToString(result, get_crypto())
+    # print(str(response))
     return response
 
 def addPlot(result, id="generic", close=True, debug=False):
diff --git a/optimizerapi/securepickle/__init__.py b/optimizerapi/securepickle/__init__.py
index 7582f09..b59a9a9 100644
--- a/optimizerapi/securepickle/__init__.py
+++ b/optimizerapi/securepickle/__init__.py
@@ -1 +1,2 @@
-from .pickler import pickleToString, unpickleFromString
\ No newline at end of file
+from .pickler import pickleToString, unpickleFromString
+from .secure import get_crypto
\ No newline at end of file
diff --git a/optimizerapi/securepickle/pickler.py b/optimizerapi/securepickle/pickler.py
index efa9878..ad029c7 100644
--- a/optimizerapi/securepickle/pickler.py
+++ b/optimizerapi/securepickle/pickler.py
@@ -1,17 +1,11 @@
 import codecs
 import pickle
-from .secure import create_key, load_key
-from cryptography.fernet import Fernet
 
-create_key()
-key = load_key()
-f = Fernet(key)
-
-def pickleToString(obj):
-    pickled = codecs.encode(f.encrypt(pickle.dumps(obj)), "base64").decode()
+def pickleToString(obj, crypto):
+    pickled = codecs.encode(crypto.encrypt(pickle.dumps(obj)), "base64").decode()
     return pickled
 
-def unpickleFromString(pickled):
-    unpickled = pickle.loads(f.decrypt(codecs.decode(pickled.encode(), "base64")))
+def unpickleFromString(pickled, crypto):
+    unpickled = pickle.loads(crypto.decrypt(codecs.decode(pickled.encode(), "base64")))
     return unpickled
 
diff --git a/optimizerapi/securepickle/secure.py b/optimizerapi/securepickle/secure.py
index b4ed0cb..c9aee33 100644
--- a/optimizerapi/securepickle/secure.py
+++ b/optimizerapi/securepickle/secure.py
@@ -1,14 +1,11 @@
 from cryptography.fernet import Fernet
+import os
 
-def is_initialized():
-    pass
-
-def load_key():
-    with open('mykey.key', 'rb') as mykey:
-        key = mykey.read()
-    return key
-
-def create_key():
-    key = Fernet.generate_key()
-    with open('mykey.key', 'wb') as mykey:
-        mykey.write(key)
\ No newline at end of file
+def get_crypto(key=None):
+    if key == None: key = os.getenv("PICKLE_KEY", None)
+    if key == None: 
+        print("No key found, generating new key")
+        key = Fernet.generate_key()
+        os.environ["PICKLE_KEY"] = key.decode("utf-8")
+        print("To reuse key for future server runs, set environment variable PICKLE_KEY=" + os.environ["PICKLE_KEY"])
+    return Fernet(key)
\ No newline at end of file
diff --git a/tests/test_securepickle.py b/tests/test_securepickle.py
index a8ae6d7..c367fc5 100644
--- a/tests/test_securepickle.py
+++ b/tests/test_securepickle.py
@@ -1,31 +1,15 @@
-from optimizerapi.securepickle import *
-import pickle
+from optimizerapi.securepickle import pickleToString, unpickleFromString, get_crypto
 
 def test_pickleToString():
-    encoded = pickleToString("myString")
+    f = get_crypto()
+    encoded = pickleToString("myString", f)
     assert encoded != "myString"
 
 def test_unpickleFromString():
-    encoded = pickleToString("myString")
-    decoded = unpickleFromString(encoded)
+    f = get_crypto()
+    encoded = pickleToString("myString", f)
+    decoded = unpickleFromString(encoded, f)
     assert decoded == "myString"
 
-# def test_load_key():
-#     create_key()
-#     key = load_key()
-#     assert key != ""
-#     assert len(key) == 44
-
-# def test_encrypt():
-#     create_key()
-#     key = load_key()
-#     f = Fernet(key)
-#     encoded = pickle.dumps("myString")
-#     encrypted = f.encrypt(encoded)
-#     assert encoded != encrypted
-#     assert type(encrypted) == str
-#     decrypted = f.decrypt(encrypted)
-#     assert encoded == decrypted
-