#!/usr/bin/env python3
"""Point your protocol at the class challenge server. Given; do not change it.

selfcheck.py tests your code against the receiver, sender and channel that came
in this folder. This points it at a server none of you wrote, over real UDP, on
a path that drops, duplicates and reorders. There are four challenges:

    A-gbn   your Go-Back-N SENDER delivers a payload to the server's receiver
    A-sr    your Selective Repeat SENDER delivers a payload to the server's receiver
    B-gbn   the server sends to your Go-Back-N RECEIVER
    B-sr    the server sends to your Selective Repeat RECEIVER

Each one returns a flag that is yours alone. Run them like this:

    python3 challenge.py A-gbn --student <your id>
    python3 challenge.py B-sr  --student <your id> --server hostname:8422

Whatever happens, this prints a TRACE block. Paste it into FINDINGS.md. The
trace is worth marks on its own, whether or not a flag comes out: it is how you
show what your protocol did. A run that fails with a useful trace is worth more
than no run at all.

If the class server is not answering, keep the error and the state you
reached, and tell the instructor. This package ships no local peer.
"""

import argparse
import hashlib
import importlib
import json
import socket
import sys
import threading
import time

import stopwait

MSS = 1000
PAYLOAD_BYTES = 24000
CHALLENGES = ("A-gbn", "A-sr", "B-gbn", "B-sr")
CONTROL_TIMEOUT = 150       # must outlast the server's session budget, or we give up first


def payload_for(student, nbytes=PAYLOAD_BYTES):
    """The bytes challenge A asks you to deliver. The server derives the same ones."""
    out, counter = b"", 0
    while len(out) < nbytes:
        out += hashlib.sha256(("p3-payload|%s|%d" % (student, counter)).encode()).digest()
        counter += 1
    return out[:nbytes]


class Tee:
    """Your protocol's UDP socket, with a note kept of what went past.

    Your sender and receiver use this exactly as they use a socket. It does
    three things a bare socket does not:

    - it counts what crossed the wire, for the trace;
    - it keeps the FINACK payload instead of discarding it, because in
      challenge A the server puts your flag there. packet.py has always allowed
      a FINACK to carry a payload; nothing else in this project reads it;
    - it will not let a run block forever. A receiver that waits on a packet
      that is never coming would otherwise hang with nothing to show, and the
      trace is the part that is worth marks. Past the deadline, recvfrom raises
      socket.timeout, your code unwinds, and the trace still gets printed.
    """

    def __init__(self, sock, budget=None):
        self._sock = sock
        self._budget = budget
        self._deadline = None if budget is None else time.monotonic() + budget + 5
        self.sent = self.received = 0
        self.finack_payload = None
        self.stopped = False

    def _remaining(self):
        if self._deadline is None:
            return None
        return self._deadline - time.monotonic()

    def sendto(self, raw, peer):
        self.sent += 1
        return self._sock.sendto(raw, peer)

    def recvfrom(self, n):
        left = self._remaining()
        if left is not None and left <= 0:
            raise socket.timeout(
                "the server ended this run" if self.stopped
                else "the run went past its %s second budget" % self._budget)
        raw, peer = self._sock.recvfrom(n)
        self.received += 1
        if len(raw) >= 9 and raw[0] == 3 and len(raw) > 9:      # FINACK with a payload
            self.finack_payload = raw[9:]
        return raw, peer

    def stop(self):
        """End the run wherever it is, even mid-recvfrom.

        Moving the deadline is not enough on its own: a receiver already
        blocked in recvfrom sits there until something arrives. Closing the
        socket from this thread does not reliably wake that call either, so
        instead one byte is sent to our own port. Every protocol in this
        project ignores a packet that fails to decode, so their loop goes round
        once more, and this time the deadline check ends the run.
        """
        self.stopped = True
        self._deadline = time.monotonic()
        try:
            port = self._sock.getsockname()[1]
            self._sock.sendto(b"\x00", ("127.0.0.1", port))
        except OSError:
            pass

    def settimeout(self, t):
        """Whatever your code asks for, never past the deadline."""
        left = self._remaining()
        if left is None:
            return self._sock.settimeout(t)
        capped = max(0.001, left) if t is None else max(0.001, min(t, left))
        return self._sock.settimeout(capped)

    def getsockname(self):
        return self._sock.getsockname()

    def close(self):
        try:
            return self._sock.close()
        except OSError:
            return None


