#!/usr/bin/env python3
"""Send one CoAP request over UDP and show both datagrams byte by byte.

    ./coap-inspect.py get "/hello?name=Imran"
    ./coap-inspect.py post /hello Imran
    ./coap-inspect.py get /health

No CoAP library: the packet is built and decoded right here, so every byte on
screen is accounted for. COAP_HOST / COAP_PORT retarget it (127.0.0.1:5683).
"""
import os
import random
import socket
import sys
import time

TYPES = {0: "CON confirmable", 1: "NON non-confirmable", 2: "ACK acknowledgement", 3: "RST reset"}
METHODS = {"get": 1, "post": 2, "put": 3, "delete": 4}
CODES = {
    "0.00": "Empty", "0.01": "GET", "0.02": "POST", "0.03": "PUT", "0.04": "DELETE",
    "2.01": "Created", "2.02": "Deleted", "2.03": "Valid", "2.04": "Changed", "2.05": "Content",
    "4.00": "Bad Request", "4.04": "Not Found", "4.05": "Method Not Allowed",
    "5.00": "Internal Server Error",
}
OPTIONS = {3: "Uri-Host", 7: "Uri-Port", 11: "Uri-Path", 12: "Content-Format", 15: "Uri-Query", 17: "Accept"}
FORMATS = {0: "text/plain", 50: "application/json"}

TTY = sys.stdout.isatty()


def bold(text):
    return f"\033[1m{text}\033[0m" if TTY else text


def encode_option(delta, value):
    """One option: a byte of 4-bit delta + 4-bit length, then any extended bytes, then the value."""
    def nibble(n):
        if n < 13:
            return n, b""
        if n < 269:
            return 13, bytes([n - 13])
        return 14, (n - 269).to_bytes(2, "big")

    d, d_ext = nibble(delta)
    l, l_ext = nibble(len(value))
    return bytes([(d << 4) | l]) + d_ext + l_ext + value


def build_request(method, path, query, payload):
    token = os.urandom(4)
    header = bytes([(1 << 6) | (0 << 4) | len(token), METHODS[method]])  # ver 1, CON, TKL, code
    header += random.randrange(1 << 16).to_bytes(2, "big") + token        # message id, token

    # Options must be in ascending number order; each stores the gap from the previous one.
    options = [(11, s.encode()) for s in path.split("/") if s]
    options += [(15, q.encode()) for q in query.split("&") if q]
    body, last = b"", 0
    for number, value in options:
        body += encode_option(number - last, value)
        last = number
    if payload:
        body += b"\xff" + payload
    return header + body


def decode(pkt):
    """Return (rows, meta). Each row is (start, end, field, meaning) over pkt[start:end]."""
    if len(pkt) < 4:
        raise ValueError(f"{len(pkt)} bytes is shorter than the 4-byte CoAP header")

    rows = []
    b0 = pkt[0]
    ver, typ, tkl = b0 >> 6, (b0 >> 4) & 0b11, b0 & 0b1111
    rows.append((0, 1, "ver|type|tkl",
                 f"{ver:02b} {typ:02b} {tkl:04b} -> version {ver}, {TYPES[typ]}, token {tkl} bytes"))

    cls, detail = pkt[1] >> 5, pkt[1] & 0b11111
    code = f"{cls}.{detail:02d}"
    rows.append((1, 2, "code", f"{cls:03b} {detail:05b} -> {code} {CODES.get(code, '')}"))

    mid = int.from_bytes(pkt[2:4], "big")
    rows.append((2, 4, "message id", f"{mid} (0x{mid:04x})"))

    token = pkt[4:4 + tkl]
    if tkl:
        rows.append((4, 4 + tkl, "token", f"0x{token.hex()}"))

    i, number = 4 + tkl, 0
    while i < len(pkt):
        if pkt[i] == 0xFF:
            body = pkt[i + 1:]
            rows.append((i, i + 1, "payload marker", "0xff -> the rest is the body"))
            rows.append((i + 1, len(pkt), "payload",
                         f"{len(body)} bytes: {body.decode('utf-8', 'replace')}"))
            break

        start = i
        delta, length = pkt[i] >> 4, pkt[i] & 0b1111
        i += 1
        extended = []
        for which, n in (("delta", delta), ("length", length)):
            if n == 13:
                n, i = pkt[i] + 13, i + 1
            elif n == 14:
                n, i = int.from_bytes(pkt[i:i + 2], "big") + 269, i + 2
            extended.append(n)
        delta, length = extended
        number += delta
        value = pkt[i:i + length]
        i += length

        if number in (12, 17):
            n = int.from_bytes(value, "big") if value else 0
            shown = f"{n} = {FORMATS.get(n, 'unknown format')}"
        else:
            shown = f'"{value.decode("utf-8", "replace")}"'
        rows.append((start, i, OPTIONS.get(number, f"option {number}"),
                     f"delta {delta} -> option {number}, {length} bytes: {shown}"))

    return rows, {"type": typ, "code": code, "mid": mid, "token": token}


