server.py

#!/usr/bin/env python3
"""
p2pool-pool — Stratum proxy that bridges miners to a P2Pool node.

Architecture:
  Miners connect here via Monero Stratum (the same protocol XMRig uses with
  P2Pool directly). They authenticate with just a worker name — no wallet
  address required. p2pool-pool holds a single upstream connection to P2Pool
  using the operator's wallet, forwards jobs to all connected miners, collects
  their shares, and tracks per-worker hashrate.

  The hashrate data is written to hashrate.json so that p2pool-splitter can
  pay each worker proportionally based on actual measured contribution rather
  than a static config file.

  Flow:
    1. Miner connects and sends "login" with a worker name
    2. p2pool-pool responds with a session ID and the current job
    3. When P2Pool sends a new job, p2pool-pool broadcasts it to all miners
    4. Miner finds a share and sends "submit"
    5. p2pool-pool records the share for hashrate tracking, forwards to P2Pool
    6. Every 60s, p2pool-pool writes hashrate.json for p2pool-splitter

  Monero Stratum protocol (used by XMRig and P2Pool):
    login   — miner authenticates (worker name as login, password ignored)
    job     — server pushes a new mining job to the miner
    submit  — miner submits a found share
    keepalived — heartbeat (miner → server, server replies with keepalived)

Usage: python server.py --config config.yaml
"""

import argparse
import asyncio
import json
import logging
import time
import uuid
from collections import defaultdict, deque
from pathlib import Path

try:
    import yaml
except ImportError:
    yaml = None

__version__ = "1.0.0"

# ---------------------------------------------------------------------------
# Logging
# ---------------------------------------------------------------------------

def setup_logging(verbose: bool = False, log_file: str | None = None) -> logging.Logger:
    level = logging.DEBUG if verbose else logging.INFO
    handlers: list[logging.Handler] = [logging.StreamHandler()]
    if log_file:
        handlers.append(logging.FileHandler(log_file))
    logging.basicConfig(
        level=level,
        format="%(asctime)s [%(levelname)s] %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
        handlers=handlers,
    )
    return logging.getLogger("p2pool-pool")

log = logging.getLogger("p2pool-pool")


# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------

class Config:
    def __init__(self, path: str):
        data = self._load(path)

        # Stratum server — what miners connect to
        self.listen_host: str = data.get("listen_host", "0.0.0.0")
        self.listen_port: int = int(data.get("listen_port", 3333))

        # Upstream P2Pool Stratum — where we forward work and shares.
        # Run P2Pool on a different port (e.g. 3334) to avoid conflict with us.
        self.p2pool_host: str = data.get("p2pool_host", "127.0.0.1")
        self.p2pool_port: int = int(data.get("p2pool_port", 3334))

        # Operator's Monero wallet address — used as the login when connecting
        # to P2Pool. All P2Pool payouts go to this wallet; p2pool-splitter
        # then distributes to individual workers.
        self.operator_wallet: str = data.get("operator_wallet", "")
        if not self.operator_wallet:
            raise ValueError("operator_wallet is required in config")

        # How often (seconds) to write hashrate.json for p2pool-splitter
        self.hashrate_export_interval: int = int(data.get("hashrate_export_interval", 60))

        # Path to write the live hashrate snapshot
        self.hashrate_file: str = data.get("hashrate_file", "hashrate.json")

        # How many seconds of share history to use when estimating hashrate.
        # Longer = smoother values but slower to react to miners joining/leaving.
        self.hashrate_window: int = int(data.get("hashrate_window", 600))  # 10 min

        self.verbose: bool = bool(data.get("verbose", False))
        self.log_file: str | None = data.get("log_file")

    def _load(self, path: str) -> dict:
        p = Path(path)
        if not p.exists():
            raise FileNotFoundError(f"Config not found: {path}")
        text = p.read_text()
        if path.endswith(".json"):
            return json.loads(text)
        if yaml:
            return yaml.safe_load(text)
        raise ImportError("PyYAML required: pip install pyyaml")


# ---------------------------------------------------------------------------
# Share tracker — records accepted shares and computes per-worker hashrate
# ---------------------------------------------------------------------------

