"""Capture one TCP connection between two Mininet hosts and print every segment.

Given; do not change it. Run inside the course container, from this folder:

    python3 capture.py             # h1 sends 3000 bytes to h2 over TCP, h2 answers
    python3 capture.py slow        # h2 reads slowly: watch the window close
    python3 capture.py udp         # the same 3000 bytes over UDP, for comparison
    python3 capture.py show FILE   # print a saved capture again
    python3 capture.py cli         # the network with a mininet> prompt

Add --relative to any of them to print sequence and acknowledgment numbers
counted from the start of the connection, the way Wireshark shows them.

The network: h1 (10.0.0.1) and h2 (10.0.0.2) on one switch, s1, which delays
every packet 10 ms, so a round trip takes 20 ms.
tcpdump runs on h1 and saves every packet to captures/<mode>.pcap, which you
can also open in Wireshark on your own computer.

Two TCP options are turned off on both hosts so the numbers in the capture
are the ones the lecture uses: timestamps (which would take 12 bytes of every
segment) and window scaling (which would multiply the window field). The
maximum segment size is set to 1000 bytes so the arithmetic is easy.

References: tcpdump https://man7.org/linux/man-pages/man8/tcpdump.8.html,
tshark https://www.wireshark.org/docs/man-pages/tshark.html,
Python sockets https://docs.python.org/3/library/socket.html
"""

import os
import socket
import subprocess
import sys
import time

PORT = 5000
MSS = 1000
SEND_BYTES = 3000          # three full segments
SLOW_BYTES = 10000         # more than h2's small receive buffer holds
SLOW_BUFFER = 4096         # h2's receive buffer in slow mode (Linux doubles it)
SLOW_WAIT = 2.0            # seconds h2 waits before it starts reading
FLAGS = [(0x02, "SYN"), (0x10, "ACK"), (0x08, "PSH"), (0x01, "FIN"), (0x04, "RST")]
HERE = os.path.dirname(os.path.abspath(__file__))

def server(mode):
    if mode == "udp":
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        s.bind(("0.0.0.0", PORT))
        got = 0
        while got < SEND_BYTES:
            data, peer = s.recvfrom(65535)
            got += len(data)
        s.sendto(f"got {got} bytes\n".encode(), peer)
        return
    listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    listener.setsockopt(socket.IPPROTO_TCP, socket.TCP_MAXSEG, MSS)
    if mode == "slow":
        # set before listen() so the window h2 advertises in its SYN-ACK is already small
        listener.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, SLOW_BUFFER)
    listener.bind(("0.0.0.0", PORT))
    listener.listen(1)
    conn, _ = listener.accept()
    if mode == "slow":
        time.sleep(SLOW_WAIT)
    got = 0
    while True:
        data = conn.recv(65535)
        if not data:               # h1 closed its sending direction (its FIN)
            break
        got += len(data)
    conn.sendall(f"got {got} bytes\n".encode())
    conn.close()


