Cleanup encrypt, decrypt + new TKM format with metadata included in the ciphertext

encrypt.py

@@ -1,15 +1,17 @@
-import base64
-import json
+from pathlib import Path
 
+import click
 from cryptography.fernet import Fernet
-from eth_utils import keccak
+from eth_utils import keccak  # type: ignore
 from nucypher.blockchain.eth.agents import CoordinatorAgent
 from nucypher.blockchain.eth.registry import InMemoryContractRegistry
 from nucypher.characters.lawful import Enrico
-from nucypher.policy.conditions.lingo import ConditionLingo
+from nucypher.policy.conditions.lingo import ConditionLingo, Lingo
 from nucypher.utilities.logging import GlobalLoggerSettings
 from nucypher_core.ferveo import DkgPublicKey
 
+from revealer_bot.tmk import TMK, Payload, decrypt, encapsulate
+
 ######################
 # Boring setup stuff #
 ######################
@@ -18,91 +20,127 @@
 GlobalLoggerSettings.set_log_level(log_level_name=LOG_LEVEL)
 GlobalLoggerSettings.start_console_logging()
 
-staking_provider_uri = "<your staking provider uri>"
-network = "lynx"
-
-coordinator_provider_uri = "<your coordinator provider uri>"
-coordinator_network = "mumbai"
-
 #####################
 # Scully the Symmet
 #####################
 
 
-def keygen():
+def keygen() -> bytes:
     _secret = Fernet.generate_key()
     return _secret
 
 