class ShareTracker:
    """
    Records the timestamp and difficulty of every accepted share per worker.

    Hashrate is estimated as: sum(share_difficulty) / elapsed_window_seconds.
    This is the same method used by real pools — each share represents a
    known amount of hashing work (proportional to its difficulty), so
    dividing by time gives an accurate average hash rate.

    Old entries outside the rolling window are pruned automatically.
    """

    def __init__(self, window_seconds: int):
        self.window = window_seconds
        # worker_name → deque of (unix_timestamp, difficulty) tuples
        self._shares: dict[str, deque] = defaultdict(deque)

    def record(self, worker: str, difficulty: float) -> None:
        """Record one accepted share for a worker at the current time."""
        self._shares[worker].append((time.time(), difficulty))
        self._prune(worker)

    def _prune(self, worker: str) -> None:
        """Drop share records older than the hashrate window."""
        cutoff = time.time() - self.window
        q = self._shares[worker]
        while q and q[0][0] < cutoff:
            q.popleft()

    def hashrates(self) -> dict[str, float]:
        """
        Return estimated hashrate (H/s) for every active worker.
        Workers with no shares in the current window are excluded.
        """
        result = {}
        now = time.time()
        for worker, q in self._shares.items():
            self._prune(worker)
            if not q:
                continue
            elapsed = now - q[0][0]
            if elapsed < 1:
                elapsed = 1  # avoid division by zero on first share
            total_difficulty = sum(d for _, d in q)
            result[worker] = total_difficulty / elapsed
        return result

    def snapshot(self) -> dict:
        """Return a JSON-serialisable snapshot for hashrate.json."""
        rates = self.hashrates()
        return {
            "updated_at": int(time.time()),
            "window_seconds": self.window,
            "workers": [
                {"name": name, "hashrate": round(rate, 2)}
                for name, rate in sorted(rates.items())
            ],
        }


# ---------------------------------------------------------------------------
# Upstream connection to P2Pool (Monero Stratum protocol)
# ---------------------------------------------------------------------------

