#!/usr/bin/env python3
"""
krgn-gbt-proxy.py - drop-in JSON-RPC proxy for Kerrigan mining pools.

From Kerrigan v1.3.2 (mainnet block 148000) the coinbase special-tx payload
carries a BLS miner-identity signature over vtx[0].vout[0].scriptPubKey. The
daemon signs whatever script it was handed via getblocktemplate's "pooladdress"
parameter, and signs OP_TRUE when that parameter is absent. A pool that rebuilds
the coinbase with its own payout script at vout[0] but never sent "pooladdress"
produces a block the network rejects with bad-cbtx-miner-identity-sig.

This proxy sits between the pool and kerrigand and injects "pooladdress" into
every getblocktemplate request. No pool-side code change is required. Every
other method, and batch requests, pass through byte-for-byte.

The pool must still place its payout output at vout[0] and must pay the same
address configured here. The proxy cannot fix output ordering.

Stdlib only. Python 3.6+.
"""

import argparse
import base64
import json
import logging
import sys
import urllib.error
import urllib.parse
import urllib.request
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

LOG = logging.getLogger("krgn-gbt-proxy")

CONFIG = {
    "upstream_url": None,
    "upstream_auth": None,   # "Basic ..." built from --upstream credentials, or None
    "pooladdress": None,
    "timeout": 30,
}


def patch_gbt(req):
    """Inject pooladdress into one getblocktemplate request object.

    Returns (injected, overridden, previous_value). Mutates req in place.
    Any rules/capabilities/longpollid the pool sent are preserved untouched.
    """
    if not isinstance(req, dict) or req.get("method") != "getblocktemplate":
        return (False, False, None)

    params = req.get("params")
    if params is None or (isinstance(params, list) and len(params) == 0):
        params = [{}]
        req["params"] = params
    elif not isinstance(params, list):
        LOG.warning("getblocktemplate params is %s, not a list; passing through untouched",
                    type(params).__name__)
        return (False, False, None)

    tmpl = params[0]
    if tmpl is None:
        tmpl = {}
        params[0] = tmpl
    if not isinstance(tmpl, dict):
        LOG.warning("getblocktemplate params[0] is %s, not an object; passing through "
                    "untouched (blocks will be REJECTED post-fork if the coinbase is rebuilt)",
                    type(tmpl).__name__)
        return (False, False, None)

    prev = tmpl.get("pooladdress")
    tmpl["pooladdress"] = CONFIG["pooladdress"]
    if prev is None:
        return (True, False, None)
    if prev == CONFIG["pooladdress"]:
        return (False, False, prev)
    # Deliberate override. The signed script must match the script the pool
    # actually pays at vout[0]; this proxy is configured with that address, so
    # its value wins. A mismatch means the pool config and the proxy config
    # disagree and someone must fix one of them.
    return (True, True, prev)


def response_height(resp_obj, req_id):
    """Best-effort block height out of a (possibly batched) GBT response."""
    items = resp_obj if isinstance(resp_obj, list) else [resp_obj]
    for it in items:
        if not isinstance(it, dict):
            continue
        if req_id is not None and it.get("id") != req_id:
            continue
        res = it.get("result")
        if isinstance(res, dict) and "height" in res:
            return res["height"]
        if it.get("error"):
            return "err:%s" % (it["error"].get("message", it["error"]),)
    return "?"


