#!/usr/bin/env python3
"""Opens the live captures of the holo lab in Wireshark.

A page cannot start a program, so this starts Wireshark for it. The page
streams a live capture to the WebSocket this listens on, at 127.0.0.1:8740,
and every capture streamed opens a Wireshark of its own, which reads it from
its standard input as it would a live capture, as `wireshark -k -i -` does.
Closing that Wireshark stops the capture on the page, and stopping the capture
on the page ends it in Wireshark, which keeps its window open on what it
captured. A Wireshark outlives the page and this both.

Wireshark captures through dumpcap, even from a pipe, and a Linux distribution
may leave dumpcap to the members of the wireshark group to run.

Only the pages of the lab may connect: those served from this machine, as
serve.py serves them, those of the site the lab is deployed to, and those of
the origins given with --origin.

    ./wireshark-bridge.py [--wireshark <path>] [--origin <origin>]...

It takes nothing beyond the standard library, so it runs as the page serves
it.
"""

import argparse
import asyncio
import base64
import hashlib
import os
import queue
import shutil
import struct
import subprocess
import sys
import threading
import urllib.parse

# Where the page looks for this.
HOST, PORT = "127.0.0.1", 8740
# The site the lab is deployed to, whose previews are on subdomains of it.
SITE = "holo-lab.pages.dev"
# What the key of a handshake is hashed with to accept it, per RFC 6455.
GUID = b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
# Opcodes of the frames, and the codes a connection is closed with.
OP_CONTINUATION, OP_BINARY, OP_CLOSE, OP_PING, OP_PONG = 0, 2, 8, 9, 10
CLOSE_NORMAL, CLOSE_ERROR = 1000, 1011
# Most a frame from the page carries: more than a capture holds before it is
# taken from, which is what the page sends at once at most.
MAX_FRAME = 64 * 1024 * 1024
# Most a Wireshark may leave unread before it is taken to have stopped
# capturing, as it does when its capture is stopped with its window open.
MAX_BACKLOG = 16 * 1024 * 1024
# How often whether a Wireshark is still there is looked at.
WATCH_SECONDS = 0.5


class Wireshark:
    """A Wireshark reading a live capture from its standard input.

    What it is fed is written by a thread of its own, since a Wireshark that
    stops reading blocks whoever writes to it.
    """

    def __init__(self, path, title):
        # A session of its own keeps it from going with this on a ^C.
        self.process = subprocess.Popen(
            [path, "-o", f"gui.window_title:{title}", "-k", "-i", "-"],
            stdin=subprocess.PIPE,
            stdout=subprocess.DEVNULL,
            start_new_session=True,
        )
        self.queue = queue.SimpleQueue()
        self.backlog = 0
        self.broken = False
        threading.Thread(target=self.write, daemon=True).start()

    def write(self):
        stdin = self.process.stdin
        while (data := self.queue.get()) is not None:
            try:
                stdin.write(data)
                stdin.flush()
            except OSError:
                self.broken = True
                break
            self.backlog -= len(data)
        try:
            stdin.close()
        except OSError:
            pass

    def feed(self, data):
        self.backlog += len(data)
        self.queue.put(data)

    def end(self):
        """Ends what it is fed, which ends its capture and leaves it open."""
        self.queue.put(None)

    def gone(self):
        """Says why it reads no more, or nothing while it does."""
        if self.process.poll() is not None:
            return "Wireshark closed"
        if self.broken or self.backlog > MAX_BACKLOG:
            return "Wireshark stopped reading the capture"
        return None


def find_wireshark():
    candidates = [
        shutil.which("wireshark"),
        "/Applications/Wireshark.app/Contents/MacOS/Wireshark",
        os.path.join(
            os.environ.get("ProgramFiles", r"C:\Program Files"),
            "Wireshark",
            "Wireshark.exe",
        ),
    ]
    return next((path for path in candidates if path and os.path.isfile(path)), None)


