plm-lottery/app/electrum/client.py

import asyncio
import itertools
import json
import ssl


class ElectrumError(Exception):
    pass


class ElectrumClient:
    """Minimal asyncio Electrum protocol client: line-delimited JSON-RPC over TLS.

    Push notifications (blockchain.headers.subscribe, blockchain.scripthash.subscribe)
    arrive under the *same* method name as the subscribe call, multiplexed for every
    scripthash subscribed — callers read `notifications(method)` and, for scripthash
    pushes, dispatch on `params[0]` (the scripthash) themselves.
    """

    def __init__(self, host: str, port: int, use_ssl: bool = True):
        self.host = host
        self.port = port
        self.use_ssl = use_ssl
        self._reader: asyncio.StreamReader | None = None
        self._writer: asyncio.StreamWriter | None = None
        self._id_counter = itertools.count(1)
        self._pending: dict[int, asyncio.Future] = {}
        self._subscriptions: dict[str, asyncio.Queue] = {}
        self._read_task: asyncio.Task | None = None

    async def connect(self) -> None:
        # Electrum servers commonly present self-signed certs; the protocol's trust
        # model is server consensus, not TLS PKI, so we only use SSL for transport
        # encryption and don't verify the certificate chain/hostname.
        ssl_context = None
        if self.use_ssl:
            ssl_context = ssl.create_default_context()
            ssl_context.check_hostname = False
            ssl_context.verify_mode = ssl.CERT_NONE
        self._reader, self._writer = await asyncio.open_connection(self.host, self.port, ssl=ssl_context)
        self._read_task = asyncio.create_task(self._read_loop())
        await self.request("server.version", ["plm-lottery", "1.4"])

    async def close(self) -> None:
        if self._read_task is not None:
            self._read_task.cancel()
        if self._writer is not None:
            self._writer.close()
            try:
                await asyncio.wait_for(self._writer.wait_closed(), timeout=2)
            except (ssl.SSLError, TimeoutError, asyncio.TimeoutError):
                pass  # some Electrum servers don't send a clean TLS close_notify

    async def request(self, method: str, params: list | None = None) -> object:
        if self._writer is None:
            raise ElectrumError("not connected")
        request_id = next(self._id_counter)
        future: asyncio.Future = asyncio.get_event_loop().create_future()
        self._pending[request_id] = future
        payload = json.dumps({"id": request_id, "method": method, "params": params or []}) + "\n"
        self._writer.write(payload.encode())
        await self._writer.drain()
        return await future

    def notifications(self, method: str) -> asyncio.Queue:
        return self._subscriptions.setdefault(method, asyncio.Queue())

    async def subscribe_headers(self) -> dict:
        self.notifications("blockchain.headers.subscribe")
        return await self.request("blockchain.headers.subscribe")

    async def subscribe_scripthash(self, scripthash: str) -> str | None:
        self.notifications("blockchain.scripthash.subscribe")
        return await self.request("blockchain.scripthash.subscribe", [scripthash])

    async def listunspent(self, scripthash: str) -> list[dict]:
        return await self.request("blockchain.scripthash.listunspent", [scripthash])

    async def broadcast(self, raw_tx_hex: str) -> str:
        return await self.request("blockchain.transaction.broadcast", [raw_tx_hex])

    async def get_transaction(self, txid: str, verbose: bool = False) -> object:
        return await self.request("blockchain.transaction.get", [txid, verbose])

    async def _read_loop(self) -> None:
        assert self._reader is not None
        try:
            while True:
                line = await self._reader.readline()
                if not line:
                    break
                message = json.loads(line)
                self._dispatch(message)
        finally:
            error = ElectrumError("connection closed")
            for future in self._pending.values():
                if not future.done():
                    future.set_exception(error)
            self._pending.clear()

    def _dispatch(self, message: dict) -> None:
        message_id = message.get("id")
        if message_id is not None and message_id in self._pending:
            future = self._pending.pop(message_id)
            if future.done():
                return
            if message.get("error"):
                future.set_exception(ElectrumError(message["error"]))
            else:
                future.set_result(message.get("result"))
        elif "method" in message:
            queue = self._subscriptions.get(message["method"])
            if queue is not None:
                queue.put_nowait(message.get("params"))