Datatype conversion to bytes is either done implicit or in the algorithms, as pickles were python only. Making it for users impossible to work from different languages.

Author: frankcorneliusmartin

vantage/node/encryption.py

@@ -90,20 +90,8 @@ def public_key_str(self):
         """
         return prepare_bytes_for_transport(self.public_key_bytes)
 
-    def encrypt_obj_to_base64(self, msg: dict, public_key_base64: str) -> str:
-        """ Encrypt dictonairy `msg` using `public_key_base64`.
-        """
-        msg_str = pickle.dumps(msg)
-        return self.encrypt_bytes_to_base64(msg_str, public_key_base64)
-
-    def encrypt_str_to_base64(self, msg: str, public_key_base64: str) -> str:
-        """ Encrypt string `msg` using `public_key_base64`.
-        """
-        msg_bytes = msg.encode(cs.STRING_ENCODING)
-        return self.encrypt_bytes_to_base64(msg_bytes, public_key_base64)
-
-    def encrypt_bytes_to_base64(
-    self, msg: bytes, public_key_base64: str) -> str:
+    def encrypt_bytes_to_base64(self, msg: bytes, 
+        public_key_base64: str) -> str:
         """ Encrypt a `msg` using `public_key_base64`.
 
             :param msg: message to be encrypted
@@ -189,22 +177,6 @@ def decrypt_bytes_from_base64(self, msg: str) -> bytes:
             return self.decrypt_bytes(msg_bytes)
         else:
             return b""
-
-    def decrypt_str_from_base64(self, msg: str) -> str:
-        """ Decrypt base64 `msg` using our private key
-
-            :param msg: string utf-8 encoded base64 encrypted msg
-        """
-        msg_bytes = self.decrypt_bytes_from_base64(msg)
-        return msg_bytes.decode(cs.STRING_ENCODING)
-
-    def decrypt_obj_from_base64(self, msg: str) -> dict:
-        """ Decrypt base64 `msg` using our private key.
-
-            :param msg: dict utf-8 encoded base64 encrypted msg
-        """
-        msg_str = self.decrypt_bytes_from_base64(msg)
-        return pickle.loads(msg_str)
 
     def __load_private_key(self, private_key_file=None):
         """ Load a private key file into this instance.

vantage/node/proxy_server.py

@@ -16,7 +16,11 @@
 
 from flask import Flask, request, jsonify
 
-from vantage.util import logger_name
+from vantage.util import (
+    logger_name, 
+    unpack_bytes_from_transport, 
+    prepare_bytes_for_transport
+)
 from vantage.node.server_io import ClientNodeProtocol
 from vantage.node.encryption import Cryptor
 
@@ -52,7 +56,7 @@ def proxy_task():
         TODO we might not want to use the SERVER_IO, as we only use the
             encryption property of it
         TODO if no public key is present, we should have some sort of 
-            faalback
+            fallback
     """
     assert app.config["SERVER_IO"], "Server IO not initialized"
 
@@ -78,7 +82,7 @@ def proxy_task():
     log.debug(f"{n_organizations} organizations, attemping to encrypt")
     encrypted_organizations = []
     for organization in organizations:
-        input_ = organization.get("input", None)
+        input_ = organization.get("input", b"")
         if not input_:
             log.error("No input for organization?!")
             return
@@ -94,8 +98,10 @@ def proxy_task():
 
         public_key = response.json().get("public_key")
 
-        encrypted_input = server_io.cryptor.encrypt_obj_to_base64(
-            input_, public_key)
+        input_unpacked = unpack_bytes_from_transport(input_)
+        encrypted_input = server_io.cryptor.encrypt_bytes_to_base64(
+            input_unpacked, public_key)
+
         log.debug(f"should be unreadable={encrypted_input}")
         organization["input"] = encrypted_input
         encrypted_organizations.append(organization)

vantage/node/server_io.py

