feat: code quality (#31)

* feat: code quality
This commit is contained in:
dni ⚡ 2024-08-30 13:07:33 +02:00 committed by GitHub
parent d656d41b90
commit a8eb139360
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
25 changed files with 3192 additions and 237 deletions

View file

@ -124,27 +124,29 @@ def decode(hrp, addr):
hrpgot, data, spec = bech32_decode(addr)
if hrpgot != hrp:
return (None, None)
decoded = convertbits(data[1:], 5, 8, False)
decoded = convertbits(data[1:], 5, 8, False) # type: ignore
if decoded is None or len(decoded) < 2 or len(decoded) > 40:
return (None, None)
if data[0] > 16:
if data[0] > 16: # type: ignore
return (None, None)
if data[0] == 0 and len(decoded) != 20 and len(decoded) != 32:
if data[0] == 0 and len(decoded) != 20 and len(decoded) != 32: # type: ignore
return (None, None)
if (
data[0] == 0
data[0] == 0 # type: ignore
and spec != Encoding.BECH32
or data[0] != 0
or data[0] != 0 # type: ignore
and spec != Encoding.BECH32M
):
return (None, None)
return (data[0], decoded)
return (data[0], decoded) # type: ignore
def encode(hrp, witver, witprog):
"""Encode a segwit address."""
spec = Encoding.BECH32 if witver == 0 else Encoding.BECH32M
ret = bech32_encode(hrp, [witver] + convertbits(witprog, 8, 5), spec)
wit_prog = convertbits(witprog, 8, 5)
assert wit_prog
ret = bech32_encode(hrp, [witver, *wit_prog], spec)
if decode(hrp, ret) == (None, None):
return None
return ret

View file

@ -3,7 +3,7 @@ import time
from dataclasses import dataclass, field
from enum import IntEnum
from hashlib import sha256
from typing import List
from typing import Optional
from secp256k1 import PublicKey
@ -21,14 +21,14 @@ class EventKind(IntEnum):
@dataclass
class Event:
content: str = None
public_key: str = None
created_at: int = None
content: Optional[str] = None
public_key: Optional[str] = None
created_at: Optional[int] = None
kind: int = EventKind.TEXT_NOTE
tags: List[List[str]] = field(
tags: list[list[str]] = field(
default_factory=list
) # Dataclasses require special handling when the default value is a mutable type
signature: str = None
signature: Optional[str] = None
def __post_init__(self):
if self.content is not None and not isinstance(self.content, str):
@ -40,7 +40,7 @@ class Event:
@staticmethod
def serialize(
public_key: str, created_at: int, kind: int, tags: List[List[str]], content: str
public_key: str, created_at: int, kind: int, tags: list[list[str]], content: str
) -> bytes:
data = [0, public_key, created_at, kind, tags, content]
data_str = json.dumps(data, separators=(",", ":"), ensure_ascii=False)
@ -48,7 +48,7 @@ class Event:
@staticmethod
def compute_id(
public_key: str, created_at: int, kind: int, tags: List[List[str]], content: str
public_key: str, created_at: int, kind: int, tags: list[list[str]], content: str
):
return sha256(
Event.serialize(public_key, created_at, kind, tags, content)
@ -57,6 +57,9 @@ class Event:
@property
def id(self) -> str:
# Always recompute the id to reflect the up-to-date state of the Event
assert self.public_key
assert self.created_at
assert self.content
return Event.compute_id(
self.public_key, self.created_at, self.kind, self.tags, self.content
)
@ -70,6 +73,8 @@ class Event:
self.tags.append(["e", event_id])
def verify(self) -> bool:
assert self.public_key
assert self.signature
pub_key = PublicKey(
bytes.fromhex("02" + self.public_key), True
) # add 02 for schnorr (bip340)
@ -96,9 +101,9 @@ class Event:
@dataclass
class EncryptedDirectMessage(Event):
recipient_pubkey: str = None
cleartext_content: str = None
reference_event_id: str = None
recipient_pubkey: Optional[str] = None
cleartext_content: Optional[str] = None
reference_event_id: Optional[str] = None
def __post_init__(self):
if self.content is not None:

View file

@ -1,12 +1,13 @@
import base64
import secrets
from typing import Optional
import secp256k1
from cffi import FFI
from cryptography.hazmat.primitives import padding
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from . import bech32
from .bech32 import Encoding, bech32_decode, bech32_encode, convertbits
from .event import EncryptedDirectMessage, Event, EventKind
@ -15,44 +16,51 @@ class PublicKey:
self.raw_bytes = raw_bytes
def bech32(self) -> str:
converted_bits = bech32.convertbits(self.raw_bytes, 8, 5)
return bech32.bech32_encode("npub", converted_bits, bech32.Encoding.BECH32)
converted_bits = convertbits(self.raw_bytes, 8, 5)
return bech32_encode("npub", converted_bits, Encoding.BECH32)
def hex(self) -> str:
return self.raw_bytes.hex()
def verify_signed_message_hash(self, hash: str, sig: str) -> bool:
def verify_signed_message_hash(self, message_hash: str, sig: str) -> bool:
pk = secp256k1.PublicKey(b"\x02" + self.raw_bytes, True)
return pk.schnorr_verify(bytes.fromhex(hash), bytes.fromhex(sig), None, True)
return pk.schnorr_verify(
bytes.fromhex(message_hash), bytes.fromhex(sig), None, True
)
@classmethod
def from_npub(cls, npub: str):
"""Load a PublicKey from its bech32/npub form"""
hrp, data, spec = bech32.bech32_decode(npub)
raw_public_key = bech32.convertbits(data, 5, 8)[:-1]
hrp, data, spec = bech32_decode(npub)
raw_data = convertbits(data, 5, 8)
assert raw_data
raw_public_key = raw_data[:-1]
return cls(bytes(raw_public_key))
class PrivateKey:
def __init__(self, raw_secret: bytes = None) -> None:
def __init__(self, raw_secret: Optional[bytes] = None) -> None:
if raw_secret is not None:
self.raw_secret = raw_secret
else:
self.raw_secret = secrets.token_bytes(32)
sk = secp256k1.PrivateKey(self.raw_secret)
assert sk.pubkey
self.public_key = PublicKey(sk.pubkey.serialize()[1:])
@classmethod
def from_nsec(cls, nsec: str):
"""Load a PrivateKey from its bech32/nsec form"""
hrp, data, spec = bech32.bech32_decode(nsec)
raw_secret = bech32.convertbits(data, 5, 8)[:-1]
hrp, data, spec = bech32_decode(nsec)
raw_data = convertbits(data, 5, 8)
assert raw_data
raw_secret = raw_data[:-1]
return cls(bytes(raw_secret))
def bech32(self) -> str:
converted_bits = bech32.convertbits(self.raw_secret, 8, 5)
return bech32.bech32_encode("nsec", converted_bits, bech32.Encoding.BECH32)
converted_bits = convertbits(self.raw_secret, 8, 5)
return bech32_encode("nsec", converted_bits, Encoding.BECH32)
def hex(self) -> str:
return self.raw_secret.hex()
@ -83,6 +91,8 @@ class PrivateKey:
)
def encrypt_dm(self, dm: EncryptedDirectMessage) -> None:
assert dm.cleartext_content
assert dm.recipient_pubkey
dm.content = self.encrypt_message(
message=dm.cleartext_content, public_key_hex=dm.recipient_pubkey
)
@ -105,14 +115,14 @@ class PrivateKey:
return unpadded_data.decode()
def sign_message_hash(self, hash: bytes) -> str:
def sign_message_hash(self, message_hash: bytes) -> str:
sk = secp256k1.PrivateKey(self.raw_secret)
sig = sk.schnorr_sign(hash, None, raw=True)
sig = sk.schnorr_sign(message_hash, None, raw=True)
return sig.hex()
def sign_event(self, event: Event) -> None:
if event.kind == EventKind.ENCRYPTED_DIRECT_MESSAGE and event.content is None:
self.encrypt_dm(event)
self.encrypt_dm(event) # type: ignore
if event.public_key is None:
event.public_key = self.public_key.hex()
event.signature = self.sign_message_hash(bytes.fromhex(event.id))
@ -121,7 +131,9 @@ class PrivateKey:
return self.raw_secret == other.raw_secret
def mine_vanity_key(prefix: str = None, suffix: str = None) -> PrivateKey:
def mine_vanity_key(
prefix: Optional[str] = None, suffix: Optional[str] = None
) -> PrivateKey:
if prefix is None and suffix is None:
raise ValueError("Expected at least one of 'prefix' or 'suffix' arguments")

View file

@ -2,7 +2,6 @@ import asyncio
import json
import time
from queue import Queue
from typing import List
from loguru import logger
from websocket import WebSocketApp
@ -21,14 +20,14 @@ class Relay:
self.error_counter: int = 0
self.error_threshold: int = 100
self.error_list: List[str] = []
self.notice_list: List[str] = []
self.error_list: list[str] = []
self.notice_list: list[str] = []
self.last_error_date: int = 0
self.num_received_events: int = 0
self.num_sent_events: int = 0
self.num_subscriptions: int = 0
self.queue = Queue()
self.queue: Queue = Queue()
def connect(self):
self.ws = WebSocketApp(
@ -63,9 +62,10 @@ class Relay:
def publish(self, message: str):
self.queue.put(message)
def publish_subscriptions(self, subscriptions: List[Subscription] = []):
def publish_subscriptions(self, subscriptions: list[Subscription]):
for s in subscriptions:
json_str = json.dumps(["REQ", s.id] + s.filters)
assert s.filters
json_str = json.dumps(["REQ", s.id, *s.filters])
self.publish(json_str)
async def queue_worker(self):
@ -84,14 +84,14 @@ class Relay:
logger.warning(f"[Relay: {self.url}] Closing queue worker.")
return
def close_subscription(self, id: str) -> None:
def close_subscription(self, sub_id: str) -> None:
try:
self.publish(json.dumps(["CLOSE", id]))
self.publish(json.dumps(["CLOSE", sub_id]))
except Exception as e:
logger.debug(f"[Relay: {self.url}] Failed to close subscription: {e}")
def add_notice(self, notice: str):
self.notice_list = [notice] + self.notice_list
self.notice_list = [notice, *self.notice_list]
def _on_open(self, _):
logger.info(f"[Relay: {self.url}] Connected.")
@ -110,7 +110,7 @@ class Relay:
self.message_pool.add_message(message, self.url)
def _on_error(self, _, error):
logger.warning(f"[Relay: {self.url}] Error: '{str(error)}'")
logger.warning(f"[Relay: {self.url}] Error: '{error!s}'")
self._append_error_message(str(error))
self.close()
@ -122,5 +122,5 @@ class Relay:
def _append_error_message(self, message):
self.error_counter += 1
self.error_list = [message] + self.error_list
self.error_list = [message, *self.error_list]
self.last_error_date = int(time.time())

View file

@ -1,7 +1,7 @@
from typing import List
from typing import Optional
class Subscription:
def __init__(self, id: str, filters: List[str] = None) -> None:
def __init__(self, id: str, filters: Optional[list[str]] = None) -> None:
self.id = id
self.filters = filters