-with open("manzana.mp3", "rb") as tony:
-    definitely_tony = tony.read()
-
-
-def encapsulate(secret):
-    f = Fernet(secret)
-    capsule = f.encrypt(definitely_tony)
-    return capsule
-
-
-plaintext_of_sym_key = keygen()
-secret_hash = keccak(plaintext_of_sym_key)
-bulk_ciphertext = encapsulate(plaintext_of_sym_key)
-
-###############
-# Enrico
-###############
-
-print("--------- Threshold Encryption ---------")
-
-coordinator_agent = CoordinatorAgent(
-    provider_uri=coordinator_provider_uri,
-    registry=InMemoryContractRegistry.from_latest_publication(network=coordinator_network),
+@click.command()
+@click.option(
+    "--input-file", type=str, default="manzana.mp3", help="Path to the file to be encrypted"
 )
-ritual_id = 15  # got this from a side channel
-ritual = coordinator_agent.get_ritual(ritual_id)
-enrico = Enrico(encrypting_key=DkgPublicKey.from_bytes(bytes(ritual.public_key)))
-
-print(
-    f"Fetched DKG public key {bytes(enrico.policy_pubkey).hex()} "
-    f"for ritual #{ritual_id} "
-    f"from Coordinator {coordinator_agent.contract.address}"
+@click.option(
+    "--ritual-id",
+    type=int,
+    help="Ritual ID obtained from a side channel",
+    default=15,
+    show_default=True,
 )
-
-eth_balance_condition = {
-    "version": ConditionLingo.VERSION,
-    "condition": {
-        "chain": 80001,
-        "method": "eth_getBalance",
-        "parameters": ["0x210eeAC07542F815ebB6FD6689637D8cA2689392", "latest"],
-        "returnValueTest": {"comparator": "==", "value": 0},
-    },
-}
-
-ciphertext_of_sym_key = enrico.encrypt_for_dkg(
-    plaintext=plaintext_of_sym_key, conditions=eth_balance_condition
+@click.option(
+    "--coordinator-provider-uri", type=str, help="URI of the coordinator provider", required=True
 )
-
-tmk = {
-    "bulk_ciphertext": base64.b64encode(bytes(bulk_ciphertext)).decode(),  # Encrypted Tony
-    "encrypted_sym_key": bytes(ciphertext_of_sym_key).hex(),
-    "conditions": eth_balance_condition,
-    "filename": "manzana.mp3",
-}
-
-################
-# Sanity check #
-################
-
-f = Fernet(plaintext_of_sym_key)
-hopefully_tony = f.decrypt(bulk_ciphertext)
-assert hopefully_tony == definitely_tony
-
-##################
-
-tmk_json = json.dumps(tmk)
-
-filename = "tony.tmk"
-with open(filename, "w") as file:
-    data = tmk_json
-    file.write(data)
-    print(f"Wrote {len(data)} bytes to {filename}")
+@click.option(
+    "--coordinator-network",
+    default="mumbai",
+    help="Network for the coordinator",
+    show_default=True,
+    type=click.Choice(["mumbai", "rinkeby", "mainnet", "goerli", "ropsten", "kovan"]),
+)
+@click.option("--chain", type=int, help="Ethereum chain ID", default=80001, show_default=True)
+@click.option(
+    "--eth-address",
+    type=str,
+    help="Ethereum address for balance check",
+    default="0x210eeAC07542F815ebB6FD6689637D8cA2689392",
+    show_default=True,
+)
+@click.option(
+    "--eth-minimum-balance",
+    type=float,
+    help="Ethereum minimum balance condition",
+    default=0,
+    show_default=True,
+)
+@click.option("--output-file", type=str, default="tony.tmk", help="Output file for encrypted data")
+def main(
+    input_file: str,
+    ritual_id: int,
+    coordinator_provider_uri: str,
+    coordinator_network: str,
+    chain: int,
+    eth_address: str,
+    eth_minimum_balance: float,
+    output_file: str,
+):
+    file_path = Path(input_file)
+
+    with open(file_path, "rb") as f:
+        file_content = f.read()
+
+    payload = Payload(file_content=file_content, metadata={"filename": file_path.name})
+
+    plaintext_of_sym_key = keygen()
+    secret_hash = keccak(plaintext_of_sym_key)
+    bulk_ciphertext = encapsulate(plaintext_of_sym_key, payload.to_bytes())
+
+    print("--------- Threshold Encryption ---------")
+
+    coordinator_agent = CoordinatorAgent(
+        provider_uri=coordinator_provider_uri,
+        registry=InMemoryContractRegistry.from_latest_publication(network=coordinator_network),
+    )
+    ritual = coordinator_agent.get_ritual(ritual_id)
+    enrico = Enrico(encrypting_key=DkgPublicKey.from_bytes(bytes(ritual.public_key)))
+
+    print(
+        f"Fetched DKG public key {bytes(enrico.policy_pubkey).hex()} "  # type: ignore
+        f"for ritual #{ritual_id} "
+        f"from Coordinator {coordinator_agent.contract.address}"
+    )
+
+    eth_balance_condition: Lingo = {
+        "version": ConditionLingo.VERSION,
+        "condition": {
+            "chain": chain,
+            "method": "eth_getBalance",
+            "parameters": [eth_address, "latest"],
+            "returnValueTest": {"comparator": ">=", "value": eth_minimum_balance},
+        },
+    }
+
+    ciphertext_of_sym_key = enrico.encrypt_for_dkg(
+        plaintext=plaintext_of_sym_key, conditions=eth_balance_condition
+    )
+
+    tmk = TMK(
+        bulk_ciphertext=bulk_ciphertext,
+        encrypted_sym_key=bytes(ciphertext_of_sym_key),
+        conditions=eth_balance_condition,
+    )
+
+    with open(output_file, "wb") as file:
+        data = tmk.to_bytes()
+        file.write(data)
+        print(f"Wrote {len(data)} bytes to {output_file}")
+
+    print("Keccak hash of plaintext sym key: ", secret_hash.hex())
+    ################
+    # Sanity check #
+    ################
+
+    hopefully_tmk = TMK.from_bytes(data)
+    hopefully_cleartext = decrypt(
+        ciphertext=hopefully_tmk.bulk_ciphertext, plaintext_of_symkey=plaintext_of_sym_key
+    )
+    hopefully_payload = Payload.from_bytes(hopefully_cleartext)
+    assert hopefully_payload.metadata["filename"] == payload.metadata["filename"]
+    assert hopefully_payload.file_content == payload.file_content
+
+
+if __name__ == "__main__":
+    main()

revealer_bot/decryption_action.py

@@ -1,13 +1,11 @@
-import base64
 import io
-import json
 
 import discord
 import requests
-from cryptography.fernet import Fernet
 from nucypher_core import ferveo
 
 from revealer_bot.bob_and_other_networky_things import bob
+from revealer_bot.tmk import TMK, Payload, decrypt
 
 
 async def decrypt_attached_tmk(message):
@@ -21,29 +19,29 @@ async def decrypt_attached_tmk(message):
         attachment_response = requests.get(url)
 
         try:
-            json_str_repr_of_tmk = str(attachment_response.content, encoding="utf-8")
-            tmk_dict = json.loads(json_str_repr_of_tmk)
-        except json.JSONDecodeError:
+            tmk = TMK.from_bytes(attachment_response.content)
+        except Exception:
             await message.reply("wrong file type or something")
             return
-        print("--------- Threshold Decryption ---------")
 
-        ciphertext_of_sym_key = bytes.fromhex(tmk_dict["encrypted_sym_key"])
-        ciphertext_to_decrypt_with_threshold = ferveo.Ciphertext.from_bytes(ciphertext_of_sym_key)
+        print("--------- Threshold Decryption ---------")
+        ciphertext_to_decrypt_with_threshold = ferveo.Ciphertext.from_bytes(tmk.encrypted_sym_key)
 
         ######### BAAAAAAHB ########
 
         plaintext_of_symkey = bob.threshold_decrypt(
             ritual_id=15,  # Cuz 15
             ciphertext=ciphertext_to_decrypt_with_threshold,
-            conditions=tmk_dict["conditions"],
+            conditions=tmk.conditions,
         )
 
-        f = Fernet(bytes(plaintext_of_symkey))
-        bulk_ciphertext = base64.b64decode(tmk_dict["bulk_ciphertext"])
-        hopefully_tony = f.decrypt(bulk_ciphertext)
+        cleartext = decrypt(
+            ciphertext=tmk.bulk_ciphertext, plaintext_of_symkey=plaintext_of_symkey
+        )
+        payload = Payload.from_bytes(cleartext)
+        filelike = io.BytesIO(payload.file_content)
 
-        filelike = io.BytesIO(hopefully_tony)
         await message.reply(
-            "Here's what I found.", file=discord.File(filelike, filename=tmk_dict["filename"])
+            "Here's what I found.",
+            file=discord.File(filelike, filename=payload.metadata["filename"]),
         )

revealer_bot/tmk.py

@@ -0,0 +1,60 @@
+from dataclasses import asdict, dataclass
+from typing import Any, NewType, cast
+
+import msgpack
+from cryptography.fernet import Fernet
+from nucypher.policy.conditions.lingo import Lingo
+
+ClearText = NewType("ClearText", bytes)
+BulkCipherText = NewType("BulkCipherText", bytes)
+
+
+@dataclass
+class Payload:
+    file_content: bytes
+    metadata: dict[str, Any]
+
+    def to_bytes(self) -> ClearText:
+        # Serialize the dictionary to bytes using MessagePack
+        data_dict = asdict(self)
+        serialized_data: ClearText = msgpack.packb(data_dict)  # type: ignore
+        return serialized_data
+
+    @classmethod
+    def from_bytes(cls, serialized_data: ClearText) -> "Payload":
+        # Deserialize the bytes to a dictionary using MessagePack
+        data_dict: dict[str, Any] = msgpack.unpackb(serialized_data)
+        return cls(**data_dict)
+
+
+def decrypt(ciphertext: BulkCipherText, plaintext_of_symkey: bytes) -> ClearText:
+    f = Fernet(plaintext_of_symkey)
+    cleartext: ClearText = cast(ClearText, f.decrypt(ciphertext))
+    return cleartext
+
+
+def encapsulate(plaintext_of_symkey: bytes, cleartext: ClearText) -> BulkCipherText:
+    f = Fernet(plaintext_of_symkey)
+    capsule = cast(BulkCipherText, f.encrypt(cleartext))
+    return capsule
+
+
+@dataclass
+class TMK:
+    bulk_ciphertext: BulkCipherText
+    encrypted_sym_key: bytes
+    conditions: Lingo
+
+    def __init__(
+        self, bulk_ciphertext: BulkCipherText, encrypted_sym_key: bytes, conditions: Lingo
+    ) -> None:
+        self.bulk_ciphertext = bulk_ciphertext
+        self.encrypted_sym_key = encrypted_sym_key
+        self.conditions = conditions
+
+    def to_bytes(self) -> bytes:
+        return msgpack.packb(asdict(self))  # type: ignore
+
+    @classmethod
+    def from_bytes(cls, serialized_data: bytes) -> "TMK":
+        return cls(**msgpack.unpackb(serialized_data))