From 7b8181b77a4f5fce43cca7a684e2378329b271cd Mon Sep 17 00:00:00 2001 From: Roman Date: Mon, 14 Apr 2025 14:53:50 -0700 Subject: [PATCH 1/5] add timelock module --- bittensor/core/timelock.py | 131 +++++++++++++++++++++++++++++++++++++ 1 file changed, 131 insertions(+) create mode 100644 bittensor/core/timelock.py diff --git a/bittensor/core/timelock.py b/bittensor/core/timelock.py new file mode 100644 index 0000000000..5c620ac895 --- /dev/null +++ b/bittensor/core/timelock.py @@ -0,0 +1,131 @@ +import struct +import time +from typing import Optional, Union + +from bittensor_commit_reveal import ( + encrypt as _btr_encrypt, + decrypt as _btr_decrypt, + get_latest_round, + get_reveal_round_signature, +) + +TLE_ENCRYPTED_DATA_SUFFIX = b"AES_GCM_" + + +def encrypt( + data: Union[bytes, str], n_blocks: int, block_time: Union[int, float] = 12.0 +) -> tuple[bytes, int]: + """Encrypts data using TimeLock Encryption + + Arguments: + data: Any bytes data to be encrypted. + n_blocks: Number of blocks to encrypt. + block_time: Time in seconds for each block. Default is `12.0` seconds. + + Returns: + tuple: A tuple containing the encrypted data and reveal TimeLock reveal round. + + Raises: + PyValueError: If failed to encrypt data. + + Usage: + data = "From Cortex to Bittensor" + + # default usage + encrypted_data, reveal_round = encrypt(data, 10) + + # passing block_time for fast-blocks node + encrypted_data, reveal_round = encrypt(data, 15, block_time=0.25) + + encrypted_data, reveal_round = encrypt(data, 5) + + + Note: + For using this function with fast-blocks node you need to set block_time to 0.25 seconds. + data, round = encrypt(data, n_blocks, block_time=0.25) + """ + if isinstance(data, str): + data = data.encode() + return _btr_encrypt(data, n_blocks, block_time) + + +def decrypt( + encrypted_data: bytes, no_errors: bool = True, return_str: bool = False +) -> Optional[Union[bytes, str]]: + """Decrypts encrypted data using TimeLock Decryption + + Arguments: + encrypted_data: Encrypted data to be decrypted. + no_errors: If True, no errors will be raised during decryption. + return_str: convert decrypted data to string if `True`. Default is `False`. + + Returns: + decrypted_data: Decrypted data, when reveled round is reached. + + Usage: + # default usage + decrypted_data = decrypt(encrypted_data) + + # passing no_errors=False for raising errors during decryption + decrypted_data = decrypt(encrypted_data, no_errors=False) + + # passing return_str=True for returning decrypted data as string + decrypted_data = decrypt(encrypted_data, return_str=True) + """ + result = _btr_decrypt(encrypted_data, no_errors) + if result is None: + return None + if return_str: + return result.decode() + return result + + +def wait_reveal_and_decrypt( + encrypted_data: bytes, + reveal_round: Optional[int] = None, + no_errors: bool = True, + return_str: bool = False, +) -> bytes: + """ + Waits for reveal round and decrypts data using TimeLock Decryption. + + Arguments: + encrypted_data: Encrypted data to be decrypted. + reveal_round: Reveal round to wait for. If None, will be parsed from encrypted data. + no_errors: If True, no errors will be raised during decryption. + return_str: convert decrypted data to string if `True`. Default is `False`. + + Raises: + struct.error: If failed to parse reveal round from encrypted data. + TypeError: If reveal_round is None or wrong type. + IndexError: If provided encrypted_data does not contain reveal round. + + Returns: + bytes: Decrypted data. + + Usage: + import bittensor as bt + encrypted, reveal_round = bt.timelock.encrypt("Cortex is power", 3) + """ + if reveal_round is None: + try: + reveal_round = struct.unpack( + " Date: Mon, 14 Apr 2025 14:54:07 -0700 Subject: [PATCH 2/5] add import to main level --- bittensor/utils/easy_imports.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bittensor/utils/easy_imports.py b/bittensor/utils/easy_imports.py index 59ebeda7ba..cc81efe3d2 100644 --- a/bittensor/utils/easy_imports.py +++ b/bittensor/utils/easy_imports.py @@ -27,7 +27,7 @@ ) from bittensor_wallet.wallet import display_mnemonic_msg, Wallet # noqa: F401 -from bittensor.core import settings +from bittensor.core import settings, timelock # noqa: F401 from bittensor.core.async_subtensor import AsyncSubtensor from bittensor.core.axon import Axon from bittensor.core.chain_data import ( # noqa: F401 From 2a8edb029253d0779716efae806970cb68d0f124 Mon Sep 17 00:00:00 2001 From: Roman Date: Mon, 14 Apr 2025 14:54:19 -0700 Subject: [PATCH 3/5] add integration tests --- tests/integration_tests/test_timelock.py | 83 ++++++++++++++++++++++++ 1 file changed, 83 insertions(+) create mode 100644 tests/integration_tests/test_timelock.py diff --git a/tests/integration_tests/test_timelock.py b/tests/integration_tests/test_timelock.py new file mode 100644 index 0000000000..33e31db782 --- /dev/null +++ b/tests/integration_tests/test_timelock.py @@ -0,0 +1,83 @@ +import struct +import time + +import pytest + +from bittensor.core import timelock + + +def test_encrypt_returns_valid_tuple(): + """Test that encrypt() returns a (bytes, int) tuple.""" + encrypted, reveal_round = timelock.encrypt("Bittensor", n_blocks=1) + assert isinstance(encrypted, bytes) + assert isinstance(reveal_round, int) + assert reveal_round > 0 + + +def test_encrypt_with_fast_block_time(): + """Test encrypt() with fast-blocks mode (block_time = 0.25s).""" + encrypted, reveal_round = timelock.encrypt("Fast mode", 5, block_time=0.25) + assert isinstance(encrypted, bytes) + assert isinstance(reveal_round, int) + + +def test_decrypt_returns_bytes_or_none(): + """Test that decrypt() returns bytes after reveal round, or None before.""" + data = b"Decode me" + encrypted, reveal_round = timelock.encrypt(data, 1) + + current_round = timelock.get_latest_round() + if current_round < reveal_round: + decrypted = timelock.decrypt(encrypted) + assert decrypted is None + else: + decrypted = timelock.decrypt(encrypted) + assert decrypted == data + + +def test_decrypt_raises_if_no_errors_false_and_invalid_data(): + """Test that decrypt() raises an error on invalid data when no_errors=False.""" + with pytest.raises(Exception): + timelock.decrypt(b"corrupt data", no_errors=False) + + +def test_decrypt_with_return_str(): + """Test decrypt() with return_str=True returns a string.""" + plaintext = "Stringified!" + encrypted, _ = timelock.encrypt(plaintext, 1, block_time=0.25) + result = timelock.decrypt(encrypted, no_errors=True, return_str=True) + if result is not None: + assert isinstance(result, str) + + +def test_get_latest_round_is_monotonic(): + """Test that get_latest_round() is monotonic over time.""" + r1 = timelock.get_latest_round() + time.sleep(3) + r2 = timelock.get_latest_round() + assert r2 >= r1 + + +def test_wait_reveal_and_decrypt_auto_round(): + """Test wait_reveal_and_decrypt() without explicit reveal_round.""" + msg = "Reveal and decrypt test" + encrypted, _ = timelock.encrypt(msg, 1) + result = timelock.wait_reveal_and_decrypt(encrypted, return_str=True) + assert result == msg + + +def test_wait_reveal_and_decrypt_manual_round(): + """Test wait_reveal_and_decrypt() with explicit reveal_round.""" + msg = "Manual round decryption" + encrypted, reveal_round = timelock.encrypt(msg, 1) + result = timelock.wait_reveal_and_decrypt(encrypted, reveal_round, return_str=True) + assert result == msg + + +def test_unpack_reveal_round_struct(): + """Test that reveal_round can be extracted from encrypted data.""" + encrypted, reveal_round = timelock.encrypt("parse test", 1) + parsed = struct.unpack( + " Date: Mon, 14 Apr 2025 14:54:31 -0700 Subject: [PATCH 4/5] bumping deps version --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 639dbd9942..9adbd17eec 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,7 +34,7 @@ dependencies = [ "pydantic>=2.3, <3", "scalecodec==1.2.11", "uvicorn", - "bittensor-commit-reveal>=0.3.1", + "bittensor-commit-reveal>=0.4.0", "bittensor-wallet>=3.0.8", "async-substrate-interface>=1.1.0" ] From 999c51ba88683ea300e42f3cfb7f631228efad98 Mon Sep 17 00:00:00 2001 From: Roman Date: Mon, 14 Apr 2025 14:57:23 -0700 Subject: [PATCH 5/5] remove wrong import --- bittensor/core/timelock.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/bittensor/core/timelock.py b/bittensor/core/timelock.py index 5c620ac895..b40c21949e 100644 --- a/bittensor/core/timelock.py +++ b/bittensor/core/timelock.py @@ -6,7 +6,6 @@ encrypt as _btr_encrypt, decrypt as _btr_decrypt, get_latest_round, - get_reveal_round_signature, ) TLE_ENCRYPTED_DATA_SUFFIX = b"AES_GCM_" @@ -126,6 +125,5 @@ def wait_reveal_and_decrypt( "decrypt", "encrypt", "get_latest_round", - "get_reveal_round_signature", "wait_reveal_and_decrypt", ]