#!/usr/bin/env python3
"""Load test for the ferrix_bgp web API: N students announce at the same moment.

Each simulated student owns one prefix (student n gets the n-th /48, numbered
in hex like the lab's route6 objects) and announces it under its own owner
name. All students are released together by a barrier, so the announcements
really are concurrent. The script then checks that the daemon did the right
thing, not merely that it answered:

  1. every announcement was accepted, with latency percentiles
  2. the local table holds exactly these routes, under the right owners
  3. every route was turned into an UPDATE for each established session
  4. the event stream delivered one "announced" event per route, none lost
  5. ownership holds under load: nobody could replace another's prefix
  6. concurrent withdrawal leaves nothing behind
  7. the daemon's own counters show the matching number of BGP UPDATEs sent,
     for the announcements and again for the withdrawals

SAFETY. Pointed at a router that peers with the real Internet, this announces
real prefixes. By default every route therefore carries the large community
210036:0:0 ("announce nowhere"), so the production router accepts the routes
but exports them to nobody, and the script withdraws everything it announced
even when it fails half-way. Pass --announce-for-real only if you mean it.

Standard library only. Examples:

  tools/loadtest.py 128
  tools/loadtest.py 128 --url https://lab.example.net --user student
  tools/loadtest.py 64 --rounds 5 --keep
"""

import argparse
import base64
import getpass
import json
import os
import ssl
import statistics
import sys
import threading
import time
import urllib.error
import urllib.request

QUIET_COMMUNITY = "210036:0:0"
OWNER_PREFIX = "loadtest-"


class Api:
    def __init__(self, base, user=None, password=None, insecure=False, timeout=15):
        self.base = base.rstrip("/")
        self.timeout = timeout
        self.headers = {"content-type": "application/json"}
        if user:
            token = base64.b64encode(f"{user}:{password}".encode()).decode()
            self.headers["authorization"] = f"Basic {token}"
        self.ctx = ssl._create_unverified_context() if insecure else None

    def call(self, method, path, body=None):
        """Returns (status, parsed JSON or text, seconds)."""
        data = json.dumps(body).encode() if body is not None else None
        req = urllib.request.Request(self.base + path, data=data, method=method, headers=self.headers)
        start = time.perf_counter()
        try:
            with urllib.request.urlopen(req, timeout=self.timeout, context=self.ctx) as resp:
                status, raw = resp.status, resp.read()
        except urllib.error.HTTPError as e:
            status, raw = e.code, e.read()
        except Exception as e:  # connection refused, timeout, TLS failure, ...
            return 0, f"{type(e).__name__}: {e}", time.perf_counter() - start
        elapsed = time.perf_counter() - start
        try:
            return status, json.loads(raw) if raw else None, elapsed
        except ValueError:
            return status, raw.decode(errors="replace"), elapsed

    def open_stream(self, path):
        req = urllib.request.Request(self.base + path, headers={**self.headers, "accept": "text/event-stream"})
        return urllib.request.urlopen(req, timeout=self.timeout, context=self.ctx)


def updates_sent(api):
    """BGP UPDATE messages the daemon has sent so far, from /metricz; None if unavailable."""
    status, text, _ = api.call("GET", "/metricz")
    if status != 200 or not isinstance(text, str):
        return None
    for line in text.splitlines():
        if line.startswith('bgp_messages_sent_total{type="update"}'):
            return int(float(line.rsplit(" ", 1)[1]))
    return 0  # the counter has no sample until the first UPDATE


def established(api):
    status, peers, _ = api.call("GET", "/api/peers")
    return [p for p in peers if p.get("state") == "Established"] if status == 200 else []


def check_wire(api, report, before, routes, sessions_before, what):
    """The advertiser works asynchronously, so give the counter a moment to settle."""
    if before is None:
        print(f"  [skip] /metricz unavailable, cannot count UPDATEs for the {what}")
        return
    expected = routes * len(sessions_before)
    deadline = time.time() + 5
    while time.time() < deadline and (updates_sent(api) or 0) - before < expected:
        time.sleep(0.1)
    delta = (updates_sent(api) or 0) - before
    if len(established(api)) != len(sessions_before):
        print(f"  [skip] sessions changed during the test; {delta} UPDATEs sent for the {what}")
    else:
        report.check(delta == expected, f"{delta} BGP UPDATEs sent for the {what} "
                     f"({routes} routes x {len(sessions_before)} session(s))", [f"expected {expected}"])


