rostrum 14.0.1

import asyncio
from enum import Enum
import json
import ssl
import os

from contextlib import asynccontextmanager
from typing import List

from test_framework.electrumconnection import ElectrumConnection


def ssl_context():
    base_dir = os.path.dirname(os.path.abspath(__file__))
    cert_path = os.path.join(base_dir, "data", "dummy_server_cert.pem")
    key_path = os.path.join(base_dir, "data", "dummy_server_key.pem")

    context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
    context.load_cert_chain(certfile=cert_path, keyfile=key_path)
    context.check_hostname = False
    context.verify_mode = ssl.CERT_NONE
    return context


class Protocol(Enum):
    TCP = (1,)
    SSL = 2


# A dummy electrum server implementation for testing peer discovery.
# pylint: disable=too-many-instance-attributes
class DummyElectrumServer:
    def __init__(self, genesis_hash, protocol_min, protocol_max):
        self.genesis_hash = genesis_hash
        self.protocol_min = protocol_min
        self.protocol_max = protocol_max
        self.servers = []
        self.tcp_port = None
        self.ssl_port = None

        self.got_remote_connection = False
        self.requested_peers = False

        # Keep a simple in-memory peer database:
        # Each entry will be a dict with basic info gleaned from the "server.add_peer" call.
        self.peers = []

    async def start(self, protocols: List[Protocol] | None = None):

        if protocols is None:
            protocols = [Protocol.TCP]

        assert len(protocols) > 0

        if Protocol.TCP in protocols:
            s = await asyncio.start_server(self.handle_client, "127.0.0.1", 0)
            self.tcp_port = s.sockets[0].getsockname()[1]
            print("dummy: running on tcp port", self.tcp_port)
            self.servers.append(s)

        if Protocol.SSL in protocols:
            s = await asyncio.start_server(
                self.handle_client, "127.0.0.1", 0, ssl=ssl_context()
            )
            self.ssl_port = s.sockets[0].getsockname()[1]
            print("dummy: running on ssl port", self.ssl_port)
            self.servers.append(s)

    async def stop(self):
        for i, s in enumerate(self.servers):
            s.close()
            await s.wait_closed()
            print(f"dummy: server stopped [{i + 1} of {len(self.servers)}]")

    async def handle_client(self, reader, writer):
        print("dummy: connection from remote server")
        self.got_remote_connection = True
        while True:
            line = await reader.readline()
            print("rostrum -> dummy:", line)
            if not line:
                break
            try:
                request = json.loads(line.decode("utf-8").strip())
                result = self.handle_request(request)
                response = json.dumps(
                    {"jsonrpc": "2.0", "result": result, "id": request.get("id", 0)}
                ).encode("utf-8")
                print("dummy -> rostrum:", response)
                writer.write(response + b"\n")
                await writer.drain()
            except Exception as e:
                print("dummy: disconnecting peer", e)
                break
        writer.close()
        await writer.wait_closed()

    def handle_request(self, request):
        method = request.get("method", "")
        if method == "server.ping":
            return None
        if method == "server.features":
            return self.features()
        if method == "server.version":
            return ["Mock Electrum v0.0", self.protocol_max]
        if method == "server.add_peer":
            return self.add_peer(request.get("params", [{}])[0])
        if method == "server.peers.subscribe":
            return self.subscribe_peers()

        return "unsupported method"

    def features(self):
        return {
            "genesis_hash": self.genesis_hash,
            "hash_function": "sha256",
            "server_version": "Mock Electrum v0.0",
            "protocol_min": self.protocol_min,
            "protocol_max": self.protocol_max,
            "hosts": {
                "localhost": {"tcp_port": self.tcp_port, "ssl_port": self.ssl_port}
            },
        }

    def add_peer(self, peer_info):
        """
        Minimal dummy logic:
          1) Check if peer's genesis_hash matches ours (in real code, might also check height).
          2) Append peer to our local database if accepted.
          3) Return True if accepted, otherwise False.
        """
        # Very basic checks: if the peer info is empty or has no hosts, reject
        if not peer_info or "hosts" not in peer_info:
            return False

        # If there's a mismatch of genesis_hash with ours, we can either reject or ignore. Here we reject:
        if peer_info.get("genesis_hash", "") != self.genesis_hash:
            return False

        # Accept and store
        self.peers.append(peer_info)
        return True

    def subscribe_peers(self):
        self.requested_peers = True
        return []


@asynccontextmanager
async def dummy_electrum(protocols: List[Protocol] | None = None):
    cli = ElectrumConnection()
    await cli.connect()
    features = await cli.call("server.features")
    dummy = DummyElectrumServer(
        features["genesis_hash"], features["protocol_min"], features["protocol_max"]
    )
    await dummy.start(protocols)
    try:
        yield dummy, cli
    finally:
        await dummy.stop()
        cli.disconnect()