Merge pull request #2824 from opentensor/feat/roman/add-timelock

Add `timelock` module

Author: basfroman

bittensor/core/timelock.py

@@ -0,0 +1,129 @@
+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,
+)
+
+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(
+                "<Q", encrypted_data.split(TLE_ENCRYPTED_DATA_SUFFIX)[-1]
+            )[0]
+        except (struct.error, TypeError, IndexError):
+            raise ValueError("Failed to parse reveal round from encrypted data.")
+
+    while get_latest_round() <= reveal_round:
+        # sleep Drand QuickNet period time (3 sec)
+        time.sleep(3)
+
+    return decrypt(encrypted_data, no_errors, return_str)
+
+
+__all__ = [
+    "decrypt",
+    "encrypt",
+    "get_latest_round",
+    "wait_reveal_and_decrypt",
+]

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

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"
 ]

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(
+        "<Q", encrypted.split(timelock.TLE_ENCRYPTED_DATA_SUFFIX)[-1]
+    )[0]
+    assert parsed == reveal_round