class Handler(BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"
    server_version = "krgn-gbt-proxy"
    sys_version = ""

    def log_message(self, fmt, *args):
        pass  # access logging is handled explicitly below

    def _send(self, code, body, ctype="application/json"):
        self.send_response(code)
        self.send_header("Content-Type", ctype)
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def do_GET(self):
        self._send(405, b'{"result":null,"error":{"code":-32600,'
                        b'"message":"POST only"},"id":null}')

    def do_POST(self):
        try:
            length = int(self.headers.get("Content-Length") or 0)
        except ValueError:
            length = 0
        raw = self.rfile.read(length) if length else b""

        # Prefer the pool's own Authorization header; fall back to --upstream creds.
        auth = self.headers.get("Authorization") or CONFIG["upstream_auth"]

        try:
            obj = json.loads(raw.decode("utf-8")) if raw else None
        except (ValueError, UnicodeDecodeError):
            obj = None

        gbt_notes = []
        if obj is not None:
            items = obj if isinstance(obj, list) else [obj]
            for it in items:
                injected, overridden, prev = patch_gbt(it)
                if isinstance(it, dict) and it.get("method") == "getblocktemplate":
                    gbt_notes.append((it.get("id"), injected, overridden, prev))
            body = json.dumps(obj).encode("utf-8")
        else:
            body = raw  # not JSON we understand; forward the original bytes

        headers = {"Content-Type": "application/json"}
        if auth:
            headers["Authorization"] = auth

        req = urllib.request.Request(CONFIG["upstream_url"], data=body,
                                     headers=headers, method="POST")
        try:
            with urllib.request.urlopen(req, timeout=CONFIG["timeout"]) as r:
                status, out = r.status, r.read()
        except urllib.error.HTTPError as e:
            # Upstream error responses (401, 500 with a JSON-RPC error body,
            # ...) are returned to the pool as-is.
            status, out = e.code, e.read()
        except Exception as e:
            LOG.error("upstream %s: %s", CONFIG["upstream_url"], e)
            status = 502
            out = json.dumps({"result": None,
                              "error": {"code": -32603,
                                        "message": "proxy upstream failure: %s" % e},
                              "id": None}).encode("utf-8")

        if gbt_notes:
            try:
                parsed = json.loads(out.decode("utf-8"))
            except Exception:
                parsed = None
            for rid, injected, overridden, prev in gbt_notes:
                h = response_height(parsed, rid) if parsed is not None else "?"
                if overridden:
                    LOG.warning("getblocktemplate height=%s pooladdress=OVERRIDDEN "
                                "(pool sent %s, proxy forced %s)",
                                h, prev, CONFIG["pooladdress"])
                elif injected:
                    LOG.info("getblocktemplate height=%s pooladdress=INJECTED (%s)",
                             h, CONFIG["pooladdress"])
                else:
                    LOG.info("getblocktemplate height=%s pooladdress=already-set (%s)",
                             h, prev if prev is not None else "n/a")

        self._send(status, out)


def parse_upstream(url):
    p = urllib.parse.urlsplit(url)
    if p.scheme not in ("http", "https"):
        raise ValueError("--upstream must be http:// or https://")
    auth = None
    if p.username is not None:
        raw = "%s:%s" % (urllib.parse.unquote(p.username),
                         urllib.parse.unquote(p.password or ""))
        auth = "Basic " + base64.b64encode(raw.encode("utf-8")).decode("ascii")
    netloc = p.hostname or "127.0.0.1"
    if ":" in netloc:
        netloc = "[%s]" % netloc
    if p.port:
        netloc += ":%d" % p.port
    clean = urllib.parse.urlunsplit((p.scheme, netloc, p.path or "/", p.query, ""))
    return clean, auth


def main():
    ap = argparse.ArgumentParser(
        description="Inject pooladdress into getblocktemplate for Kerrigan v1.3.2+.")
    ap.add_argument("--listen", default="127.0.0.1:7199",
                    help="host:port to listen on (default 127.0.0.1:7199)")
    ap.add_argument("--upstream", required=True,
                    help="http://user:pass@host:port of kerrigand")
    ap.add_argument("--pooladdress", required=True,
                    help="the address the pool actually pays itself at vout[0]")
    ap.add_argument("--timeout", type=float, default=30.0)
    ap.add_argument("--quiet", action="store_true", help="log warnings and errors only")
    args = ap.parse_args()

    logging.basicConfig(level=logging.WARNING if args.quiet else logging.INFO,
                        format="%(asctime)s %(levelname)s %(message)s")

    host, _, port = args.listen.rpartition(":")
    if not host or not port.isdigit():
        ap.error("--listen must be host:port")

    CONFIG["upstream_url"], CONFIG["upstream_auth"] = parse_upstream(args.upstream)
    CONFIG["pooladdress"] = args.pooladdress
    CONFIG["timeout"] = args.timeout

    srv = ThreadingHTTPServer((host, int(port)), Handler)
    srv.daemon_threads = True
    LOG.info("listening on %s:%s -> %s  pooladdress=%s",
             host, port, CONFIG["upstream_url"], CONFIG["pooladdress"])
    try:
        srv.serve_forever()
    except KeyboardInterrupt:
        pass
    finally:
        srv.server_close()


if __name__ == "__main__":
    sys.exit(main())
