From 88df1823e25fb07c7774752e0ff03d6870872d54 Mon Sep 17 00:00:00 2001
From: Jakob Langdal <jakob.langdal@alexandra.dk>
Date: Tue, 20 Apr 2021 21:21:56 +0200
Subject: [PATCH] Preliminary encryption support

---
 optimizerapi/optimizer.py             | 11 ++--------
 optimizerapi/securepickle/__init__.py |  1 +
 optimizerapi/securepickle/pickler.py  | 17 +++++++++++++++
 optimizerapi/securepickle/secure.py   | 14 ++++++++++++
 requirements.txt                      |  3 ++-
 tests/test_securepickle.py            | 31 +++++++++++++++++++++++++++
 6 files changed, 67 insertions(+), 10 deletions(-)
 create mode 100644 optimizerapi/securepickle/__init__.py
 create mode 100644 optimizerapi/securepickle/pickler.py
 create mode 100644 optimizerapi/securepickle/secure.py
 create mode 100644 tests/test_securepickle.py

diff --git a/optimizerapi/optimizer.py b/optimizerapi/optimizer.py
index 560998d..248f325 100644
--- a/optimizerapi/optimizer.py
+++ b/optimizerapi/optimizer.py
@@ -1,5 +1,3 @@
-import codecs
-import pickle
 import json
 from json_tricks import dumps
 from ProcessOptimizer import Optimizer, expected_minimum
@@ -10,6 +8,8 @@
 import base64
 import io 
 from numbers import Number
+from .securepickle import pickleToString, unpickleFromString
+
 import numpy
 numpy.random.seed(42)
 """ProcessOptimizer web request handler
@@ -18,13 +18,6 @@
 The handler functions are mapped to the OpenAPI specification through the "operationId" field
 in the specification.yml file found in the folder "openapi" in the root of this project.
 """
-def pickleToString(obj):
-    pickled = codecs.encode(pickle.dumps(obj), "base64").decode()
-    return pickled
-
-def unpickleFromString(pickled):
-    unpickled = pickle.loads(codecs.decode(pickled.encode(), "base64"))
-    return unpickled
 
 def run(body) -> dict:
     """Executes the ProcessOptimizer
diff --git a/optimizerapi/securepickle/__init__.py b/optimizerapi/securepickle/__init__.py
new file mode 100644
index 0000000..7582f09
--- /dev/null
+++ b/optimizerapi/securepickle/__init__.py
@@ -0,0 +1 @@
+from .pickler import pickleToString, unpickleFromString
\ No newline at end of file
diff --git a/optimizerapi/securepickle/pickler.py b/optimizerapi/securepickle/pickler.py
new file mode 100644
index 0000000..efa9878
--- /dev/null
+++ b/optimizerapi/securepickle/pickler.py
@@ -0,0 +1,17 @@
+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()
+    return pickled
+
+def unpickleFromString(pickled):
+    unpickled = pickle.loads(f.decrypt(codecs.decode(pickled.encode(), "base64")))
+    return unpickled
+
diff --git a/optimizerapi/securepickle/secure.py b/optimizerapi/securepickle/secure.py
new file mode 100644
index 0000000..b4ed0cb
--- /dev/null
+++ b/optimizerapi/securepickle/secure.py
@@ -0,0 +1,14 @@
+from cryptography.fernet import Fernet
+
+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
diff --git a/requirements.txt b/requirements.txt
index 32b37b6..351437f 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -2,4 +2,5 @@ connexion==2.7.0
 connexion[swagger-ui]
 Flask==1.1.2
 ProcessOptimizer==0.6.1
-json-tricks==3.15.5
\ No newline at end of file
+json-tricks==3.15.5
+cryptography==3.4.7
\ No newline at end of file
diff --git a/tests/test_securepickle.py b/tests/test_securepickle.py
new file mode 100644
index 0000000..a8ae6d7
--- /dev/null
+++ b/tests/test_securepickle.py
@@ -0,0 +1,31 @@
+from optimizerapi.securepickle import *
+import pickle
+
+def test_pickleToString():
+    encoded = pickleToString("myString")
+    assert encoded != "myString"
+
+def test_unpickleFromString():
+    encoded = pickleToString("myString")
+    decoded = unpickleFromString(encoded)
+    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
+
+
+