Mining refactor with new features

nostr/key.py

@@ -108,19 +108,6 @@ def sign_delegation(self, delegation: Delegation) -> None:
     def __eq__(self, other):
         return self.raw_secret == other.raw_secret
 
-def mine_vanity_key(prefix: str = None, suffix: str = None) -> PrivateKey:
-    if prefix is None and suffix is None:
-        raise ValueError("Expected at least one of 'prefix' or 'suffix' arguments")
-
-    while True:
-        sk = PrivateKey()
-        if prefix is not None and not sk.public_key.bech32()[5:5+len(prefix)] == prefix:
-            continue
-        if suffix is not None and not sk.public_key.bech32()[-len(suffix):] == suffix:
-            continue
-        break
-
-    return sk
 
 ffi = FFI()
 @ffi.callback("int (unsigned char *, const unsigned char *, const unsigned char *, void *)")

nostr/pow.py

@@ -25,30 +25,110 @@ def count_leading_zero_bits(hex_str: str) -> int:
 
     return total
 
+
+def _guess_event(content: str, public_key: str, kind: int, tags: list=[]) -> Event:
+    event = Event(public_key, content, kind, tags)
+    num_leading_zero_bits = count_leading_zero_bits(event.id)
+    return num_leading_zero_bits, event
+
+
 def mine_event(content: str, difficulty: int, public_key: str, kind: int, tags: list=[]) -> Event:
     all_tags = [["nonce", "1", str(difficulty)]]
     all_tags.extend(tags)
 
-    created_at = int(time.time())
-    event_id = Event.compute_id(public_key, created_at, kind, all_tags, content)
-    num_leading_zero_bits = count_leading_zero_bits(event_id)
+    num_leading_zero_bits, event = _guess_event(content, public_key, kind, all_tags)
 
     attempts = 1
     while num_leading_zero_bits < difficulty:
         attempts += 1
         all_tags[0][1] = str(attempts)
-        created_at = int(time.time())
-        event_id = Event.compute_id(public_key, created_at, kind, all_tags, content)
-        num_leading_zero_bits = count_leading_zero_bits(event_id)
+        num_leading_zero_bits, event = _guess_event(content, public_key, kind, all_tags)
+        num_leading_zero_bits = count_leading_zero_bits(event.id)
 
-    return Event(public_key, content, created_at, kind, all_tags, event_id)
+    return event
 
-def mine_key(difficulty: int) -> PrivateKey:
+
+def _guess_key():
     sk = PrivateKey()
     num_leading_zero_bits = count_leading_zero_bits(sk.public_key.hex())
+    return num_leading_zero_bits, sk
+
+
+def mine_key(difficulty: int) -> PrivateKey:
+    num_leading_zero_bits, sk = _guess_key()
 
     while num_leading_zero_bits < difficulty:
-        sk = PrivateKey()
-        num_leading_zero_bits = count_leading_zero_bits(sk.public_key.hex())
+        num_leading_zero_bits, sk = _guess_key()
 
     return sk
