|
| 1 | +import struct |
| 2 | +import time |
| 3 | +from typing import Optional, Union |
| 4 | + |
| 5 | +from bittensor_commit_reveal import ( |
| 6 | + encrypt as _btr_encrypt, |
| 7 | + decrypt as _btr_decrypt, |
| 8 | + get_latest_round, |
| 9 | +) |
| 10 | + |
| 11 | +TLE_ENCRYPTED_DATA_SUFFIX = b"AES_GCM_" |
| 12 | + |
| 13 | + |
| 14 | +def encrypt( |
| 15 | + data: Union[bytes, str], n_blocks: int, block_time: Union[int, float] = 12.0 |
| 16 | +) -> tuple[bytes, int]: |
| 17 | + """Encrypts data using TimeLock Encryption |
| 18 | +
|
| 19 | + Arguments: |
| 20 | + data: Any bytes data to be encrypted. |
| 21 | + n_blocks: Number of blocks to encrypt. |
| 22 | + block_time: Time in seconds for each block. Default is `12.0` seconds. |
| 23 | +
|
| 24 | + Returns: |
| 25 | + tuple: A tuple containing the encrypted data and reveal TimeLock reveal round. |
| 26 | +
|
| 27 | + Raises: |
| 28 | + PyValueError: If failed to encrypt data. |
| 29 | +
|
| 30 | + Usage: |
| 31 | + data = "From Cortex to Bittensor" |
| 32 | +
|
| 33 | + # default usage |
| 34 | + encrypted_data, reveal_round = encrypt(data, 10) |
| 35 | +
|
| 36 | + # passing block_time for fast-blocks node |
| 37 | + encrypted_data, reveal_round = encrypt(data, 15, block_time=0.25) |
| 38 | +
|
| 39 | + encrypted_data, reveal_round = encrypt(data, 5) |
| 40 | +
|
| 41 | +
|
| 42 | + Note: |
| 43 | + For using this function with fast-blocks node you need to set block_time to 0.25 seconds. |
| 44 | + data, round = encrypt(data, n_blocks, block_time=0.25) |
| 45 | + """ |
| 46 | + if isinstance(data, str): |
| 47 | + data = data.encode() |
| 48 | + return _btr_encrypt(data, n_blocks, block_time) |
| 49 | + |
| 50 | + |
| 51 | +def decrypt( |
| 52 | + encrypted_data: bytes, no_errors: bool = True, return_str: bool = False |
| 53 | +) -> Optional[Union[bytes, str]]: |
| 54 | + """Decrypts encrypted data using TimeLock Decryption |
| 55 | +
|
| 56 | + Arguments: |
| 57 | + encrypted_data: Encrypted data to be decrypted. |
| 58 | + no_errors: If True, no errors will be raised during decryption. |
| 59 | + return_str: convert decrypted data to string if `True`. Default is `False`. |
| 60 | +
|
| 61 | + Returns: |
| 62 | + decrypted_data: Decrypted data, when reveled round is reached. |
| 63 | +
|
| 64 | + Usage: |
| 65 | + # default usage |
| 66 | + decrypted_data = decrypt(encrypted_data) |
| 67 | +
|
| 68 | + # passing no_errors=False for raising errors during decryption |
| 69 | + decrypted_data = decrypt(encrypted_data, no_errors=False) |
| 70 | +
|
| 71 | + # passing return_str=True for returning decrypted data as string |
| 72 | + decrypted_data = decrypt(encrypted_data, return_str=True) |
| 73 | + """ |
| 74 | + result = _btr_decrypt(encrypted_data, no_errors) |
| 75 | + if result is None: |
| 76 | + return None |
| 77 | + if return_str: |
| 78 | + return result.decode() |
| 79 | + return result |
| 80 | + |
| 81 | + |
| 82 | +def wait_reveal_and_decrypt( |
| 83 | + encrypted_data: bytes, |
| 84 | + reveal_round: Optional[int] = None, |
| 85 | + no_errors: bool = True, |
| 86 | + return_str: bool = False, |
| 87 | +) -> bytes: |
| 88 | + """ |
| 89 | + Waits for reveal round and decrypts data using TimeLock Decryption. |
| 90 | +
|
| 91 | + Arguments: |
| 92 | + encrypted_data: Encrypted data to be decrypted. |
| 93 | + reveal_round: Reveal round to wait for. If None, will be parsed from encrypted data. |
| 94 | + no_errors: If True, no errors will be raised during decryption. |
| 95 | + return_str: convert decrypted data to string if `True`. Default is `False`. |
| 96 | +
|
| 97 | + Raises: |
| 98 | + struct.error: If failed to parse reveal round from encrypted data. |
| 99 | + TypeError: If reveal_round is None or wrong type. |
| 100 | + IndexError: If provided encrypted_data does not contain reveal round. |
| 101 | +
|
| 102 | + Returns: |
| 103 | + bytes: Decrypted data. |
| 104 | +
|
| 105 | + Usage: |
| 106 | + import bittensor as bt |
| 107 | + encrypted, reveal_round = bt.timelock.encrypt("Cortex is power", 3) |
| 108 | + """ |
| 109 | + if reveal_round is None: |
| 110 | + try: |
| 111 | + reveal_round = struct.unpack( |
| 112 | + "<Q", encrypted_data.split(TLE_ENCRYPTED_DATA_SUFFIX)[-1] |
| 113 | + )[0] |
| 114 | + except (struct.error, TypeError, IndexError): |
| 115 | + raise ValueError("Failed to parse reveal round from encrypted data.") |
| 116 | + |
| 117 | + while get_latest_round() <= reveal_round: |
| 118 | + # sleep Drand QuickNet period time (3 sec) |
| 119 | + time.sleep(3) |
| 120 | + |
| 121 | + return decrypt(encrypted_data, no_errors, return_str) |
| 122 | + |
| 123 | + |
| 124 | +__all__ = [ |
| 125 | + "decrypt", |
| 126 | + "encrypt", |
| 127 | + "get_latest_round", |
| 128 | + "wait_reveal_and_decrypt", |
| 129 | +] |
0 commit comments