def unrunnable_dumpcap(wireshark):
    """Returns the dumpcap Wireshark captures through when it is there and is
    not ours to run, which leaves Wireshark unable to capture."""
    dumpcap = os.path.join(os.path.dirname(os.path.realpath(wireshark)), "dumpcap")
    runnable = os.access(dumpcap, os.X_OK)
    if os.name == "posix" and os.path.isfile(dumpcap) and not runnable:
        return dumpcap
    return None


def allowed(origin, origins):
    if origin in origins:
        return True
    url = urllib.parse.urlsplit(origin)
    host = url.hostname or ""
    if url.scheme == "http" and host in ("localhost", "127.0.0.1", "::1"):
        return True
    return url.scheme == "https" and (host == SITE or host.endswith("." + SITE))


def frame(opcode, payload=b""):
    """A frame as this sends it, whole and unmasked."""
    size = len(payload)
    if size < 126:
        head = struct.pack("!BB", 0x80 | opcode, size)
    elif size < 1 << 16:
        head = struct.pack("!BBH", 0x80 | opcode, 126, size)
    else:
        head = struct.pack("!BBQ", 0x80 | opcode, 127, size)
    return head + payload


def close_frame(code, reason):
    # A close frame carries 125 bytes at most, two of them the code, and a
    # reason that is not whole UTF-8 fails the connection instead.
    reason = reason.encode()[:123].decode(errors="ignore").encode()
    return frame(OP_CLOSE, struct.pack("!H", code) + reason)