class P2PoolUpstream:
    """
    Manages one persistent Monero Stratum connection to the P2Pool node.

    Monero Stratum login flow:
      → {"method": "login", "params": {"login": "<wallet>", "pass": "x",
                                        "agent": "p2pool-pool/1.0.0"}, "id": 1}
      ← {"id": 1, "result": {"id": "<session>", "job": {...}, "status": "OK"}}

    Subsequent job notifications from P2Pool arrive as:
      ← {"method": "job", "params": {"blob": "...", "job_id": "...",
                                      "target": "...", "height": ...,
                                      "seed_hash": "..."}}

    Share submission:
      → {"method": "submit", "params": {"id": "<session>", "job_id": "...",
                                         "nonce": "...", "result": "..."}, "id": 2}
      ← {"id": 2, "result": {"status": "OK"}}   (or error)
    """

    def __init__(self, host: str, port: int, wallet: str):
        self.host = host
        self.port = port
        self.wallet = wallet
        self.on_job = None                      # callback(job_dict) set by StratumProxy
        self._writer: asyncio.StreamWriter | None = None
        self._msg_id = 0
        self._pending: dict[int, asyncio.Future] = {}
        self.session_id: str = ""               # assigned by P2Pool after login
        self.current_job: dict | None = None    # latest job params from P2Pool
        self.difficulty: float = 1.0            # current upstream difficulty

    def _next_id(self) -> int:
        self._msg_id += 1
        return self._msg_id

    async def connect(self) -> None:
        """
        Open a TCP connection to P2Pool and perform the Monero Stratum login.
        Starts a background task to continuously read incoming messages.
        """
        log.info(f"Connecting to P2Pool at {self.host}:{self.port}")
        reader, self._writer = await asyncio.open_connection(self.host, self.port)
        asyncio.ensure_future(self._read_loop(reader))

        # Send login — P2Pool authenticates by wallet address.
        # Register the pending future BEFORE sending so _dispatch can resolve it.
        # On success, P2Pool responds with a session ID and the first job.
        login_id = self._next_id()
        loop = asyncio.get_running_loop()
        login_future: asyncio.Future = loop.create_future()
        self._pending[login_id] = login_future

        await self._send({
            "id": login_id,
            "method": "login",
            "params": {
                "login": self.wallet,
                "pass": "x",
                "agent": f"p2pool-pool/{__version__}",
                "rigid": "",    # optional rig ID
            },
        })
        log.info("Login sent to P2Pool — waiting for job")

        # Block until P2Pool confirms login (delivers session ID + first job).
        # _dispatch handles the response and resolves this future.
        try:
            await asyncio.wait_for(login_future, timeout=30)
        except asyncio.TimeoutError:
            raise ConnectionError("P2Pool login timed out — check wallet and port")

    async def submit(self, session_id: str, job_id: str,
                     nonce: str, result: str) -> bool:
        """
        Forward a miner's share to P2Pool.

        Monero Stratum submit params:
          id      — the miner's P2Pool session ID (we use our own upstream session)
          job_id  — the job this share was found for
          nonce   — 8-hex-char nonce the miner found
          result  — 64-hex-char hash result

        Returns True if P2Pool responds with status "OK".
        """
        msg_id = self._next_id()
        loop = asyncio.get_running_loop()
        future: asyncio.Future = loop.create_future()
        self._pending[msg_id] = future

        await self._send({
            "id": msg_id,
            "method": "submit",
            "params": {
                "id": self.session_id,   # our upstream session, not the miner's
                "job_id": job_id,
                "nonce": nonce,
                "result": result,
            },
        })

        try:
            return await asyncio.wait_for(future, timeout=10)
        except asyncio.TimeoutError:
            log.warning("P2Pool share submission timed out")
            self._pending.pop(msg_id, None)
            return False

    async def _send(self, obj: dict) -> None:
        if self._writer:
            self._writer.write((json.dumps(obj) + "\n").encode())
            await self._writer.drain()

    async def _read_loop(self, reader: asyncio.StreamReader) -> None:
        """
        Read newline-delimited JSON messages from P2Pool and dispatch:
          - login response   → store session ID, deliver first job
          - "job" method     → store current job, call on_job callback
          - submit response  → resolve the waiting Future for that msg ID
        """
        buffer = b""
        while True:
            try:
                chunk = await reader.read(4096)
                if not chunk:
                    log.warning("P2Pool upstream connection closed")
                    return
                buffer += chunk
                while b"\n" in buffer:
                    line, buffer = buffer.split(b"\n", 1)
                    line = line.strip()
                    if line:
                        self._dispatch(json.loads(line.decode()))
            except Exception as e:
                log.error(f"P2Pool read error: {e}")
                return

    def _dispatch(self, msg: dict) -> None:
        """Route an incoming P2Pool message to the appropriate handler."""
        method = msg.get("method")
        msg_id = msg.get("id")

        if method == "job":
            # Unsolicited job notification — P2Pool found a new block template
            job = msg.get("params", {})
            self._store_job(job)
            if self.on_job:
                asyncio.ensure_future(self.on_job(job))

        elif msg_id is not None and msg_id in self._pending:
            # Response to one of our requests (login or submit)
            future = self._pending.pop(msg_id)
            result = msg.get("result", {})
            error = msg.get("error")

            if error:
                log.warning(f"P2Pool error (id={msg_id}): {error}")
                if not future.done():
                    future.set_result(False)
                return

            # Login response — result contains session ID and first job
            if isinstance(result, dict) and "id" in result:
                self.session_id = result["id"]
                log.info(f"P2Pool login OK — session {self.session_id[:8]}…")
                job = result.get("job")
                if job:
                    self._store_job(job)
                    if self.on_job:
                        asyncio.ensure_future(self.on_job(job))
                if not future.done():
                    future.set_result(True)
            else:
                # Submit response — result is {"status": "OK"} or similar
                accepted = isinstance(result, dict) and result.get("status") == "OK"
                if not future.done():
                    future.set_result(accepted)

    def _store_job(self, job: dict) -> None:
        """Cache the latest job and extract its difficulty."""
        self.current_job = job
        # P2Pool encodes difficulty in the "target" field as a 4- or 8-byte
        # little-endian hex value. A smaller target = higher difficulty.
        target_hex = job.get("target", "")
        if target_hex:
            try:
                self.difficulty = self._target_to_difficulty(target_hex)
            except Exception:
                pass  # keep the last known difficulty on parse failure
        log.debug(f"Job stored: id={job.get('job_id','?')} "
                  f"height={job.get('height','?')} diff={self.difficulty:.0f}")

    @staticmethod
    def _target_to_difficulty(target_hex: str) -> float:
        """
        Convert a Monero Stratum target hex string to a difficulty value.

        P2Pool sends target as an 8-hex-char (4-byte) little-endian integer.
        Difficulty = 2^32 / target_uint32  (standard Monero formula).
        """
        # Pad to 8 chars if shorter
        target_hex = target_hex.zfill(8)
        target_bytes = bytes.fromhex(target_hex[:8])
        target_uint32 = int.from_bytes(target_bytes, byteorder="little")
        if target_uint32 == 0:
            return 1.0
        return (2 ** 32) / target_uint32


