Naut-Peer/harness/tracker_ffi.py

from __future__ import annotations

import ctypes as C
import os
import random
import socket
import ssl
import struct
import time
from dataclasses import dataclass
from urllib.parse import urlsplit, urlunsplit
from urllib.request import Request, urlopen

TRACKER_OK = 0
TRACKER_EVENT_NONE = 0
TRACKER_EVENT_COMPLETED = 1
TRACKER_EVENT_STARTED = 2
TRACKER_EVENT_STOPPED = 3
TRACKER_ADDR_IPV4 = 4
TRACKER_ADDR_IPV6 = 6
TRACKER_MAX_PEERS = 256
TRACKER_MAX_URL_DATA = 512
DHT_MAX_TRANSACTION = 16
DHT_MAX_TOKEN = 64
DHT_MAX_NODES = 256
DHT_MAX_ERROR = 128
DHT_MSG_QUERY = 1
DHT_MSG_RESPONSE = 2
DHT_QUERY_GET_PEERS = 3
DEFAULT_DHT_BOOTSTRAP = (
    ("router.bittorrent.com", 6881),
    ("dht.transmissionbt.com", 6881),
    ("router.utorrent.com", 6881),
)


class TrackerPeer(C.Structure):
    _fields_ = [
        ("family", C.c_uint8),
        ("addr", C.c_uint8 * 16),
        ("port", C.c_uint16),
        ("peer_id", C.c_uint8 * 20),
        ("has_peer_id", C.c_uint8),
    ]


class TrackerAnnounceRequest(C.Structure):
    _fields_ = [
        ("info_hash", C.c_uint8 * 20),
        ("peer_id", C.c_uint8 * 20),
        ("port", C.c_uint16),
        ("uploaded", C.c_uint64),
        ("downloaded", C.c_uint64),
        ("left", C.c_uint64),
        ("numwant", C.c_int32),
        ("key", C.c_uint32),
        ("ip4", C.c_uint32),
        ("event", C.c_int),
        ("compact", C.c_uint8),
        ("no_peer_id", C.c_uint8),
        ("has_key", C.c_uint8),
        ("has_ip4", C.c_uint8),
        ("ip", C.c_char * 64),
        ("tracker_id", C.c_char * 128),
        ("url_data", C.c_char * TRACKER_MAX_URL_DATA),
    ]


class TrackerAnnounceResponse(C.Structure):
    _fields_ = [
        ("interval", C.c_uint32),
        ("min_interval", C.c_uint32),
        ("complete", C.c_uint32),
        ("incomplete", C.c_uint32),
        ("tracker_id", C.c_char_p),
        ("peers", C.POINTER(TrackerPeer)),
        ("peer_count", C.c_size_t),
        ("compact", C.c_uint8),
    ]


class DHTNode(C.Structure):
    _fields_ = [
        ("id", C.c_uint8 * 20),
        ("family", C.c_uint8),
        ("addr", C.c_uint8 * 16),
        ("port", C.c_uint16),
    ]


class DHTMessage(C.Structure):
    _fields_ = [
        ("type", C.c_int),
        ("query", C.c_int),
        ("transaction", C.c_uint8 * DHT_MAX_TRANSACTION),
        ("transaction_len", C.c_size_t),
        ("id", C.c_uint8 * 20),
        ("target", C.c_uint8 * 20),
        ("info_hash", C.c_uint8 * 20),
        ("port", C.c_uint16),
        ("implied_port", C.c_uint8),
        ("want_ipv4", C.c_uint8),
        ("want_ipv6", C.c_uint8),
        ("token", C.c_uint8 * DHT_MAX_TOKEN),
        ("token_len", C.c_size_t),
        ("nodes", DHTNode * DHT_MAX_NODES),
        ("node_count", C.c_size_t),
        ("peers", TrackerPeer * TRACKER_MAX_PEERS),
        ("peer_count", C.c_size_t),
        ("error_code", C.c_int),
        ("error_message", C.c_char * DHT_MAX_ERROR),
    ]


@dataclass
class TrackerResult:
    tracker: str
    ok: bool
    protocol: str
    peers: list[tuple[str, int]]
    interval: int = 0
    complete: int = 0
    incomplete: int = 0
    error: str = ""
    elapsed_ms: float = 0.0