class EventCounter(threading.Thread):
    """Follows /api/events and records what it sees until stopped."""

    def __init__(self, api):
        super().__init__(daemon=True)
        self.api, self.events, self.error = api, [], None
        self.connected, self._stop = threading.Event(), threading.Event()

    def run(self):
        try:
            stream = self.api.open_stream("/api/events")
        except Exception as e:
            self.error = f"{type(e).__name__}: {e}"
            self.connected.set()
            return
        self.connected.set()
        try:
            for raw in stream:
                if self._stop.is_set():
                    break
                line = raw.decode(errors="replace").strip()
                if line.startswith("data:"):
                    try:
                        self.events.append(json.loads(line[5:]))
                    except ValueError:
                        pass
        except Exception as e:
            if not self._stop.is_set():
                self.error = f"{type(e).__name__}: {e}"

    def stop(self):
        self._stop.set()


def percentile(values, p):
    ordered = sorted(values)
    return ordered[min(len(ordered) - 1, int(round(p / 100 * (len(ordered) - 1))))]


def latency_line(latencies):
    if not latencies:
        return "no samples"
    ms = [v * 1000 for v in latencies]
    return (f"min {min(ms):.0f} ms, median {statistics.median(ms):.0f} ms, "
            f"p95 {percentile(ms, 95):.0f} ms, max {max(ms):.0f} ms")


def run_concurrently(jobs, worker):
    """Runs worker(job) for every job at the same instant; returns results in order."""
    results = [None] * len(jobs)
    barrier = threading.Barrier(len(jobs))

    def target(i, job):
        try:
            barrier.wait(timeout=30)
        except threading.BrokenBarrierError:
            results[i] = (0, "barrier broken", 0.0)
            return
        results[i] = worker(job)

    threads = [threading.Thread(target=target, args=(i, job)) for i, job in enumerate(jobs)]
    wall = time.perf_counter()
    for t in threads:
        t.start()
    for t in threads:
        t.join()
    return results, time.perf_counter() - wall


class Report:
    def __init__(self):
        self.failures = 0

    def check(self, ok, text, detail=None):
        print(f"  [{'ok' if ok else 'FAIL'}] {text}")
        if not ok:
            self.failures += 1
            for line in (detail or [])[:8]:
                print(f"         {line}")
            if detail and len(detail) > 8:
                print(f"         ... and {len(detail) - 8} more")
        return ok