+
+
+bech32_chars = '023456789acdefghjklmnpqrstuvwxyz'
+
+
+def _guess_vanity_key():
+    sk = PrivateKey()
+    vk = sk.public_key.bech32()
+    return sk, vk
+
+
+def mine_vanity_key(prefix: str = None, suffix: str = None) -> PrivateKey:
+    if prefix is None and suffix is None:
+        raise ValueError("Expected at least one of 'prefix' or 'suffix' arguments")
+    for pattern in [prefix, suffix]:
+        if pattern is not None:
+            missing_chars = [c for c in pattern if c not in bech32_chars]
+            if len(missing_chars) > 0:
+                raise ValueError(
+                    f"{missing_chars} not in valid list of bech32 chars: ({bech32_chars})"
+                    )
+    while True:
+        sk, vk = _guess_vanity_key()
+        if prefix is not None and not vk[5:5+len(prefix)] == prefix:
+            continue
+        if suffix is not None and not vk[-len(suffix):] == suffix:
+            continue
+        break
+
+    return sk
+
+
+def get_hashrate(operation, n_guesses: int = 1e4, **operation_kwargs):
+    n_guesses = int(n_guesses)
+    def _time_operation():
+        start = time.perf_counter()
+        operation(**operation_kwargs)
+        end = time.perf_counter()
+        return end - start
+    t = sum([_time_operation() for _ in range(n_guesses)]) / n_guesses
+    hashrate = 1 / t
+    return hashrate
+
+
+def expected_time(n_pattern: int, n_options: int, hashrate: float):
+    p = 1 / n_options
+    expected_guesses = 1 / (p ** n_pattern)
+    time_seconds = expected_guesses / hashrate
+    return time_seconds
+
+
+def estimate_event_time(content: str, difficulty: int, public_key: str,
+                        kind: int, tags: list=[]) -> float:
+    all_tags = [["nonce", "1", str(difficulty)]]
+    all_tags.extend(tags)
+    hashrate = get_hashrate(_guess_event,
+                             content=content,
+                             public_key=public_key,
+                             kind=kind,
+                             tags=all_tags)
+    return expected_time(difficulty, 2, hashrate)
+
+
+def estimate_vanity_time(n_pattern: int):
+    hashrate = get_hashrate(_guess_vanity_key)
+    return expected_time(n_pattern, len(bech32_chars), hashrate)
+
+
+def estimate_key_time(difficulty: int) -> float:
+    hashrate = get_hashrate(_guess_key)
+    return expected_time(difficulty, 2, hashrate)

test/test_key.py

@@ -1,3 +1,4 @@
+import pytest
 from nostr.key import PrivateKey
 
 

test/test_pow.py

@@ -0,0 +1,67 @@
+import pytest
+from nostr.key import PrivateKey
+from nostr.event import EventKind
+from nostr import pow
+from nostr.pow import mine_event, mine_key, mine_vanity_key
+
+def test_mine_event():
+    """ test mining an event with specific difficulty """
+    public_key = PrivateKey().public_key.hex()
+    difficulty = 8
+    event = mine_event(content='test',difficulty=difficulty,
+                       public_key=public_key, kind=EventKind.TEXT_NOTE)
+    assert pow.count_leading_zero_bits(event.id) >= difficulty
+
+
+def test_mine_key():
+    """ test mining a public key with specific difficulty """
+    difficulty = 8
+    sk = mine_key(difficulty)
+    assert pow.count_leading_zero_bits(sk.public_key.hex()) >= difficulty
+
+
+def test_time_estimates():
+    """ test functions to estimate POW time """
+    public_key = PrivateKey().public_key.hex()
+
+    # test successful run of all estimators
+    pow.estimate_event_time('test',
+                            8,
+                            public_key,
+                            EventKind.TEXT_NOTE)
+    pow.estimate_key_time(8)
+    pow.estimate_vanity_time(8)
+
+
+def test_mine_vanity_key():
+    """ test vanity key mining """
+    pattern = '23'
+
+    # mine a valid pattern as prefix
+    sk = mine_vanity_key(prefix=pattern)
+    sk.public_key.bech32()
+    assert sk.public_key.bech32().startswith(f'npub1{pattern}')
+
+    # mine a valid pattern as suffix
+    sk = mine_vanity_key(suffix=pattern)
+    assert sk.public_key.bech32().endswith(pattern)
+
+    # mine an invalid pattern
+    pattern = '1'
+    expected_error = "not in valid list of bech32 chars"
+    with pytest.raises(ValueError) as e:
+        sk = mine_vanity_key(prefix=pattern)
+    assert expected_error in str(e)
+
+
+def test_expected_pow_times():
+    """ sense check expected calculations using known patterns """
+    # assume constant hashrate
+    hashrate = 10000
+
+    # calculate expected time to get a 32-difficulty bit key
+    e1 = pow.expected_time(32, 2, hashrate)
+
+    # caluclate 8 leading 0 hex key, which is equivalent
+    e2 = pow.expected_time(8, 16, hashrate)
+    assert e1 == e2