Naut-Peer/harness/harness.py

"""
Test/driver harness for the C peer.

The Python harness parses .torrent metadata and can discover peers through the
sibling torrent-tracker C library's DHT and HTTP/UDP tracker helpers. The C peer
does the fast part: pull the requested blocks from one peer. This harness owns
what to download, reassembles pieces, and verifies SHA-1 hashes -- the peer
never hashes or persists anything.
"""
from __future__ import annotations

import argparse
import hashlib
import os
import secrets
import time

from peer_ffi import Peer, PeerConfig, STATE_ERROR, STATE_NAMES, ERROR_NAMES
from torrent_meta import Metadata, load_metadata, load_torrent
from tracker_ffi import DHTClient, TrackerClient


def make_peer_id() -> bytes:
    return b"-PC0001-" + secrets.token_bytes(12)


def discover_peers(torrent_path: str, max_wait: float = 20.0) -> list[tuple[str, int]]:
    """Discover peer endpoints through DHT first, then torrent trackers."""
    tf = load_torrent(torrent_path)
    peer_id = make_peer_id()
    key = secrets.randbits(32)
    deadline = time.time() + max_wait
    seen: set[tuple[str, int]] = set()

    dht_budget = max(0.5, min(6.0, max_wait / 2.0))
    dht_result = DHTClient().lookup(tf.metadata.info_hash, timeout=dht_budget)
    seen.update(dht_result.peers)
    if seen:
        return sorted(seen)

    client = TrackerClient()
    for url in tf.trackers:
        if time.time() >= deadline:
            break
        result = client.announce(
            url, tf.metadata, peer_id, port=6881, key=key, numwant=50,
            event="started", timeout=max(0.5, deadline - time.time()))
        if result.ok:
            seen.update(result.peers)
        if seen:
            break
    return sorted(seen)


class Downloader:
    def __init__(self, meta: Metadata, *, peer_id: bytes | None = None,
                 num_slots: int = 0, max_pipeline: int = 0,
                 request_timeout_ms: int = 0, recv_buffer_bytes: int = 0,
                 lib_path: str | None = None):
        self.meta = meta
        cfg = PeerConfig()
        cfg.info_hash[:] = meta.info_hash
        cfg.peer_id[:] = peer_id or make_peer_id()
        cfg.piece_length = meta.piece_length
        cfg.total_size = meta.total_size
        cfg.num_pieces = meta.num_pieces
        cfg.num_slots = num_slots
        cfg.max_pipeline = max_pipeline
        cfg.request_timeout_ms = request_timeout_ms
        cfg.recv_buffer_bytes = recv_buffer_bytes
        self.peer = Peer(cfg, lib_path)

        self.buffers: list[bytearray | None] = [None] * meta.num_pieces
        self.received = [0] * meta.num_pieces
        self.done = bytearray(meta.num_pieces)
        self.done_count = 0

    def download(self, ip: str, port: int, pieces=None, priorities=None,
                 timeout: float = 60.0, progress_every: float = 1.0,
                 output: str | None = None) -> bytes | None:
        meta = self.meta
        pieces = list(pieces) if pieces is not None else list(range(meta.num_pieces))

        out_fh = open(output, "wb") if output else None
        if out_fh:
            out_fh.truncate(meta.total_size)

        # Build the priority vector: caller-supplied scheme, or uniform "1" over
        # the wanted pieces (0 = not wanted). The peer masks this with what the
        # remote actually has, so a peer missing a piece is simply skipped.
        prio = bytearray(priorities) if priorities is not None \
            else bytearray(meta.num_pieces)
        for i in pieces:
            self.buffers[i] = bytearray(meta.piece_len(i))
            if priorities is None:
                prio[i] = 1

        self.peer.start(ip, port)
        self.peer.set_priorities(prio)

        want = len(pieces)
        deadline = time.time() + timeout
        last_print = 0.0
        last_progress_bytes = 0
        last_progress_time = time.time()

        while self.done_count < want:
            st = self.peer.status()
            if st.state == STATE_ERROR:
                raise RuntimeError(f"peer error: {ERROR_NAMES[st.error]}")

            descs = self.peer.poll_ready()
            if not descs:
                self.peer.wait(100)
                now = time.time()
                if st.bytes_received != last_progress_bytes:
                    last_progress_bytes = st.bytes_received
                    last_progress_time = now
                if now > deadline and now - last_progress_time > timeout:
                    raise TimeoutError(
                        f"stalled: {self.done_count}/{want} pieces, "
                        f"state={STATE_NAMES[st.state]}")
                if now - last_print >= progress_every:
                    last_print = now
                    print(f"  {self.done_count}/{want} pieces  "
                          f"{st.rate_bps/1e6:.1f} MB/s  "
                          f"outstanding={st.outstanding} free={st.free_slots}")
                continue

            for d in descs:
                buf = self.buffers[d.piece]
                buf[d.begin:d.begin + d.len] = self.peer.block_data(d.slot, d.len)
                self.peer.release(d.slot)
                self.received[d.piece] += d.len
                if (not self.done[d.piece]
                        and self.received[d.piece] >= meta.piece_len(d.piece)):
                    digest = hashlib.sha1(bytes(buf)).digest()
                    if digest != meta.piece_hashes[d.piece]:
                        raise ValueError(f"piece {d.piece} hash mismatch")
                    self.done[d.piece] = 1
                    self.done_count += 1
                    # Drop priority so a verified piece is no longer a
                    # selection candidate (the peer also won't re-request it).
                    self.peer.set_priority(d.piece, 0)
                    if out_fh:
                        out_fh.seek(d.piece * meta.piece_length)
                        out_fh.write(buf)
                    if not output:
                        pass  # keep in memory for return
                    else:
                        self.buffers[d.piece] = None  # free once flushed

        if out_fh:
            out_fh.close()
            return None
        return b"".join(bytes(self.buffers[i]) for i in pieces)

    def close(self):
        self.peer.stop()
        self.peer.close()


def main() -> int:
    ap = argparse.ArgumentParser(description="Drive the C peer to download a torrent.")
    ap.add_argument("torrent", help="path to .torrent file")
    ap.add_argument("--peer", help="explicit peer ip:port (skip tracker)")
    ap.add_argument("--output", "-o", help="write downloaded data here")
    ap.add_argument("--slots", type=int, default=0, help="arena slots (16 KiB each)")
    ap.add_argument("--pipeline", type=int, default=0, help="max outstanding requests")
    ap.add_argument("--timeout", type=float, default=120.0)
    args = ap.parse_args()

    meta = load_metadata(args.torrent)
    print(f"torrent: {meta.name}  {meta.total_size} bytes  "
          f"{meta.num_pieces} pieces x {meta.piece_length}")

    if args.peer:
        host, port = args.peer.rsplit(":", 1)
        endpoints = [(host, int(port))]
    else:
        print("discovering peers via tracker/DHT...")
        endpoints = discover_peers(args.torrent)
        if not endpoints:
            print("no peers found")
            return 1
        print(f"found {len(endpoints)} peer(s); using {endpoints[0]}")

    dl = Downloader(meta, num_slots=args.slots, max_pipeline=args.pipeline)
    try:
        ip, port = endpoints[0]
        t0 = time.time()
        dl.download(ip, port, timeout=args.timeout, output=args.output)
        dt = time.time() - t0
        mb = meta.total_size / 1e6
        print(f"done: {mb:.1f} MB in {dt:.2f}s = {mb/dt:.1f} MB/s")
    finally:
        dl.close()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())