# ---------------------------------------------------------------------------
# Per-miner session (Monero Stratum)
# ---------------------------------------------------------------------------

class MinerSession:
    """
    Represents one connected miner.

    Handles the Monero Stratum protocol:
      login      — miner identifies with a worker name
      submit     — miner submits a found share
      keepalived — heartbeat ping (we reply immediately)

    On login, we respond with a unique session ID and the current job.
    On submit, we record the share for hashrate tracking and forward
    to P2Pool upstream.
    """

    def __init__(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter,
                 upstream: P2PoolUpstream, tracker: ShareTracker):
        self.reader = reader
        self.writer = writer
        self.upstream = upstream
        self.tracker = tracker
        self.worker_name: str = "unknown"
        self.session_id: str = str(uuid.uuid4()).replace("-", "")[:32]
        self.authorized: bool = False
        addr = writer.get_extra_info("peername")
        self.addr = f"{addr[0]}:{addr[1]}" if addr else "?"

    async def _send(self, obj: dict) -> None:
        try:
            self.writer.write((json.dumps(obj) + "\n").encode())
            await self.writer.drain()
        except Exception:
            pass

    async def push_job(self, job: dict) -> None:
        """Push a new mining job to this miner."""
        if not self.authorized:
            return
        await self._send({"method": "job", "params": job})

    async def run(self) -> None:
        """Main read loop — process all incoming Stratum messages from this miner."""
        log.info(f"Miner connected from {self.addr}")
        buffer = b""
        try:
            while True:
                chunk = await self.reader.read(4096)
                if not chunk:
                    break
                buffer += chunk
                while b"\n" in buffer:
                    line, buffer = buffer.split(b"\n", 1)
                    line = line.strip()
                    if line:
                        try:
                            await self._dispatch(json.loads(line.decode()))
                        except Exception as e:
                            log.debug(f"[{self.worker_name}] Message error: {e}")
        except Exception as e:
            log.debug(f"[{self.worker_name}] Connection dropped: {e}")
        finally:
            log.info(f"Miner disconnected: {self.worker_name} ({self.addr})")
            try:
                self.writer.close()
            except Exception:
                pass

    async def _dispatch(self, msg: dict) -> None:
        method = msg.get("method")
        msg_id = msg.get("id")

        if method == "login":
            await self._handle_login(msg_id, msg.get("params", {}))
        elif method == "submit":
            await self._handle_submit(msg_id, msg.get("params", {}))
        elif method == "keepalived":
            # Heartbeat — miner just wants to know we're still here
            await self._send({"id": msg_id, "result": {"status": "KEEPALIVED"},
                              "error": None})
        else:
            log.debug(f"[{self.worker_name}] Unknown method: {method}")

    async def _handle_login(self, msg_id: int, params: dict) -> None:
        """
        Authenticate the miner. We use the login field as the worker name
        so p2pool-splitter can map hashrate to payout addresses.

        Response includes:
          id  — session ID we assign to this miner
          job — the current P2Pool job so the miner can start immediately
        """
        self.worker_name = params.get("login", "worker") or "worker"
        self.authorized = True

        job = self.upstream.current_job or {}
        await self._send({
            "id": msg_id,
            "result": {
                "id": self.session_id,
                "job": job,
                "status": "OK",
            },
            "error": None,
        })
        log.info(f"Worker authorized: {self.worker_name} ({self.addr})")

    async def _handle_submit(self, msg_id: int, params: dict) -> None:
        """
        Process a share submission from the miner.

        Monero Stratum submit params (dict):
          id      — miner's session ID
          job_id  — which job this share belongs to
          nonce   — 8-hex-char nonce
          result  — 64-hex-char hash result

        Steps:
          1. Record the share for hashrate tracking (always, regardless of difficulty)
          2. Forward to P2Pool upstream and relay the accept/reject response
        """
        job_id = params.get("job_id", "")
        nonce  = params.get("nonce", "")
        result = params.get("result", "")

        if not job_id or not nonce or not result:
            await self._send({"id": msg_id, "result": None,
                              "error": {"code": -1, "message": "Missing params"}})
            return

        # Always count the share for hashrate tracking.
        # Use the upstream difficulty since we don't do per-miner vardiff yet.
        self.tracker.record(self.worker_name, self.upstream.difficulty)

        # Forward to P2Pool and wait for the response
        accepted = await self.upstream.submit(
            self.session_id, job_id, nonce, result
        )

        status = "OK" if accepted else "REJECTED"
        log.debug(f"[{self.worker_name}] Share {status} "
                  f"(job {job_id[:8]}… nonce {nonce})")

        await self._send({
            "id": msg_id,
            "result": {"status": status},
            "error": None,
        })


