Encrypt pickled result

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):

optimizerapi/securepickle/__init__.py

@@ -1 +1,2 @@
-from .pickler import pickleToString, unpickleFromString
+from .pickler import pickleToString, unpickleFromString
+from .secure import get_crypto

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
 

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

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
-