@dataclass
class DHTResult:
    peers: list[tuple[str, int]]
    nodes_queried: int
    nodes_discovered: int
    elapsed_ms: float
    error: str = ""


def _default_lib_path() -> str:
    env = os.environ.get("TORRENT_TRACKER_LIB")
    if env:
        return env
    root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    candidates = [
        os.path.join(root, "..", "torrent-tracker", "build", "libtorrenttracker.so"),
        os.path.join(root, "build", "libtorrenttracker.so"),
    ]
    for path in candidates:
        if os.path.exists(path):
            return os.path.abspath(path)
    return os.path.abspath(candidates[0])


def _load(lib_path: str | None = None) -> C.CDLL:
    lib = C.CDLL(lib_path or _default_lib_path())
    lib.tracker_http_write_announce_query.restype = C.c_int
    lib.tracker_http_write_announce_query.argtypes = [
        C.POINTER(TrackerAnnounceRequest), C.c_char_p, C.c_size_t, C.POINTER(C.c_size_t)]
    lib.tracker_http_parse_announce_response.restype = C.c_int
    lib.tracker_http_parse_announce_response.argtypes = [
        C.c_void_p, C.c_size_t, C.POINTER(TrackerPeer), C.c_size_t,
        C.POINTER(TrackerAnnounceResponse)]
    lib.tracker_udp_write_connect_request.restype = C.c_int
    lib.tracker_udp_write_connect_request.argtypes = [
        C.c_uint32, C.c_void_p, C.c_size_t, C.POINTER(C.c_size_t)]
    lib.tracker_udp_parse_connect_response.restype = C.c_int
    lib.tracker_udp_parse_connect_response.argtypes = [
        C.c_void_p, C.c_size_t, C.c_uint32, C.POINTER(C.c_uint64)]
    lib.tracker_udp_write_announce_request.restype = C.c_int
    lib.tracker_udp_write_announce_request.argtypes = [
        C.c_uint64, C.c_uint32, C.POINTER(TrackerAnnounceRequest),
        C.c_void_p, C.c_size_t, C.POINTER(C.c_size_t)]
    lib.tracker_udp_parse_announce_response.restype = C.c_int
    lib.tracker_udp_parse_announce_response.argtypes = [
        C.c_void_p, C.c_size_t, C.c_uint32, C.c_int, C.POINTER(TrackerPeer),
        C.c_size_t, C.POINTER(TrackerAnnounceResponse)]
    lib.dht_write_get_peers_query.restype = C.c_int
    lib.dht_write_get_peers_query.argtypes = [
        C.c_void_p, C.c_size_t, C.c_void_p, C.c_void_p, C.c_uint8, C.c_uint8,
        C.c_void_p, C.c_size_t, C.POINTER(C.c_size_t)]
    lib.dht_parse_message.restype = C.c_int
    lib.dht_parse_message.argtypes = [C.c_void_p, C.c_size_t, C.POINTER(DHTMessage)]
    lib.dht_write_peers_response.restype = C.c_int
    lib.dht_write_peers_response.argtypes = [
        C.c_void_p, C.c_size_t, C.c_void_p, C.c_void_p, C.c_size_t,
        C.POINTER(TrackerPeer), C.c_size_t, C.c_void_p, C.c_size_t,
        C.POINTER(C.c_size_t)]
    return lib