def one_round(api, args, students, report, round_no):
    body_for = lambda s: {
        "prefix": s["prefix"], "owner": s["owner"],
        "communities": args.community, "large_communities": s["large"],
    }

    events = EventCounter(api)
    if not args.no_events:
        events.start()
        events.connected.wait(timeout=10)
        time.sleep(0.3)  # let the subscription settle before the burst

    # ---- 1. everybody announces at once -----------------------------------
    print(f"\nRound {round_no}: {len(students)} students announce concurrently")
    sent_before, sessions_before = updates_sent(api), established(api)
    results, wall = run_concurrently(students, lambda s: api.call("POST", "/api/routes/local", body_for(s)))
    accepted = [s for s, r in zip(students, results) if r[0] == 200]
    rejected = [f"{s['owner']} {s['prefix']}: HTTP {r[0]} {str(r[1])[:160]}" for s, r in zip(students, results) if r[0] != 200]
    report.check(not rejected, f"all {len(students)} announcements accepted in {wall:.2f} s "
                 f"({len(students) / wall:.0f}/s)", rejected)
    print(f"         latency: {latency_line([r[2] for r in results if r[0] == 200])}")
    if not accepted:
        print("  [skip] nothing was accepted, so there is nothing further to verify")
        if not args.no_events:
            events.stop()
        return accepted

    # ---- 2. the table holds exactly what was announced --------------------
    status, table, _ = api.call("GET", "/api/routes/local")
    table = table if status == 200 and isinstance(table, list) else []
    by_prefix = {r["prefix"]: r for r in table}
    missing = [f"{s['prefix']} missing" for s in accepted if s["prefix"] not in by_prefix]
    wrong_owner = [f"{s['prefix']} owned by {by_prefix[s['prefix']]['owner']!r}, expected {s['owner']!r}"
                   for s in accepted if s["prefix"] in by_prefix and by_prefix[s["prefix"]]["owner"] != s["owner"]]
    wrong_attrs = [f"{s['prefix']} large communities {by_prefix[s['prefix']]['large_communities']}"
                   for s in accepted if s["prefix"] in by_prefix
                   and sorted(by_prefix[s["prefix"]]["large_communities"]) != sorted(s["large"])]
    report.check(not (missing or wrong_owner or wrong_attrs),
                 "local table holds every route with the right owner and communities",
                 missing + wrong_owner + wrong_attrs)

    # ---- 3. each route became an UPDATE on every session -------------------
    up = established(api)
    if not up:
        print("  [skip] no established BGP session, so nothing was put on the wire")
    else:
        sample = accepted if len(accepted) <= 32 else accepted[:: max(1, len(accepted) // 32)]
        bad = []
        for s in sample:
            code, wire, _ = api.call("GET", "/api/routes/local/wire?prefix=" + urllib.request.quote(s["prefix"], safe=""))
            if code != 200 or not wire:
                bad.append(f"{s['prefix']}: HTTP {code}, {len(wire) if isinstance(wire, list) else 0} session(s)")
            elif not all(w["hex"].startswith("ff" * 16) for w in wire):
                bad.append(f"{s['prefix']}: malformed UPDATE")
        report.check(not bad, f"UPDATE built for {len(up)} established session(s) "
                     f"(checked {len(sample)} of {len(accepted)} routes)", bad)
        check_wire(api, report, sent_before, len(accepted), sessions_before, "announcements")

    # ---- 4. the event stream lost nothing ----------------------------------
    if not args.no_events:
        deadline = time.time() + 5
        want = {s["prefix"] for s in accepted}
        seen = lambda: {e["route"]["prefix"] for e in events.events if e.get("type") == "local_announced"}
        while time.time() < deadline and not want <= seen():
            time.sleep(0.1)
        lagged = [e for e in events.events if e.get("type") == "lagged"]
        lost = sorted(want - seen())
        detail = [f"no event for {p}" for p in lost] + [f"stream lagged: missed {e['missed']} on {e['stream']}" for e in lagged]
        if events.error:
            detail.insert(0, f"event stream: {events.error}")
        report.check(not detail, f"event stream delivered all {len(want)} announcements", detail)

    # ---- 5. ownership holds under load -------------------------------------
    if len(accepted) >= 2 and not args.no_policy:
        # Everyone tries, at once, to take over their neighbour's prefix.
        thieves = [{**accepted[i], "prefix": accepted[(i + 1) % len(accepted)]["prefix"]} for i in range(len(accepted))]
        results, _ = run_concurrently(thieves, lambda s: api.call("POST", "/api/routes/local", body_for(s)))
        leaked = [f"{t['owner']} took {t['prefix']}: HTTP {r[0]}" for t, r in zip(thieves, results) if r[0] != 409]
        status, table, _ = api.call("GET", "/api/routes/local")
        owners = {r["prefix"]: r["owner"] for r in table} if status == 200 else {}
        leaked += [f"{s['prefix']} now owned by {owners.get(s['prefix'])!r}" for s in accepted if owners.get(s["prefix"]) != s["owner"]]
        report.check(not leaked, f"{len(thieves)} concurrent takeover attempts all refused (409), owners unchanged", leaked)

    if not args.no_events:
        events.stop()
    return accepted


def withdraw_all(api, students, report, verify=True):
    if not students:
        return
    sent_before, sessions_before = (updates_sent(api), established(api)) if verify else (None, [])
    results, wall = run_concurrently(
        students, lambda s: api.call("DELETE", "/api/routes/local", {"prefix": s["prefix"], "owner": s["owner"]}))
    failed = [f"{s['prefix']}: HTTP {r[0]} {r[1]}" for s, r in zip(students, results) if r[0] != 200]
    if not verify:
        return
    report.check(not failed, f"all {len(students)} withdrawals accepted in {wall:.2f} s", failed)
    status, table, _ = api.call("GET", "/api/routes/local")
    left = [r["prefix"] for r in table if r["owner"].startswith(OWNER_PREFIX)] if status == 200 else ["(could not read table)"]
    report.check(not left, "nothing left behind by the load test", left)
    if sessions_before:
        check_wire(api, report, sent_before, len(students) - len(failed), sessions_before, "withdrawals")


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("students", type=int, help="number of simulated students")
    ap.add_argument("--url", default="http://127.0.0.1:8080", help="base URL of the status page (default %(default)s)")
    ap.add_argument("--user", help="HTTP basic auth user; the password is read from $FERRIX_PASSWORD or prompted for")
    ap.add_argument("--insecure", action="store_true", help="do not verify the TLS certificate")
    ap.add_argument("--prefix-template", default="2a0d:d741:{n:x}::/48",
                    help="prefix of student n, a Python format string (default %(default)s)")
    ap.add_argument("--first", type=int, default=1, help="number of the first student (default 1)")
    ap.add_argument("--rounds", type=int, default=1, help="repeat announce/verify/withdraw this many times")
    ap.add_argument("--community", action="append", default=[], help="standard community to attach (repeatable)")
    ap.add_argument("--large-community", action="append", default=[], help="extra large community to attach (repeatable)")
    ap.add_argument("--announce-for-real", action="store_true",
                    help=f"do NOT add {QUIET_COMMUNITY}; routes may be exported to the Internet")
    ap.add_argument("--keep", action="store_true", help="leave the routes announced at the end")
    ap.add_argument("--no-events", action="store_true", help="skip the event stream check")
    ap.add_argument("--no-policy", action="store_true", help="skip the concurrent takeover check")
    args = ap.parse_args()
    if args.students < 1:
        ap.error("need at least one student")

    password = None
    if args.user:
        password = os.environ.get("FERRIX_PASSWORD") or getpass.getpass(f"Password for {args.user}: ")
    api = Api(args.url, args.user, password, args.insecure)

    large = list(args.large_community)
    if not args.announce_for_real and QUIET_COMMUNITY not in large:
        large.insert(0, QUIET_COMMUNITY)
    students = [{"owner": f"{OWNER_PREFIX}{n}", "prefix": args.prefix_template.format(n=n), "large": large}
                for n in range(args.first, args.first + args.students)]

    # ---- preflight ----------------------------------------------------------
    status, info, _ = api.call("GET", "/api/status")
    if status == 401:
        sys.exit(f"{args.url} wants a password: HTTP 401. Pass --user, and check the password.")
    if status != 200:
        sys.exit(f"Cannot reach {args.url}/api/status: HTTP {status} {str(info)[:200]}")
    print(f"Target: {args.url}  router-id {info['router_id']}  AS{info['local_asn']}  version {info['version']}  "
          f"on {info.get('hostname', '?')}")
    print(f"        {info['sessions']} established session(s), {info['local_routes']} local route(s) already present")
    print(f"Prefixes {students[0]['prefix']} .. {students[-1]['prefix']}, large communities {large or 'none'}")
    if args.announce_for_real:
        print("WARNING: --announce-for-real: these routes may be exported to the Internet.")

    status, table, _ = api.call("GET", "/api/routes/local")
    stale = [r for r in table if r["owner"].startswith(OWNER_PREFIX)] if status == 200 else []
    if stale:
        print(f"Removing {len(stale)} route(s) left over from an earlier load test")
        withdraw_all(api, [{"owner": r["owner"], "prefix": r["prefix"]} for r in stale], Report(), verify=False)
    taken = {r["prefix"]: r["owner"] for r in table if not r["owner"].startswith(OWNER_PREFIX)} if status == 200 else {}
    clash = [s for s in students if s["prefix"] in taken]
    if clash:
        print(f"Skipping {len(clash)} prefix(es) currently announced by real owners: "
              + ", ".join(f"{s['prefix']} ({taken[s['prefix']]})" for s in clash[:5]) + (" ..." if len(clash) > 5 else ""))
        students = [s for s in students if s["prefix"] not in taken]
    if not students:
        sys.exit("Nothing left to test.")

    report, announced = Report(), []
    try:
        for round_no in range(1, args.rounds + 1):
            announced = one_round(api, args, students, report, round_no)
            if round_no < args.rounds or not args.keep:
                withdraw_all(api, announced, report)
                announced = []
    except KeyboardInterrupt:
        print("\nInterrupted.")
        report.failures += 1
    finally:
        if announced and not args.keep:
            withdraw_all(api, announced, Report(), verify=False)

    status, info, _ = api.call("GET", "/api/status")
    if status == 200:
        print(f"\nAfterwards: {info['sessions']} established session(s), {info['local_routes']} local route(s)")
    print("RESULT: " + ("PASS" if report.failures == 0 else f"FAIL ({report.failures} check(s) failed)"))
    sys.exit(0 if report.failures == 0 else 1)


if __name__ == "__main__":
    main()