@@ -264,7 +264,7 @@ def refresh_token(self):
         self._access_token = response.json()["access_token"]
 
     def post_task(self, name:str, image:str, collaboration_id:int, 
-        input_:str='', description='', organization_ids:list=[]) -> dict:
+        input_:bytes=b'', description='', organization_ids:list=[]) -> dict:
         """ Post a new task at the server.
 
             It will also encrypt `input_` for each receiving 
@@ -281,28 +281,23 @@ def post_task(self, name:str, image:str, collaboration_id:int,
         """
         assert self.cryptor, "Encryption has not yet been setup!"
 
-        encryption_method = self.cryptor.encrypt_obj_to_base64
-        if type(input_) == str:
-            encryption_method = self.cryptor.encrypt_str_to_base64
-        elif type(input_) == dict:
-            encryption_method = self.cryptor.encrypt_obj_to_base64
-
         organization_json_list = []
         for org_id in organization_ids:
             pub_key = self.request(f"organization/{org_id}").get("public_key")
             pub_key = unpack_bytes_from_transport(pub_key)
             organization_json_list.append(
                 {
                     "id": org_id,
-                    "input": encryption_method(input_, pub_key)
+                    "input": self.cryptor.encrypt_bytes_to_base64(
+                        input_, pub_key)
                 }
             )
 
         return self.request('task', method='post', json={
             "name": name,
             "image": image, 
             "collaboration_id": collaboration_id,
-            "input": input_,
+            "input": input_, # TODO remove this
             "description": description,
             "organizations": organization_json_list
         })
@@ -673,7 +668,6 @@ def request_token_for_container(self, task_id: int, image: str):
                 container-token (a task results in a algorithm-
                 container at the node)
             :param image: image-name of the task
-        
         """
         self.log.debug(
             f"requesting container token for task_id={task_id} "
@@ -792,8 +786,11 @@ def patch_results(self, id: int, initiator_id: int, result: dict):
 
             self.log.debug(public_key)
 
-            result["result"] = self.cryptor.encrypt_obj_to_base64(
-                result["result"], public_key
+            results_unpacked = unpack_bytes_from_transport(
+                result["result"])
+            
+            result["result"] = self.cryptor.encrypt_bytes_to_base64(
+                results_unpacked, public_key
             )
             self.log.debug("Sending encrypted results to server")
 

vantage/tests/test_node_encryption.py

@@ -3,6 +3,7 @@
 import yaml
 import bcrypt
 import datetime
+import json
 
 # from sqlalchemy import create_engine
 # from sqlalchemy.orm import Session
@@ -71,14 +72,14 @@ def test_unpacking_transport_key(self):
         )
 
     def test_encryption_decryption(self):
-        msg = {"msg":"some message!"}
-        encrypted = self.cryptor.encrypt_obj_to_base64(
+        msg = json.dumps({"msg":"some message!"}).encode("ascii")
+        encrypted = self.cryptor.encrypt_bytes_to_base64(
             msg,
             self.cryptor.public_key_str
         )
         self.assertNotEqual(msg, encrypted)
 
-        unencrypted = self.cryptor.decrypt_obj_from_base64(
+        unencrypted = self.cryptor.decrypt_bytes_from_base64(
             encrypted
         )
         self.assertEqual(msg, unencrypted)

vantage/tests/test_proxy_server.py

@@ -17,7 +17,10 @@
 from vantage.node.encryption import Cryptor
 
 from vantage.constants import PACAKAGE_FOLDER, APPNAME, DATA_FOLDER, VERSION
-from vantage.util import unpack_bytes_from_transport
+from vantage.util import (
+    unpack_bytes_from_transport, 
+    prepare_bytes_for_transport
+)
 
 
 log = logging.getLogger(__name__.split(".")[-1])
@@ -154,13 +157,14 @@ def test_task(self):
         if not self.headers:
             self.login()
 
+        input_ = prepare_bytes_for_transport("bla".encode("ascii"))
         proxy_test = self.app.post(
             "task", 
             headers=self.headers,
             json={
                 "organizations":[{
                     "id":1,
-                    "input": "bla"
+                    "input":  input_
                 }],
                 "image": "some-image",
                 "collaboration_id": 1