class TrackerClient:
    def __init__(self, lib_path: str | None = None):
        self.lib = _load(lib_path)

    def _request(self, meta, peer_id: bytes, port: int, key: int, numwant: int,
                 event: str) -> TrackerAnnounceRequest:
        req = TrackerAnnounceRequest()
        C.memset(C.byref(req), 0, C.sizeof(req))
        req.info_hash[:] = meta.info_hash
        req.peer_id[:] = peer_id
        req.port = port
        req.left = meta.total_size
        req.numwant = numwant
        req.key = key
        req.has_key = 1
        req.compact = 1
        req.no_peer_id = 1
        req.event = {
            "completed": TRACKER_EVENT_COMPLETED,
            "started": TRACKER_EVENT_STARTED,
            "stopped": TRACKER_EVENT_STOPPED,
        }.get(event, TRACKER_EVENT_NONE)
        return req

    @staticmethod
    def _peers(peers, count: int) -> list[tuple[str, int]]:
        out: list[tuple[str, int]] = []
        for i in range(count):
            p = peers[i]
            if p.family == TRACKER_ADDR_IPV4:
                host = socket.inet_ntop(socket.AF_INET, bytes(p.addr[:4]))
            elif p.family == TRACKER_ADDR_IPV6:
                host = socket.inet_ntop(socket.AF_INET6, bytes(p.addr[:16]))
            else:
                continue
            out.append((host, int(p.port)))
        return out

    def announce_http(self, url: str, meta, peer_id: bytes, port: int, key: int,
                      numwant: int, event: str, timeout: float) -> TrackerResult:
        start = time.monotonic()
        try:
            req = self._request(meta, peer_id, port, key, numwant, event)
            query = C.create_string_buffer(2048)
            written = C.c_size_t()
            rc = self.lib.tracker_http_write_announce_query(
                C.byref(req), query, C.sizeof(query), C.byref(written))
            if rc != TRACKER_OK:
                raise RuntimeError(f"tracker_http_write_announce_query failed: {rc}")

            parts = urlsplit(url)
            q = parts.query
            suffix = query.value.decode("ascii")
            q = f"{q}&{suffix}" if q else suffix
            announce_url = urlunsplit((parts.scheme, parts.netloc, parts.path, q,
                                       parts.fragment))
            request = Request(announce_url,
                              headers={"User-Agent": "torrent-peer/0.1"})
            ctx = ssl.create_default_context()
            with urlopen(request, timeout=timeout, context=ctx) as resp:
                raw = resp.read(2 * 1024 * 1024)

            raw_buf = C.create_string_buffer(raw, len(raw))
            out_peers = (TrackerPeer * TRACKER_MAX_PEERS)()
            parsed = TrackerAnnounceResponse()
            rc = self.lib.tracker_http_parse_announce_response(
                raw_buf, len(raw), out_peers, TRACKER_MAX_PEERS, C.byref(parsed))
            if rc != TRACKER_OK:
                raise RuntimeError(f"tracker_http_parse_announce_response failed: {rc}")
            return TrackerResult(
                tracker=url,
                ok=True,
                protocol=parts.scheme,
                peers=self._peers(out_peers, parsed.peer_count),
                interval=int(parsed.interval),
                complete=int(parsed.complete),
                incomplete=int(parsed.incomplete),
                elapsed_ms=(time.monotonic() - start) * 1000.0,
            )
        except Exception as exc:
            return TrackerResult(url, False, "http", [], error=str(exc),
                                 elapsed_ms=(time.monotonic() - start) * 1000.0)

    @staticmethod
    def _url_data(url: str) -> bytes:
        parts = urlsplit(url)
        data = (parts.path or "").encode("utf-8")
        if parts.query:
            data += b"?" + parts.query.encode("utf-8")
        return data[:TRACKER_MAX_URL_DATA - 1]

    @staticmethod
    def _roundtrip(sock: socket.socket, packet: bytes, txid: int,
                   timeout: float) -> bytes:
        deadline = time.monotonic() + timeout
        delay = min(timeout, 1.0)
        while True:
            sock.send(packet)
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                raise TimeoutError("UDP tracker timed out")
            sock.settimeout(min(delay, remaining))
            try:
                raw = sock.recv(65535)
            except socket.timeout:
                delay = min(delay * 2.0, 8.0)
                continue
            if len(raw) >= 8 and struct.unpack_from("!I", raw, 4)[0] == txid:
                return raw

    def announce_udp(self, url: str, meta, peer_id: bytes, port: int, key: int,
                     numwant: int, event: str, timeout: float) -> TrackerResult:
        start = time.monotonic()
        parts = urlsplit(url)
        if not parts.hostname:
            return TrackerResult(url, False, "udp", [], error="missing UDP tracker host")
        tracker_port = parts.port or 80
        try:
            infos = socket.getaddrinfo(parts.hostname, tracker_port, 0,
                                       socket.SOCK_DGRAM)
            last_error: Exception | None = None
            for family, socktype, proto, _canon, sockaddr in infos:
                if family not in (socket.AF_INET, socket.AF_INET6):
                    continue
                try:
                    with socket.socket(family, socktype, proto) as sock:
                        sock.connect(sockaddr)

                        txid = random.getrandbits(32)
                        buf = C.create_string_buffer(2048)
                        written = C.c_size_t()
                        rc = self.lib.tracker_udp_write_connect_request(
                            txid, buf, C.sizeof(buf), C.byref(written))
                        if rc != TRACKER_OK:
                            raise RuntimeError(f"connect request failed: {rc}")
                        raw = self._roundtrip(sock, buf.raw[:written.value], txid, timeout)
                        conn_id = C.c_uint64()
                        raw_buf = C.create_string_buffer(raw, len(raw))
                        rc = self.lib.tracker_udp_parse_connect_response(
                            raw_buf, len(raw), txid, C.byref(conn_id))
                        if rc != TRACKER_OK:
                            raise RuntimeError(f"connect response failed: {rc}")

                        req = self._request(meta, peer_id, port, key, numwant, event)
                        url_data = self._url_data(url)
                        if url_data:
                            req.url_data = url_data
                        txid = random.getrandbits(32)
                        rc = self.lib.tracker_udp_write_announce_request(
                            conn_id.value, txid, C.byref(req), buf, C.sizeof(buf),
                            C.byref(written))
                        if rc != TRACKER_OK:
                            raise RuntimeError(f"announce request failed: {rc}")
                        raw = self._roundtrip(sock, buf.raw[:written.value], txid, timeout)
                        raw_buf = C.create_string_buffer(raw, len(raw))
                        out_peers = (TrackerPeer * TRACKER_MAX_PEERS)()
                        parsed = TrackerAnnounceResponse()
                        tracker_family = (TRACKER_ADDR_IPV6 if family == socket.AF_INET6
                                          else TRACKER_ADDR_IPV4)
                        rc = self.lib.tracker_udp_parse_announce_response(
                            raw_buf, len(raw), txid, tracker_family, out_peers,
                            TRACKER_MAX_PEERS, C.byref(parsed))
                        if rc != TRACKER_OK:
                            raise RuntimeError(f"announce response failed: {rc}")
                        return TrackerResult(
                            tracker=url,
                            ok=True,
                            protocol="udp",
                            peers=self._peers(out_peers, parsed.peer_count),
                            interval=int(parsed.interval),
                            complete=int(parsed.complete),
                            incomplete=int(parsed.incomplete),
                            elapsed_ms=(time.monotonic() - start) * 1000.0,
                        )
                except Exception as exc:
                    last_error = exc
                    continue
            raise RuntimeError(str(last_error or "no usable tracker address"))
        except Exception as exc:
            return TrackerResult(url, False, "udp", [], error=str(exc),
                                 elapsed_ms=(time.monotonic() - start) * 1000.0)

    def announce(self, url: str, meta, peer_id: bytes, port: int, key: int,
                 numwant: int = 50, event: str = "started",
                 timeout: float = 8.0) -> TrackerResult:
        scheme = urlsplit(url).scheme.lower()
        if scheme in ("http", "https"):
            return self.announce_http(url, meta, peer_id, port, key, numwant,
                                      event, timeout)
        if scheme == "udp":
            return self.announce_udp(url, meta, peer_id, port, key, numwant,
                                     event, timeout)
        return TrackerResult(url, False, scheme or "unknown", [],
                             error=f"unsupported tracker scheme {scheme!r}")


