diff --git a/encrypt.py b/encrypt.py index a2cc805..7149801 100644 --- a/encrypt.py +++ b/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 = "" -network = "lynx" - -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() diff --git a/revealer_bot/decryption_action.py b/revealer_bot/decryption_action.py index 033b7e7..c8d773e 100644 --- a/revealer_bot/decryption_action.py +++ b/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"]), ) diff --git a/revealer_bot/tmk.py b/revealer_bot/tmk.py new file mode 100644 index 0000000..0436d38 --- /dev/null +++ b/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))