def show(title, pkt):
    print(bold(f"\n{title}  ({len(pkt)} bytes)"))

    print(bold("\n  hex + ascii"))
    for off in range(0, len(pkt), 16):
        chunk = pkt[off:off + 16]
        hexes = " ".join(f"{b:02x}" for b in chunk)
        text = "".join(chr(b) if 32 <= b < 127 else "." for b in chunk)
        print(f"  {off:04x}  {hexes:<47}  |{text}|")

    print(bold("\n  binary"))
    for off in range(0, len(pkt), 4):
        chunk = pkt[off:off + 4]
        bits = " ".join(f"{b:08b}" for b in chunk)
        text = "".join(chr(b) if 32 <= b < 127 else "." for b in chunk)
        print(f"  {off:04x}  {bits:<35}  |{text}|")

    rows, meta = decode(pkt)
    print(bold("\n  fields"))
    print(f"  {'byte':<7} {'hex':<24} {'field':<15} meaning")
    for start, end, field, meaning in rows:
        span = f"{start}" if end - start == 1 else f"{start}-{end - 1}"
        hexes = pkt[start:end].hex(" ")
        if len(hexes) > 23:
            hexes = hexes[:20] + "..."
        print(f"  {span:<7} {hexes:<24} {field:<15} {meaning}")
    return meta


def main():
    args = sys.argv[1:]
    if not args or args[0].lower() not in METHODS:
        print(__doc__.strip())
        return 2

    method = args[0].lower()
    target = args[1] if len(args) > 1 else "/hello"
    payload = args[2].encode() if len(args) > 2 else b""
    path, _, query = target.partition("?")
    host = os.environ.get("COAP_HOST", "127.0.0.1")
    port = int(os.environ.get("COAP_PORT", "5683"))

    request = build_request(method, path, query, payload)

    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)  # SOCK_DGRAM = UDP
    sock.settimeout(3)
    sock.connect((host, port))
    local = "%s:%d" % sock.getsockname()

    req = show(f">>> REQUEST   udp {local} -> {host}:{port}", request)
    sent_at = time.perf_counter()
    sock.send(request)
    try:
        response = sock.recv(1152)
    except socket.timeout:
        print(f"\nno reply within 3s - is the CoAP server running on {host}:{port}?")
        return 1
    except ConnectionRefusedError:
        print(f"\nport unreachable - nothing is listening on udp {host}:{port}")
        return 1
    elapsed = (time.perf_counter() - sent_at) * 1000

    res = show(f"<<< RESPONSE  udp {host}:{port} -> {local}", response)

    print(bold("\n  exchange"))
    print(f"  request   {TYPES[req['type']].split()[0]}  {req['code']}  mid 0x{req['mid']:04x}  "
          f"token 0x{req['token'].hex()}  {len(request)} bytes")
    print(f"  response  {TYPES[res['type']].split()[0]}  {res['code']}  mid 0x{res['mid']:04x}  "
          f"token 0x{res['token'].hex()}  {len(response)} bytes  after {elapsed:.2f} ms")
    same_mid = "yes" if req["mid"] == res["mid"] else "NO"
    same_token = "yes" if req["token"] == res["token"] else "NO"
    kind = " (piggybacked: the ACK carries the answer)" if res["type"] == 2 else ""
    print(f"  message id matches: {same_mid}   token matches: {same_token}{kind}")
    return 0


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