refactor cape key/encrypt to accept username parameter (#111)

* refactor cape key/encrypt to accept username parameter

* add user_encrypt example

* keyword-only args to cape.key and cape.encrypt

* lint

examples/async_echo.py

@@ -9,7 +9,7 @@
 async def main(cape: pycape.Cape, deploy_path: os.PathLike, echo: str) -> str:
     function_ref = await cli.deploy(deploy_path)
     echo_arg = echo.encode()
-    echo_enc = await cape.encrypt(echo_arg, function_ref.token)
+    echo_enc = await cape.encrypt(echo_arg, token=function_ref.token)
     async with cape.function_context(function_ref):
         result = await cape.invoke(echo_enc)
     return result.decode()

examples/deploy_run_echo.py

@@ -10,8 +10,12 @@
 cape = pycape.Cape(url=CAPE_HOST)
 # Deploy Cape function
 function_ref = cli.deploy(ECHO_DEPLOY_PATH, token_expiry=100)
+print("Echo deployed:")
+print(f"\t- ID: {function_ref.id}")
+print(f"\t- Token: {function_ref.token}")
+print(f"\t- Checksum: {function_ref.checksum}")
 # Encrypt input
-message = cape.encrypt("Welcome to Cape".encode(), function_ref.token)
+message = cape.encrypt("Welcome to Cape".encode(), token=function_ref.token)
 # Run Cape function
 result = cape.run(function_ref, message)
 print(f"The result is: {result.decode()}")

examples/echo_token.json

@@ -1 +1 @@
-{"function_id":"8kpZeYKDQYZEfcXt5SHKjP","function_token":"eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJleHAiOjU1NDQ0NDQyMDQsImlhdCI6MTY2ODYyNDE4NCwiaXNzIjoiZ2l0aHVifDgwMDEwMTYiLCJzY29wZSI6ImZ1bmN0aW9uOmludm9rZSIsInN1YiI6IjhrcFplWUtEUVlaRWZjWHQ1U0hLalAifQ.r9pdo1z1-IsszN5F4fqz8TNcKbSCviBFRqhIB0zlIrAwjBKwMt8lkVIV98zcTe4rEJ7aEqyfiV3sSw7yi4RHuVzYCMasaVaGEMxBk_P5aBP7IxByyexngC5RGrYvVzlvjtOi9VIJSZvHr5IUJL06FKZR1gJMsVoRYlOQCMh8W-TCsBDJuwa7c5AshOFtuxu4HTxq-N2AR_CY-BUfHtuarXDJ0o4xRGXnGeTu37IX_PhHWeYIz68SJRqwWLrL5jxUn1Nr-JfUNWH3pN7EOdSTXKpe0qnWqiJaQTorRB36JVjwZdyeJ2y9NjWDxPUjd3efSKATz6d34xFw-JX4h1VyOA","function_checksum":"f654a25ddf65ed79bfe514178126d97e1141857edf10b781c2f8c11525163c88"}
+{"function_id":"ji3BJGBgnKDKoC9NvfCtE9","function_token":"eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJleHAiOjE3MDY4ODkzMTcsImlhdCI6MTY3NTM1MzMxNywiaXNzIjoiZ2l0aHVifDExNDY5NjYyMiIsInNjb3BlIjoiZnVuY3Rpb246aW52b2tlIiwic3ViIjoiamkzQkpHQmduS0RLb0M5TnZmQ3RFOSJ9.lr5QpJINij90pIT_nnv9Gpjk2jMrlxQhZpqg4pkMBYgZNAaJXz_sAD4ubR5Dav7jXUJWYx03hFBI8deB45GrIWTovG2kGLF5JdREXs1XY-F7YPvSxrICMSSleYWi3V1XIwOU8I5nzgp2NhpsrqO6FqsETyBqWOVP4EOogjAdas4rcp2IJE4qDQzNz7HLVLTZuNdVtcXz4YKMiK8n7VpNJkRX5MvAjqjAQ-HGLvpteybCXpK9IZYt3VF3Ge-I3Xw0b7OYa2WuCnfnh3coXnFtfWQEpaoLO5Q6kKkzCCNPtkYgoxNwsn5SxWJiY9KFWBI7ht7Kp-PYmoZAr_2iTOKmZw","function_checksum":"f654a25ddf65ed79bfe514178126d97e1141857edf10b781c2f8c11525163c88"}

examples/user_encrypt.py

@@ -0,0 +1,23 @@
+import pathlib
+
+import pycape
+
+function_json = pathlib.Path(__file__).parent.absolute() / "echo_token.json"
+function_ref = pycape.FunctionRef.from_json(function_json)
+
+cape = pycape.Cape()
+
+# Two options to encrypt for a Cape user
+
+# 1 - retrieve the user's key with cape.key, and pass the key to cape.encrypt
+capedocs_key = cape.key(username="capedocs")
+encrypted_data = cape.encrypt(b"encrypt against capedocs's key", key=capedocs_key)
+result = cape.run(function_ref, encrypted_data)
+print(result.decode())
+
+# 2 - pass username directly to cape.encrypt
+encrypted_data = cape.encrypt(
+    b"quickly encrypt my data for capedocs", username="capedocs"
+)
+result = cape.run(function_ref, encrypted_data)
+print(result.decode())

pycape/cape.py

@@ -46,6 +46,7 @@
 from typing import Optional
 from typing import Union
 
+import requests
 import synchronicity
 import websockets
 
@@ -128,7 +129,9 @@ async def connect(
     async def encrypt(
         self,
         input: bytes,
-        token: str,
+        *,
+        username: Optional[str] = None,
+        token: Optional[str] = None,
         key: Optional[bytes] = None,
         key_path: Optional[Union[str, os.PathLike]] = None,
     ) -> bytes:
@@ -144,6 +147,8 @@ async def encrypt(
 
         Args:
             input: Input bytes to encrypt.
+            username: A Github username corresponding to a Cape user who's public key
+                you want to use for the encryption. See :meth:`~Cape.key` for details.
             token: A Cape token scoped for retrieving a user's Cape public key.
                 See :meth:`~Cape.key` for details.
             key: Optional bytes for the Cape key. If None, will delegate to calling
@@ -163,7 +168,9 @@ async def encrypt(
             Exception: if the enclave threw an error while trying to fulfill the
                 connection request.
         """
-        cape_key = key or await self.key(token, key_path)
+        cape_key = key or await self.key(
+            username=username, token=token, key_path=key_path
+        )
         ctxt = cape_encrypt.encrypt(input, cape_key)
         # cape-encrypted ctxt must be b64-encoded and tagged
         ctxt = base64.b64encode(ctxt)
@@ -258,15 +265,25 @@ async def invoke(
     @_synchronizer
     async def key(
         self,
-        token: str,
+        *,
+        username: Optional[str] = None,
+        token: Optional[str] = None,
         key_path: Optional[Union[str, os.PathLike]] = None,
         pcrs: Optional[Dict[str, List[str]]] = None,
     ) -> bytes:
         """Load a Cape key from disk or download and persist an enclave-generated one.
 
+        The caller must provide one of the following arguments: ``username``, ``token``,
+        or ``key_path``.
+
         Args:
-            token: A string representing a Cape authentication token. Usually retrieved
-                from a :class:`~.function_ref.FunctionRef` via :attr:`FunctionRef.token`
+            username: An optional string representing the Github username of a Cape
+                user. The resulting public key will be associated with their account,
+                and data encrypted with this key will be available inside functions
+                that user has deployed.
+            token: An optional string representing a Cape authentication token. Usually
+                retrieved from a :class:`~.function_ref.FunctionRef` via
+                :attr:`FunctionRef.token`
             key_path: The path to the Cape key file. If the file already exists, the key
                 will be read from disk and returned. Otherwise, a Cape key will be
                 requested from the Cape platform and written to this location.
@@ -285,12 +302,26 @@ async def key(
             Exception: if the enclave threw an error while trying to fulfill the
                 connection request.
         """
+        if username is None and token is None and key_path is None:
+            raise ValueError(
+                "Must supply one of [`username`, `token`, `key_path`] arguments, but "
+                "found `None` for all of them."
+            )
+        if username is not None and token is not None:
+            raise ValueError(
+                "Provided both `username` and `token` arguments, but these are "
+                "mutually exclusive."
+            )
+
         if key_path is None:
             config_dir = pathlib.Path(cape_config.LOCAL_CONFIG_DIR)
+            key_qualifier = (
+                username or token[-200:]
+            )  # Shorten file name to avoid Errno 63 when saving file
             key_path = (
                 config_dir
                 / "encryption_keys"
-                / token[-200:]  # Shorten file name to avoid Errno 63 when saving file
+                / key_qualifier
                 / cape_config.LOCAL_CAPE_KEY_FILENAME
             )
         else:
@@ -300,7 +331,11 @@ async def key(
             with open(key_path, "rb") as f:
                 cape_key = f.read()
         else:
-            cape_key = await self._request_key(token, key_path, pcrs=pcrs)
+            if username is not None:
+                cape_key = await self._request_key_with_username(username, pcrs=pcrs)
+            else:
+                cape_key = await self._request_key_with_token(token, pcrs=pcrs)
+            await _persist_cape_key(cape_key, key_path)
 
         return cape_key
 
@@ -432,10 +467,33 @@ async def _request_invocation(self, serde_hooks, use_serdio, *args, **kwargs):
 
         return result
 
-    async def _request_key(
+    async def _request_key_with_username(
+        self,
+        username: str,
+        pcrs: Optional[Dict[str, List[str]]] = None,
+    ) -> bytes:
+        user_key_endpoint = f"{self._url}/v1/user/{username}/key"
+        response = requests.get(user_key_endpoint)
+        adoc_blob = response.json()["attestation_document"]
+        root_cert = self._root_cert or attest.download_root_cert()
+        attestation_doc = attest.parse_attestation(
+            base64.b64decode(adoc_blob), root_cert
+        )
+        if pcrs is not None:
+            attest.verify_pcrs(pcrs, attestation_doc)
+
+        user_data = attestation_doc.get("user_data")
+        user_data_dict = json.loads(user_data)
+        cape_key = user_data_dict.get("key")
+        if cape_key is None:
+            raise RuntimeError(
+                "Enclave response did not include a Cape key in attestation user data."
+            )
+        return base64.b64decode(cape_key)
+
+    async def _request_key_with_token(
         self,
         token: str,
-        key_path: pathlib.Path,
         pcrs: Optional[Dict[str, List[str]]] = None,
     ) -> bytes:
         key_endpoint = f"{self._url}/v1/key"
@@ -455,9 +513,7 @@ async def _request_key(
             raise RuntimeError(
                 "Enclave response did not include a Cape key in attestation user data."
             )
-        cape_key = base64.b64decode(cape_key)
-        await _persist_cape_key(cape_key, key_path)
-        return cape_key
+        return base64.b64decode(cape_key)
 
 
 class _EnclaveContext:

pycape/experimental/cli.py

@@ -15,6 +15,7 @@ async def deploy(
     deploy_path: Union[str, os.PathLike],
     url: Optional[str] = None,
     token_expiry: Optional[str] = None,
+    public: bool = False,
 ) -> fref.FunctionRef:
     """Deploy a directory or a zip file containing a Cape function declared in
     an app.py script.
@@ -49,6 +50,8 @@ async def deploy(
     deploy_path = pathlib.Path(deploy_path)
 
     cmd_deploy = f"cape deploy {deploy_path} -u {url} -o json"
+    if public:
+        cmd_deploy += " --public"
     out_deploy, err_deploy = _call_cape_cli(cmd_deploy)
     err_deploy = err_deploy.decode()
     out_deploy = out_deploy.decode()