class DHTClient:
    def __init__(self, lib_path: str | None = None,
                 bootstrap: tuple[tuple[str, int], ...] = DEFAULT_DHT_BOOTSTRAP):
        self.lib = _load(lib_path)
        self.bootstrap = bootstrap
        self.node_id = os.urandom(20)
        self._tx = random.randrange(1, 0xffff)

    @staticmethod
    def _peer_endpoint(peer: TrackerPeer) -> tuple[str, int] | None:
        if peer.family == TRACKER_ADDR_IPV4:
            host = socket.inet_ntop(socket.AF_INET, bytes(peer.addr[:4]))
        elif peer.family == TRACKER_ADDR_IPV6:
            host = socket.inet_ntop(socket.AF_INET6, bytes(peer.addr[:16]))
        else:
            return None
        return host, int(peer.port)

    @staticmethod
    def _node_endpoint(node: DHTNode) -> tuple[str, int] | None:
        if node.family == TRACKER_ADDR_IPV4:
            host = socket.inet_ntop(socket.AF_INET, bytes(node.addr[:4]))
        elif node.family == TRACKER_ADDR_IPV6:
            host = socket.inet_ntop(socket.AF_INET6, bytes(node.addr[:16]))
        else:
            return None
        return host, int(node.port)

    def _next_tx(self) -> bytes:
        self._tx = (self._tx + 1) & 0xffff
        return self._tx.to_bytes(2, "big")

    def _get_peers_packet(self, info_hash: bytes, tx: bytes) -> bytes:
        buf = C.create_string_buffer(2048)
        written = C.c_size_t()
        tx_buf = C.create_string_buffer(tx, len(tx))
        id_buf = C.create_string_buffer(self.node_id, len(self.node_id))
        hash_buf = C.create_string_buffer(info_hash, len(info_hash))
        rc = self.lib.dht_write_get_peers_query(
            tx_buf, len(tx), id_buf, hash_buf, 1, 1, buf, C.sizeof(buf),
            C.byref(written))
        if rc != TRACKER_OK:
            raise RuntimeError(f"dht_write_get_peers_query failed: {rc}")
        return buf.raw[:written.value]

    def lookup(self, info_hash: bytes, *, timeout: float = 6.0,
               max_queries: int = 32, max_peers: int = 100) -> DHTResult:
        start = time.monotonic()
        deadline = start + timeout
        peers: set[tuple[str, int]] = set()
        queue: list[tuple[str, int]] = list(self.bootstrap)
        seen_nodes: set[tuple[str, int]] = set()
        queried = 0
        discovered = 0
        last_error = ""

        while queue and queried < max_queries and len(peers) < max_peers:
            if time.monotonic() >= deadline:
                break
            host, port = queue.pop(0)
            if (host, port) in seen_nodes:
                continue
            seen_nodes.add((host, port))
            queried += 1

            remaining = max(0.05, deadline - time.monotonic())
            try:
                infos = socket.getaddrinfo(host, port, 0, socket.SOCK_DGRAM)
            except OSError as exc:
                last_error = str(exc)
                continue

            for family, socktype, proto, _canon, sockaddr in infos:
                if family not in (socket.AF_INET, socket.AF_INET6):
                    continue
                tx = self._next_tx()
                packet = self._get_peers_packet(info_hash, tx)
                try:
                    with socket.socket(family, socktype, proto) as sock:
                        sock.settimeout(min(1.0, remaining))
                        sock.sendto(packet, sockaddr)
                        raw, _addr = sock.recvfrom(4096)
                except OSError as exc:
                    last_error = str(exc)
                    continue

                msg = DHTMessage()
                raw_buf = C.create_string_buffer(raw, len(raw))
                rc = self.lib.dht_parse_message(raw_buf, len(raw), C.byref(msg))
                if rc != TRACKER_OK:
                    last_error = f"dht_parse_message failed: {rc}"
                    continue
                got_tx = bytes(msg.transaction[:msg.transaction_len])
                if got_tx != tx or msg.type != DHT_MSG_RESPONSE:
                    continue

                for i in range(int(msg.peer_count)):
                    ep = self._peer_endpoint(msg.peers[i])
                    if ep:
                        peers.add(ep)
                for i in range(int(msg.node_count)):
                    ep = self._node_endpoint(msg.nodes[i])
                    if ep and ep not in seen_nodes and ep not in queue:
                        queue.append(ep)
                        discovered += 1
                break

        return DHTResult(
            peers=sorted(peers),
            nodes_queried=queried,
            nodes_discovered=discovered,
            elapsed_ms=(time.monotonic() - start) * 1000.0,
            error="" if peers else last_error,
        )