# ---------------------------------------------------------------------------
# Main proxy
# ---------------------------------------------------------------------------

class StratumProxy:
    """
    Ties together the P2Pool upstream and all miner sessions.
    Broadcasts new jobs to every connected miner and periodically
    exports hashrate data for p2pool-splitter.
    """

    def __init__(self, config: Config):
        self.cfg = config
        self.upstream = P2PoolUpstream(
            config.p2pool_host, config.p2pool_port, config.operator_wallet
        )
        self.upstream.on_job = self._broadcast_job
        self.tracker = ShareTracker(config.hashrate_window)
        self.sessions: list[MinerSession] = []

    async def _broadcast_job(self, job: dict) -> None:
        """Forward a new job from P2Pool to every connected miner."""
        if not self.sessions:
            return
        log.debug(f"Broadcasting job {job.get('job_id','?')[:8]}… "
                  f"to {len(self.sessions)} miner(s)")
        await asyncio.gather(
            *(s.push_job(job) for s in self.sessions),
            return_exceptions=True,
        )

    async def _handle_miner(self, reader: asyncio.StreamReader,
                             writer: asyncio.StreamWriter) -> None:
        """Called by asyncio for each new incoming miner connection."""
        session = MinerSession(reader, writer, self.upstream, self.tracker)
        self.sessions.append(session)
        try:
            await session.run()
        finally:
            if session in self.sessions:
                self.sessions.remove(session)

    async def _export_loop(self) -> None:
        """
        Periodically write hashrate.json with live per-worker H/s estimates.
        p2pool-splitter reads this file to determine payout proportions.
        """
        while True:
            await asyncio.sleep(self.cfg.hashrate_export_interval)
            snapshot = self.tracker.snapshot()
            Path(self.cfg.hashrate_file).write_text(json.dumps(snapshot, indent=2))
            workers = snapshot["workers"]
            if workers:
                log.info("Hashrate: " + ", ".join(
                    f"{w['name']} {w['hashrate']:.0f} H/s" for w in workers
                ))
            else:
                log.info("Hashrate: no active workers")

    async def run(self) -> None:
        await self.upstream.connect()

        server = await asyncio.start_server(
            self._handle_miner, self.cfg.listen_host, self.cfg.listen_port
        )
        log.info(f"Listening for miners on "
                 f"{self.cfg.listen_host}:{self.cfg.listen_port}")

        async with server:
            await asyncio.gather(
                server.serve_forever(),
                self._export_loop(),
            )


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def main() -> None:
    parser = argparse.ArgumentParser(
        description="p2pool-pool: Stratum proxy bridging miners to P2Pool",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    parser.add_argument("--config", "-c", default="config.yaml",
                        help="Config file path [default: config.yaml]")
    parser.add_argument("--version", action="version",
                        version=f"%(prog)s {__version__}")
    args = parser.parse_args()

    try:
        cfg = Config(args.config)
    except (FileNotFoundError, ValueError, ImportError) as e:
        print(f"ERROR: {e}")
        raise SystemExit(1)

    setup_logging(cfg.verbose, cfg.log_file)
    log.info(f"p2pool-pool v{__version__} starting")
    log.info(f"Operator wallet: {cfg.operator_wallet[:16]}…")
    log.info(f"P2Pool upstream: {cfg.p2pool_host}:{cfg.p2pool_port}")

    proxy = StratumProxy(cfg)
    try:
        asyncio.run(proxy.run())
    except KeyboardInterrupt:
        log.info("Shutting down")


if __name__ == "__main__":
    main()