From d2bd877d03a916a14c9b02384e39c474be7f7539 Mon Sep 17 00:00:00 2001 From: Jakub Hejhal Date: Sat, 22 Jul 2023 00:49:00 +0200 Subject: [PATCH 1/6] cleanup and refactor encrypt.py --- encrypt.py | 171 ++++++++++++++++++++++++++++------------------------- 1 file changed, 92 insertions(+), 79 deletions(-) diff --git a/encrypt.py b/encrypt.py index bbafa66..97cf2e8 100644 --- a/encrypt.py +++ b/encrypt.py @@ -1,13 +1,16 @@ +import argparse import base64 import json +from pathlib import Path -from nucypher_core.ferveo import DkgPublicKey - +from cryptography.fernet import Fernet +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 ###################### # Boring setup stuff # @@ -17,90 +20,100 @@ 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 ##################### -from cryptography.fernet import Fernet -from eth_utils import keccak -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, clear_text: bytes) -> bytes: f = Fernet(secret) - capsule = f.encrypt(definitely_tony) + capsule = f.encrypt(clear_text) 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 - ), -) -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}" -) - -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) - -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}') + +def main(args): + file_path = Path(args.file_path) + + with open(file_path, "rb") as f: + clear_text = f.read() + + plaintext_of_sym_key = keygen() + + # TODO: what do we do with this secret hash? + # It's not used anywhere atm. + secret_hash = keccak(plaintext_of_sym_key) + bulk_ciphertext = encapsulate(plaintext_of_sym_key, clear_text) + + print("--------- Threshold Encryption ---------") + + coordinator_agent = CoordinatorAgent( + provider_uri=args.coordinator_provider_uri, + registry=InMemoryContractRegistry.from_latest_publication(network=args.coordinator_network), + ) + ritual_id = args.ritual_id + 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": args.chain, + "method": "eth_getBalance", + "parameters": [args.eth_address, "latest"], + "returnValueTest": {"comparator": "==", "value": 0}, + }, + } + + ciphertext_of_sym_key = enrico.encrypt_for_dkg( + plaintext=plaintext_of_sym_key, conditions=eth_balance_condition + ) + + 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 == clear_text + + ################## + + tmk_json = json.dumps(tmk) + + filename = args.output_file + with open(filename, "w") as file: + data = tmk_json + file.write(data) + print(f"Wrote {len(data)} bytes to {filename}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Encrypt file using threshold encryption.") + parser.add_argument("--file_path", type=str, default="manzana.mp3", help="Path to the file to be encrypted") + parser.add_argument("--ritual_id", type=int, help="Ritual ID obtained from a side channel", default=15) + parser.add_argument("--coordinator_provider_uri", type=str, help="URI of the coordinator provider", required=True) + parser.add_argument("--coordinator_network", type=str, default="mumbai", help="Network for the coordinator", choices=["mumbai", "rinkeby", "mainnet", "goerli", "ropsten", "kovan"]) + parser.add_argument("--chain", type=int, help="Ethereum chain ID", default=80001) + parser.add_argument("--eth_address", type=str, help="Ethereum address for balance check", default="0x210eeAC07542F815ebB6FD6689637D8cA2689392") + parser.add_argument("--output-file", type=str, default="tony.tmk", help="Output file for encrypted data") + + args = parser.parse_args() + main(args) From 9bc460a1c02dccfc8de9ca866f01ba09b387df07 Mon Sep 17 00:00:00 2001 From: Jakub Hejhal Date: Sat, 22 Jul 2023 01:22:51 +0200 Subject: [PATCH 2/6] add TMK format type spec --- encrypt.py | 3 ++- revealer_bot/decryption_action.py | 4 +++- revealer_bot/types.py | 8 ++++++++ 3 files changed, 13 insertions(+), 2 deletions(-) create mode 100644 revealer_bot/types.py diff --git a/encrypt.py b/encrypt.py index 97cf2e8..224d619 100644 --- a/encrypt.py +++ b/encrypt.py @@ -11,6 +11,7 @@ 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 # @@ -79,7 +80,7 @@ def main(args): plaintext=plaintext_of_sym_key, conditions=eth_balance_condition ) - tmk = { + 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, diff --git a/revealer_bot/decryption_action.py b/revealer_bot/decryption_action.py index d4ce878..a754ee4 100644 --- a/revealer_bot/decryption_action.py +++ b/revealer_bot/decryption_action.py @@ -2,6 +2,7 @@ import io import json import os +from typing import cast from cryptography.fernet import Fernet import discord @@ -10,6 +11,7 @@ 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: await message.reply("wrong file type or something") return diff --git a/revealer_bot/types.py b/revealer_bot/types.py new file mode 100644 index 0000000..b89a194 --- /dev/null +++ b/revealer_bot/types.py @@ -0,0 +1,8 @@ +from typing import TypedDict +from nucypher.policy.conditions.lingo import ConditionLingo, Lingo + +class TMK(TypedDict): + bulk_ciphertext: str # encoded as base64 + encrypted_sym_key: str # encoded as hex + conditions: Lingo + filename: str From d74e771002bc4f0a35f617d9c0b71de28ded68dd Mon Sep 17 00:00:00 2001 From: Jakub Hejhal Date: Sun, 23 Jul 2023 18:09:02 +0200 Subject: [PATCH 3/6] use click instead of argparse --- encrypt.py | 87 +++++++++++++++++++++++++----------------------------- 1 file changed, 41 insertions(+), 46 deletions(-) diff --git a/encrypt.py b/encrypt.py index 6b24b0a..3f21b5c 100644 --- a/encrypt.py +++ b/encrypt.py @@ -1,8 +1,8 @@ -import argparse import base64 import json from pathlib import Path +import click from cryptography.fernet import Fernet from eth_utils import keccak # type: ignore from nucypher.blockchain.eth.agents import CoordinatorAgent @@ -38,8 +38,39 @@ def encapsulate(secret: bytes, clear_text: bytes) -> bytes: return capsule -def main(args): - file_path = Path(args.file_path) +@click.command() +@click.option( + "--file-path", type=str, default="manzana.mp3", help="Path to the file to be encrypted" +) +@click.option("--ritual-id", type=int, help="Ritual ID obtained from a side channel", default=15) +@click.option( + "--coordinator-provider-uri", type=str, help="URI of the coordinator provider", required=True +) +@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", +) +@click.option("--output-file", type=str, default="tony.tmk", help="Output file for encrypted data") +def main( + file_path, + ritual_id, + coordinator_provider_uri, + coordinator_network, + chain, + eth_address, + output_file, +): + file_path = Path(file_path) with open(file_path, "rb") as f: clear_text = f.read() @@ -52,12 +83,9 @@ def main(args): print("--------- Threshold Encryption ---------") coordinator_agent = CoordinatorAgent( - provider_uri=args.coordinator_provider_uri, - registry=InMemoryContractRegistry.from_latest_publication( - network=args.coordinator_network - ), + provider_uri=coordinator_provider_uri, + registry=InMemoryContractRegistry.from_latest_publication(network=coordinator_network), ) - ritual_id = args.ritual_id ritual = coordinator_agent.get_ritual(ritual_id) enrico = Enrico(encrypting_key=DkgPublicKey.from_bytes(bytes(ritual.public_key))) @@ -70,9 +98,9 @@ def main(args): eth_balance_condition: Lingo = { "version": ConditionLingo.VERSION, "condition": { - "chain": args.chain, + "chain": chain, "method": "eth_getBalance", - "parameters": [args.eth_address, "latest"], + "parameters": [eth_address, "latest"], "returnValueTest": {"comparator": "==", "value": 0}, }, } @@ -100,46 +128,13 @@ def main(args): tmk_json = json.dumps(tmk) - filename = args.output_file - with open(filename, "w") as file: + with open(output_file, "w") as file: data = tmk_json file.write(data) - print(f"Wrote {len(data)} bytes to {filename}") + print(f"Wrote {len(data)} bytes to {output_file}") print("Keccak hash of plaintext sym key: ", secret_hash.hex()) if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Encrypt file using threshold encryption.") - parser.add_argument( - "--file_path", type=str, default="manzana.mp3", help="Path to the file to be encrypted" - ) - parser.add_argument( - "--ritual_id", type=int, help="Ritual ID obtained from a side channel", default=15 - ) - parser.add_argument( - "--coordinator_provider_uri", - type=str, - help="URI of the coordinator provider", - required=True, - ) - parser.add_argument( - "--coordinator_network", - type=str, - default="mumbai", - help="Network for the coordinator", - choices=["mumbai", "rinkeby", "mainnet", "goerli", "ropsten", "kovan"], - ) - parser.add_argument("--chain", type=int, help="Ethereum chain ID", default=80001) - parser.add_argument( - "--eth_address", - type=str, - help="Ethereum address for balance check", - default="0x210eeAC07542F815ebB6FD6689637D8cA2689392", - ) - parser.add_argument( - "--output-file", type=str, default="tony.tmk", help="Output file for encrypted data" - ) - - args = parser.parse_args() - main(args) + main() From 6e971854a0cd0e3fc637152cd16fe39adce1846b Mon Sep 17 00:00:00 2001 From: Jakub Hejhal Date: Sun, 23 Jul 2023 18:30:47 +0200 Subject: [PATCH 4/6] add type annotations and eth-minimum-balance cmd arg --- encrypt.py | 37 ++++++++++++++++++++++++++----------- revealer_bot/types.py | 8 +++++--- 2 files changed, 31 insertions(+), 14 deletions(-) diff --git a/encrypt.py b/encrypt.py index 3f21b5c..429f398 100644 --- a/encrypt.py +++ b/encrypt.py @@ -40,9 +40,15 @@ def encapsulate(secret: bytes, clear_text: bytes) -> bytes: @click.command() @click.option( - "--file-path", type=str, default="manzana.mp3", help="Path to the file to be encrypted" + "--input-file", type=str, default="manzana.mp3", help="Path to the file to be encrypted" +) +@click.option( + "--ritual-id", + type=int, + help="Ritual ID obtained from a side channel", + default=15, + show_default=True, ) -@click.option("--ritual-id", type=int, help="Ritual ID obtained from a side channel", default=15) @click.option( "--coordinator-provider-uri", type=str, help="URI of the coordinator provider", required=True ) @@ -59,18 +65,27 @@ def encapsulate(secret: bytes, clear_text: bytes) -> bytes: 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( - file_path, - ritual_id, - coordinator_provider_uri, - coordinator_network, - chain, - eth_address, - output_file, + 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(file_path) + file_path = Path(input_file) with open(file_path, "rb") as f: clear_text = f.read() @@ -101,7 +116,7 @@ def main( "chain": chain, "method": "eth_getBalance", "parameters": [eth_address, "latest"], - "returnValueTest": {"comparator": "==", "value": 0}, + "returnValueTest": {"comparator": ">=", "value": eth_minimum_balance}, }, } diff --git a/revealer_bot/types.py b/revealer_bot/types.py index b89a194..c635b60 100644 --- a/revealer_bot/types.py +++ b/revealer_bot/types.py @@ -1,8 +1,10 @@ from typing import TypedDict -from nucypher.policy.conditions.lingo import ConditionLingo, Lingo + +from nucypher.policy.conditions.lingo import Lingo + class TMK(TypedDict): - bulk_ciphertext: str # encoded as base64 - encrypted_sym_key: str # encoded as hex + bulk_ciphertext: str # encoded as base64 + encrypted_sym_key: str # encoded as hex conditions: Lingo filename: str From 0fab1b2e10887e92a9855abff5ed87e446dd9dbd Mon Sep 17 00:00:00 2001 From: Jakub Hejhal Date: Sun, 23 Jul 2023 18:35:14 +0200 Subject: [PATCH 5/6] rename clear_text -> cleartext --- encrypt.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/encrypt.py b/encrypt.py index 429f398..1075214 100644 --- a/encrypt.py +++ b/encrypt.py @@ -32,9 +32,9 @@ def keygen() -> bytes: return _secret -def encapsulate(secret: bytes, clear_text: bytes) -> bytes: +def encapsulate(secret: bytes, cleartext: bytes) -> bytes: f = Fernet(secret) - capsule = f.encrypt(clear_text) + capsule = f.encrypt(cleartext) return capsule @@ -88,12 +88,12 @@ def main( file_path = Path(input_file) with open(file_path, "rb") as f: - clear_text = f.read() + cleartext = f.read() plaintext_of_sym_key = keygen() secret_hash = keccak(plaintext_of_sym_key) - bulk_ciphertext = encapsulate(plaintext_of_sym_key, clear_text) + bulk_ciphertext = encapsulate(plaintext_of_sym_key, cleartext) print("--------- Threshold Encryption ---------") @@ -137,7 +137,7 @@ def main( f = Fernet(plaintext_of_sym_key) hopefully_decrypted = f.decrypt(bulk_ciphertext) - assert hopefully_decrypted == clear_text + assert hopefully_decrypted == cleartext ################## From f4e8b0bfe2bd60479a1e8f807d556e66a62b7322 Mon Sep 17 00:00:00 2001 From: Jakub Hejhal Date: Mon, 24 Jul 2023 00:13:36 +0200 Subject: [PATCH 6/6] use binary format for TMK, include metadata in ciphertext --- encrypt.py | 55 ++++++++++++---------------- revealer_bot/decryption_action.py | 30 +++++++--------- revealer_bot/tmk.py | 60 +++++++++++++++++++++++++++++++ revealer_bot/types.py | 10 ------ 4 files changed, 96 insertions(+), 59 deletions(-) create mode 100644 revealer_bot/tmk.py delete mode 100644 revealer_bot/types.py diff --git a/encrypt.py b/encrypt.py index 1075214..7149801 100644 --- a/encrypt.py +++ b/encrypt.py @@ -1,5 +1,3 @@ -import base64 -import json from pathlib import Path import click @@ -12,7 +10,7 @@ from nucypher.utilities.logging import GlobalLoggerSettings from nucypher_core.ferveo import DkgPublicKey -from revealer_bot.types import TMK +from revealer_bot.tmk import TMK, Payload, decrypt, encapsulate ###################### # Boring setup stuff # @@ -32,12 +30,6 @@ def keygen() -> bytes: return _secret -def encapsulate(secret: bytes, cleartext: bytes) -> bytes: - f = Fernet(secret) - capsule = f.encrypt(cleartext) - return capsule - - @click.command() @click.option( "--input-file", type=str, default="manzana.mp3", help="Path to the file to be encrypted" @@ -88,12 +80,13 @@ def main( file_path = Path(input_file) with open(file_path, "rb") as f: - cleartext = f.read() + file_content = f.read() - plaintext_of_sym_key = keygen() + 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, cleartext) + bulk_ciphertext = encapsulate(plaintext_of_sym_key, payload.to_bytes()) print("--------- Threshold Encryption ---------") @@ -124,31 +117,29 @@ def main( 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) + tmk = TMK( + bulk_ciphertext=bulk_ciphertext, + encrypted_sym_key=bytes(ciphertext_of_sym_key), + conditions=eth_balance_condition, + ) - with open(output_file, "w") as file: - data = tmk_json + 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__": diff --git a/revealer_bot/decryption_action.py b/revealer_bot/decryption_action.py index 00fa479..c8d773e 100644 --- a/revealer_bot/decryption_action.py +++ b/revealer_bot/decryption_action.py @@ -1,15 +1,11 @@ -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 +from revealer_bot.tmk import TMK, Payload, decrypt async def decrypt_attached_tmk(message): @@ -23,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 = cast(TMK, 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)) diff --git a/revealer_bot/types.py b/revealer_bot/types.py deleted file mode 100644 index c635b60..0000000 --- a/revealer_bot/types.py +++ /dev/null @@ -1,10 +0,0 @@ -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