def client(mode):
    total = SLOW_BYTES if mode == "slow" else SEND_BYTES
    if mode == "udp":
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        for _ in range(total // MSS):
            s.sendto(b"x" * MSS, ("10.0.0.2", PORT))
        s.settimeout(3)
        print("h2 answered:", s.recv(100).decode().strip())
        return
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.setsockopt(socket.IPPROTO_TCP, socket.TCP_MAXSEG, MSS)
    s.connect(("10.0.0.2", PORT))
    s.sendall(b"x" * total)
    s.shutdown(socket.SHUT_WR)     # "I am done sending", but h1 still reads the answer
    answer = b""
    while True:
        data = s.recv(100)
        if not data:
            break
        answer += data
    s.close()
    print("h2 answered:", answer.decode().strip())


def flag_names(value):
    bits = int(value, 16)
    return ",".join(name for bit, name in FLAGS if bits & bit) or "-"


def show(path, relative):
    """Print one line per packet: who sent it, flags, numbers, payload, window."""
    fields = ["frame.number", "frame.time_relative", "ip.src", "udp.length",
              "tcp.flags", "tcp.seq", "tcp.ack", "tcp.len", "tcp.window_size_value",
              "_ws.expert.message"]
    cmd = ["tshark", "-r", path, "-T", "fields", "-E", "separator=|", "-E", "occurrence=f",
           "-o", f"tcp.relative_sequence_numbers:{'TRUE' if relative else 'FALSE'}"]
    for f in fields:
        cmd += ["-e", f]
    out = subprocess.run(cmd, capture_output=True, text=True).stdout
    names = {"10.0.0.1": "h1 -> h2", "10.0.0.2": "h2 -> h1"}
    print(f"{'#':>3} {'time (ms)':>9}  {'direction':<9} {'flags':<12} {'seq':>10} {'ack':>10} "
          f"{'len':>5} {'win':>6}  note")
    for line in out.splitlines():
        num, t, src, udplen, flags, seq, ack, length, win, note = (line.split("|") + [""] * 10)[:10]
        ms = f"{float(t) * 1000:9.1f}"
        who = names.get(src, src)
        if udplen:
            print(f"{num:>3} {ms}  {who:<9} {'UDP':<12} {'':>10} {'':>10} {int(udplen) - 8:>5} {'':>6}")
            continue
        if not flags:
            continue
        ack_shown = ack if int(flags, 16) & 0x10 else ""
        print(f"{num:>3} {ms}  {who:<9} {flag_names(flags):<12} {seq:>10} {ack_shown:>10} "
              f"{length:>5} {win:>6}  {note.strip()}")

def build():
    from mininet.net import Mininet
    from mininet.nodelib import LinuxBridge
    from mininet.topo import Topo

    class TwoHosts(Topo):
        def build(self):
            h1 = self.addHost("h1", ip="10.0.0.1/24")
            h2 = self.addHost("h2", ip="10.0.0.2/24")
            s1 = self.addSwitch("s1")
            self.addLink(h1, s1)
            self.addLink(h2, s1)

    net = Mininet(topo=TwoHosts(), switch=LinuxBridge, controller=None, autoSetMacs=True)
    net.start()
    s1 = net.get("s1")
    for port in ("s1-eth1", "s1-eth2"):
        s1.cmd(f"tc qdisc replace dev {port} root netem delay 10ms")
    for h in net.hosts:
        intf = f"{h.name}-eth0"
        h.cmd(f"ethtool -K {intf} tso off gso off gro off lro off 2>/dev/null")
        h.cmd("sysctl -qw net.ipv4.tcp_timestamps=0 net.ipv4.tcp_window_scaling=0 net.ipv4.tcp_sack=0")
    h1, h2 = net.get("h1", "h2")
    for _ in range(5):
        if " 0 received" not in h1.cmd("ping -c 1 -W 1 10.0.0.2"):
            break
    return net, h1, h2


def run(mode, relative):
    net, h1, h2 = build()
    try:
        os.makedirs(os.path.join(HERE, "captures"), exist_ok=True)
        path = os.path.join(HERE, "captures", f"{mode}.pcap")
        dump = h1.popen(["tcpdump", "-i", "h1-eth0", "-U", "-n", "-w", path, "tcp or udp"],
                        stderr=subprocess.DEVNULL)
        time.sleep(1)                                   # let tcpdump start listening
        srv = h2.popen([sys.executable, os.path.abspath(__file__), "_server", mode])
        time.sleep(0.5)                                 # let the server start listening
        print(h1.cmd(f"{sys.executable} {os.path.abspath(__file__)} _client {mode}").strip())
        srv.wait(timeout=30)
        time.sleep(0.5)                                 # the last ACK reaches tcpdump
        dump.terminate()
        dump.wait()
        print(f"saved {os.path.relpath(path, os.getcwd())}\n")
        show(path, relative)
    finally:
        net.stop()


def main():
    args = [a for a in sys.argv[1:] if a != "--relative"]
    relative = "--relative" in sys.argv
    mode = args[0] if args else "tcp"
    if mode == "_server":
        server(args[1])
    elif mode == "_client":
        client(args[1])
    elif mode == "show":
        show(args[1], relative)
    elif mode == "cli":
        from mininet.cli import CLI
        net, _, _ = build()
        try:
            CLI(net)
        finally:
            net.stop()
    elif mode in ("tcp", "slow", "udp"):
        from mininet.log import setLogLevel
        setLogLevel("info")
        run(mode, relative)
    else:
        sys.exit(__doc__)


if __name__ == "__main__":
    main()