cleanup and refactor encrypt.py

encrypt.py

@@ -1,15 +1,19 @@
 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.types import TMK
+
 ######################
 # Boring setup stuff #
 ######################
@@ -18,91 +22,134 @@
 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):
+def encapsulate(secret: bytes, cleartext: bytes) -> bytes:
     f = Fernet(secret)
-    capsule = f.encrypt(definitely_tony)
+    capsule = f.encrypt(cleartext)
     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:
+        cleartext = f.read()
+
+    plaintext_of_sym_key = keygen()
+
+    secret_hash = keccak(plaintext_of_sym_key)
+    bulk_ciphertext = encapsulate(plaintext_of_sym_key, cleartext)
+
+    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": base64.b64encode(bytes(bulk_ciphertext)).decode(),  # Encrypted Tony
+        "encrypted_sym_key": bytes(ciphertext_of_sym_key).hex(),
+        "conditions": eth_balance_condition,
+        "filename": file_path.name,
+    }
+
+    ################
+    # Sanity check #
+    ################
+
+    f = Fernet(plaintext_of_sym_key)
+    hopefully_decrypted = f.decrypt(bulk_ciphertext)
+    assert hopefully_decrypted == cleartext
+
+    ##################
+
+    tmk_json = json.dumps(tmk)
+
+    with open(output_file, "w") as file:
+        data = tmk_json
+        file.write(data)
+        print(f"Wrote {len(data)} bytes to {output_file}")
+
+    print("Keccak hash of plaintext sym key: ", secret_hash.hex())
+
+
+if __name__ == "__main__":
+    main()

revealer_bot/decryption_action.py

@@ -1,13 +1,15 @@
 import base64
 import io
 import json
+from typing import cast
 
 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.types import TMK
 
 
 async def decrypt_attached_tmk(message):
@@ -22,7 +24,7 @@ async def decrypt_attached_tmk(message):
 
         try:
             json_str_repr_of_tmk = str(attachment_response.content, encoding="utf-8")
-            tmk_dict = json.loads(json_str_repr_of_tmk)
+            tmk_dict = cast(TMK, json.loads(json_str_repr_of_tmk))
         except json.JSONDecodeError:
             await message.reply("wrong file type or something")
             return

revealer_bot/types.py

@@ -0,0 +1,10 @@
+from typing import TypedDict
+
+from nucypher.policy.conditions.lingo import Lingo
+
+
+class TMK(TypedDict):
+    bulk_ciphertext: str  # encoded as base64
+    encrypted_sym_key: str  # encoded as hex
+    conditions: Lingo
+    filename: str