def control(host, port, messages, on_session):
    """Run the control conversation. `on_session` does the UDP half."""
    with socket.create_connection((host, port), timeout=CONTROL_TIMEOUT) as tcp:
        rfile = tcp.makefile("rb")

        def send(obj):
            tcp.sendall((json.dumps(obj) + "\n").encode())

        def recv():
            line = rfile.readline(8192)
            if not line:
                raise SystemExit("the server closed the connection without answering")
            try:
                return json.loads(line)
            except json.JSONDecodeError:
                raise SystemExit("the server sent a line that is not JSON: %r" % line[:120])

        send(messages["hello"])
        session = recv()
        if session.get("msg") == "error":
            raise SystemExit("the server refused this run: %s" % session.get("detail"))
        if session.get("msg") != "session":
            raise SystemExit("expected a session message, got %r" % session)
        result = on_session(session, send, recv)
        return session, result


def watcher(tee, recv):
    """Listen for the server's verdict while your protocol runs.

    The server gives up on a silent peer long before this client would. Without
    this, a receiver waiting on a packet that is never coming sits there until
    its own budget runs out and you see no trace at all. The moment a verdict
    arrives, the socket's deadline moves to now, your code unwinds, and the
    trace gets printed.
    """
    holder = {}

    def listen():
        try:
            holder["result"] = recv()
        except Exception as exc:                   # noqa: BLE001
            holder["error"] = exc
        tee.stop()

    thread = threading.Thread(target=listen, daemon=True)
    thread.start()
    return holder, thread


def settle(holder, thread, seconds=20):
    """The server's verdict, once the watcher has it."""
    thread.join(timeout=seconds)
    if "result" in holder:
        return holder["result"]
    if "error" in holder:
        return {"ok": False, "reason": "no verdict from the server (%s)" % holder["error"]}
    return {"ok": False, "reason": "the server sent no verdict within %d seconds" % seconds}


def run_a(args, module, session, send, recv):
    """Your sender delivers the payload the server asked for."""
    data = payload_for(args.student)
    udp = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    udp.bind(("", 0))
    tee = Tee(udp, budget=session.get("seconds", 90))
    peer = (args.host, session["udp_port"])
    send({"msg": "ready", "udp_port": udp.getsockname()[1]})
    holder, thread = watcher(tee, recv)

    started = time.monotonic()
    error = None
    stats = {}
    try:
        sender = module.Sender(tee, peer, window=args.window, timeout=args.timeout)
        stats = sender.send(data) or {}
    except NotImplementedError:
        raise SystemExit("%s.Sender.send() is not written yet. Finish it and pass "
                         "selfcheck.py before running this." % module.__name__)
    except Exception as exc:                       # noqa: BLE001 - the trace wants the reason
        error = "%s: %s" % (type(exc).__name__, exc)
    elapsed = time.monotonic() - started
    result = settle(holder, thread)
    tee.close()
    if tee.stopped and error:
        error = "cut short: the server had already given its verdict (below). %s" % error
    return {"bytes_offered": len(data), "seconds": round(elapsed, 2),
            "packets_sent": tee.sent, "packets_received": tee.received,
            "your_stats": stats, "your_error": error,
            "finack_payload": (tee.finack_payload or b"").decode(errors="replace") or None,
            "server": result}