async def read_frame(reader):
    """Reads a frame from the page, which masks every frame it sends."""
    first, second = await reader.readexactly(2)
    size = second & 0x7F
    if size == 126:
        (size,) = struct.unpack("!H", await reader.readexactly(2))
    elif size == 127:
        (size,) = struct.unpack("!Q", await reader.readexactly(8))
    if not second & 0x80 or size > MAX_FRAME:
        raise ConnectionError("not a frame a page sends")
    mask = await reader.readexactly(4)
    data = await reader.readexactly(size)
    key = int.from_bytes((mask * (size // 4 + 1))[:size], "big")
    data = (int.from_bytes(data, "big") ^ key).to_bytes(size, "big")
    return first & 0x0F, data


async def handshake(reader, writer, origins):
    """Accepts the WebSocket of a page of the lab, and returns the title of
    the capture it streams, or refuses anything else and returns None."""
    try:
        head = await asyncio.wait_for(reader.readuntil(b"\r\n\r\n"), 10)
    except (asyncio.IncompleteReadError, asyncio.LimitOverrunError,
            asyncio.TimeoutError):
        return None
    request, *lines = head.decode("latin-1").split("\r\n")
    headers = {}
    for line in lines:
        name, colon, value = line.partition(":")
        if colon:
            headers[name.strip().lower()] = value.strip()

    method, target, *_ = request.split(" ") + ["", ""]
    key = headers.get("sec-websocket-key")
    origin = headers.get("origin", "")
    if method != "GET" or not key or headers.get("upgrade", "").lower() != "websocket":
        status = "400 Bad Request"
    elif not allowed(origin, origins):
        print(f"refused {origin or 'a client with no origin'}: "
              f"--origin {origin or '<origin>'} lets it in", flush=True)
        status = "403 Forbidden"
    else:
        status = None
    if status:
        writer.write(f"HTTP/1.1 {status}\r\nContent-Length: 0\r\n\r\n".encode())
        await writer.drain()
        return None

    accept = base64.b64encode(hashlib.sha1(key.encode() + GUID).digest()).decode()
    writer.write(
        "HTTP/1.1 101 Switching Protocols\r\n"
        "Upgrade: websocket\r\n"
        "Connection: Upgrade\r\n"
        f"Sec-WebSocket-Accept: {accept}\r\n\r\n".encode()
    )
    await writer.drain()

    # Wireshark reads a % in a title as the start of a field of its own.
    query = urllib.parse.parse_qs(urllib.parse.urlsplit(target).query)
    title = query.get("title", ["holo lab"])[0]
    return "".join(c for c in title if c.isprintable() and c != "%")[:200]


async def relay(reader, writer, feed):
    """Feeds what the page sends to the given function until the page closes
    the connection, and returns the code it closed it with."""
    while True:
        opcode, data = await read_frame(reader)
        if opcode in (OP_BINARY, OP_CONTINUATION):
            feed(data)
        elif opcode == OP_PING:
            writer.write(frame(OP_PONG, data))
        elif opcode == OP_CLOSE:
            return data[:2]


async def watch(wireshark):
    while not wireshark.gone():
        await asyncio.sleep(WATCH_SECONDS)
    return wireshark.gone()


async def close(writer, relaying, code, reason):
    """Closes the connection, saying why, once the page has closed it too or
    failed to in time: closing it on what the page still sends would reset
    it, which can lose the page why."""
    writer.write(close_frame(code, reason))
    await writer.drain()
    try:
        await asyncio.wait_for(relaying, 5)
    except (asyncio.TimeoutError, ConnectionError, asyncio.IncompleteReadError):
        pass


async def serve(reader, writer, args):
    title = None
    try:
        title = await handshake(reader, writer, args.origin)
        if title is None:
            return
        problem = None
        if unrunnable_dumpcap(args.wireshark):
            problem = "Wireshark cannot run dumpcap: join the wireshark group"
        else:
            try:
                wireshark = Wireshark(args.wireshark, title)
            except OSError as error:
                problem = f"cannot start Wireshark: {error.strerror or error}"
        if problem:
            print(f"{title}: {problem}", flush=True)
            discarding = asyncio.ensure_future(relay(reader, writer, lambda data: None))
            await close(writer, discarding, CLOSE_ERROR, problem)
            return
        print(f"{title}: Wireshark started", flush=True)

        # Whichever ends first ends the other: the page stopping the capture
        # ends Wireshark's, and Wireshark going stops the page's.
        relaying = asyncio.ensure_future(relay(reader, writer, wireshark.feed))
        watching = asyncio.ensure_future(watch(wireshark))
        await asyncio.wait({relaying, watching}, return_when=asyncio.FIRST_COMPLETED)
        wireshark.end()
        watching.cancel()
        if relaying.done():
            writer.write(frame(OP_CLOSE, relaying.result()))
            await writer.drain()
            print(f"{title}: the capture ended", flush=True)
        else:
            reason = watching.result()
            print(f"{title}: {reason}", flush=True)
            await close(writer, relaying, CLOSE_NORMAL, reason)
    except (ConnectionError, asyncio.IncompleteReadError):
        if title:
            print(f"{title}: the page went", flush=True)
    finally:
        writer.close()


async def main():
    parser = argparse.ArgumentParser(
        description="Opens the live captures of the holo lab in Wireshark.")
    parser.add_argument("--wireshark", default=find_wireshark(),
                        help="the Wireshark to start, the one on PATH by default")
    parser.add_argument("--origin", action="append", default=[],
                        help="an origin whose pages may connect besides the lab's own, "
                             "such as https://lab.example.com")
    args = parser.parse_args()
    if not args.wireshark:
        sys.exit("no Wireshark found: give one with --wireshark")
    dumpcap = unrunnable_dumpcap(args.wireshark)
    if dumpcap:
        print(f"Wireshark captures through {dumpcap}, which is not yours to run.\n"
              "To capture live, join the wireshark group:\n"
              "    sudo usermod -aG wireshark $USER\n"
              "and log in again, or start this under `sg wireshark`.", flush=True)

    try:
        server = await asyncio.start_server(
            lambda reader, writer: serve(reader, writer, args), HOST, PORT)
    except OSError as error:
        sys.exit(f"cannot listen on {HOST}:{PORT}: {error.strerror}; "
                 "is the bridge running already?")
    print(f"opening the lab's live captures in {args.wireshark}, "
          f"listening on ws://{HOST}:{PORT}", flush=True)
    async with server:
        await server.serve_forever()


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        pass