def run_b(args, module, session, send, recv):
    """The server sends; your receiver has to reassemble it."""
    udp = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    udp.bind(("", 0))
    tee = Tee(udp, budget=session.get("seconds", 90))
    # Prime the path before the server sends: from behind a NAT or firewall an
    # unsolicited inbound datagram is dropped, which looks like a dead receiver.
    try:
        udp.sendto(b"", (args.host, session["udp_port"]))
    except OSError:
        pass
    send({"msg": "ready", "udp_port": udp.getsockname()[1]})
    holder, thread = watcher(tee, recv)

    started = time.monotonic()
    error, data = None, b""
    try:
        made = getattr(module, "Receiver", None) or stopwait.Receiver
        try:
            receiver = made(tee, window=args.window)
        except TypeError:
            receiver = made(tee)                    # a cumulative receiver takes no window
        data = receiver.receive()
    except NotImplementedError:
        raise SystemExit("%s.Receiver is not written yet. Finish it and pass "
                         "selfcheck.py before running this." % module.__name__)
    except Exception as exc:                       # noqa: BLE001
        error = "%s: %s" % (type(exc).__name__, exc)
    elapsed = time.monotonic() - started
    result = settle(holder, thread)
    tee.close()
    if tee.stopped and error:
        error = "cut short: the server had already given its verdict (below). %s" % error
    found = None
    for line in data.split(b"\n"):
        if line.startswith(b"FLAG: "):
            found = line[6:].decode(errors="replace").strip()
    return {"bytes_expected": session.get("bytes"), "bytes_received": len(data),
            "seconds": round(elapsed, 2), "packets_sent": tee.sent,
            "packets_received": tee.received, "your_error": error,
            "flag_in_the_data": found, "server": result}


def trace(challenge, student, info):
    """The block to paste into FINDINGS.md. This is the graded part of Stage 2."""
    lines = ["", "=" * 68,
             "TRACE  challenge %s  student %s" % (challenge, student),
             "=" * 68]
    server = info.pop("server", {}) or {}
    for key, value in info.items():
        lines.append("  %-22s %s" % (key, value))
    lines.append("  server said")
    if server.get("ok"):
        lines.append("    %-20s %s" % ("ok", True))
        lines.append("    %-20s %s" % ("flag", server.get("flag")))
    else:
        lines.append("    %-20s %s" % ("ok", False))
        lines.append("    %-20s %s" % ("reason", server.get("reason") or server.get("detail")))
    for key, value in (server.get("detail") or {}).items():
        lines.append("    %-20s %s" % (key, value))
    lines.append("=" * 68)
    return "\n".join(lines)


def main(argv=None):
    parser = argparse.ArgumentParser(description=(__doc__ or "challenge client").splitlines()[0])
    parser.add_argument("challenge", choices=CHALLENGES)
    parser.add_argument("--student", required=True, help="your UCCS username, lowercase")
    parser.add_argument("--server", default="66.42.112.186:8422",
                        help="host:port of the control channel; the default is the class server")
    parser.add_argument("--window", type=int, default=8)
    parser.add_argument("--timeout", type=float, default=0.15)
    args = parser.parse_args(argv)

    host, _, port = args.server.partition(":")
    args.host, args.port = host, int(port or 8422)
    side, protocol = args.challenge.split("-")
    module = importlib.import_module(protocol)

    runner = run_a if side == "A" else run_b
    try:
        _, info = control(args.host, args.port,
                          {"hello": {"msg": "hello", "challenge": args.challenge,
                                     "student": args.student.strip().lower()}},
                          lambda s, send, recv: runner(args, module, s, send, recv))
    except OSError as exc:
        raise SystemExit("could not reach %s:%d (%s). Keep this error and tell the "
                         "instructor; this package ships no local peer."
                         % (args.host, args.port, exc))

    print(trace(args.challenge, args.student, dict(info)))
    flag = info.get("flag_in_the_data") or (info.get("server") or {}).get("flag")
    if flag and flag != "(in the bytes this server sent)":
        print("\nFLAG  %s\n" % flag)
        return 0
    print("\nNo flag this time. The trace above is still worth marks: paste it into\n"
          "FINDINGS.md and say what you think went wrong.\n")
    return 1


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