#!/bin/bash
# WireShield pairing CLI — the same workflow as the Cockpit "Quick Connect" page,
# for servers without a browser, Cockpit, or the Houston client app.
set -euo pipefail

API_HOST=${WIRESHIELD_API_HOST:-127.0.0.1}
API_PORT=${HNE_PORT:-8420}
API_BASE="https://${API_HOST}:${API_PORT}/api/v1"
ENV_FILE=${WIRESHIELD_API_ENV:-/etc/wireshield/api.env}
AUDIT_SOURCE="$(id -un)@$(hostname -s 2>/dev/null || echo localhost)/cli"

RESP_BODY=""
RESP_CODE=0
JSON_OUT=0

die() { echo "wireshield-pair: $*" >&2; exit 1; }

usage() {
    cat <<'EOF'
Usage: wireshield-pair <command> [options]

Commands:
  check                     Pairing readiness check (coordinator, public IP, NAT, UPnP)
  create [options]          Create a 6-character pairing code and wait for the other server
  join <CODE> [options]     Join using a code created on another server
  status <CODE> [--json]    Check a pending pairing code once and exit
  list                      List paired tunnels on this server
  remove <NAME> [-y]        Tear down and delete a paired tunnel on this server
  cancel <CODE> [--force]   Cancel a pending pairing code
  diag                      Report NAT traversal / hole punching status for each tunnel

Options for create:
  --name <name>             Tunnel name (a-z0-9-, max 15 chars). Default: auto-generated
  --ttl <minutes>           Code lifetime, 1-120 (default 15)
  --port <udp-port>         WireGuard listen port (default: next free port)
  --endpoint <host:port>    Advertise this endpoint instead of the detected public IP
  --no-wait                 Print the code and exit instead of waiting for the peer
  --json                    Machine-readable output

Options for join:
  --name <name>             Tunnel name. Omit to inherit the name from the other server
  --port <udp-port>         WireGuard listen port (default: next free port)
  --endpoint <host:port>    Advertise this endpoint instead of the detected public IP
  --json                    Machine-readable output

Examples:
  sudo wireshield-pair create --name testnet --ttl 30
  sudo wireshield-pair join ABC123
  sudo wireshield-pair remove testnet
  sudo wireshield-pair diag
EOF
}

load_key() {
    [[ $EUID -eq 0 ]] || die "must be run as root"
    [[ -r $ENV_FILE ]] || die "cannot read $ENV_FILE — is the wireshield package installed?"
    # shellcheck disable=SC1090
    source "$ENV_FILE"
    [[ -n ${HNE_API_KEY:-} ]] || die "HNE_API_KEY missing from $ENV_FILE"
}

# request METHOD PATH [JSON_BODY] -> sets RESP_BODY / RESP_CODE
request() {
    local method=$1 path=$2 body=${3:-} raw
    local args=(curl -sk --max-time 60 -X "$method"
        -H 'Content-Type: application/json'
        -H "X-API-Key: $HNE_API_KEY"
        -H "X-Audit-Source: $AUDIT_SOURCE"
        -w $'\n%{http_code}')
    [[ -n $body ]] && args+=(-d "$body")
    args+=("${API_BASE}${path}")

    raw=$("${args[@]}") || die "could not reach the WireShield API at $API_BASE (is wireshield.service running?)"
    RESP_CODE=${raw##*$'\n'}
    RESP_BODY=${raw%$'\n'*}
}

read -r -d '' PY_FIELD <<'PY' || true
import json, sys
try:
    data = json.loads(sys.argv[1])
except Exception:
    raise SystemExit(0)
for part in sys.argv[2].split("."):
    data = data.get(part) if isinstance(data, dict) else None
    if data is None:
        raise SystemExit(0)
print(data)
PY

read -r -d '' PY_BODY <<'PY' || true
import json, sys
out = {}
for pair in sys.argv[1:]:
    key, _, value = pair.partition("=")
    if value == "":
        continue
    out[key] = int(value) if value.isdigit() else value
print(json.dumps(out))
PY

read -r -d '' PY_CHECK <<'PY' || true
import json, sys
d = json.loads(sys.argv[1])
print("Status:        {}".format(d.get("status")))
print("Message:       {}".format(d.get("message")))
print("Coordinator:   {} (HTTP {})".format(
    "reachable" if d.get("vps_reachable") else "unreachable", d.get("vps_http_code")))
print("Enrolled:      {}".format(d.get("enrolled")))
print("Public IP:     {}".format(d.get("public_ip") or "unknown"))
print("Local IP:      {}".format(d.get("local_ip") or "unknown"))
print("Next port:     {}".format(d.get("listen_port")))
print("NAT endpoint:  {}".format(d.get("nat_endpoint") or "not discovered"))
print("UPnP mapping:  {}".format("available" if d.get("upnp_available") else "unavailable"))
PY

read -r -d '' PY_JOINED <<'PY' || true
import json, sys
d = json.loads(sys.argv[1])
print("Tunnel established")
print("  Interface:     {}".format(d.get("interface")))
print("  Address:       {}".format(d.get("address")))
print("  Listen port:   UDP {}".format(d.get("listen_port")))
print("  Peer address:  {}".format(d.get("peer_address")))
print("  Peer endpoint: {}".format(d.get("peer_endpoint")))
PY

read -r -d '' PY_LIST <<'PY' || true
import json, sys
conns = json.loads(sys.argv[1]).get("connections", [])
if not conns:
    print("No paired tunnels.")
    raise SystemExit(0)
print("{:16} {:>6}  {:18} {:24} {}".format("INTERFACE", "PORT", "SUBNET", "PEER ENDPOINT", "STATE"))
for c in conns:
    print("{:16} {:>6}  {:18} {:24} {}".format(
        c.get("interface", ""), c.get("listen_port", ""), c.get("subnet", ""),
        c.get("peer_endpoint") or "-", c.get("peer_status") or "-"))
PY

read -r -d '' PY_RESOLVE <<'PY' || true
import json, sys
conns = json.loads(sys.argv[1]).get("connections", [])
target = sys.argv[2].strip().lower()
for c in conns:
    if target in (str(c.get("interface") or "").lower(), str(c.get("network_id") or "").lower()):
        print(c.get("network_id"))
        break
PY

read -r -d '' PY_DIAG <<'PY' || true
import json, subprocess, sys, time

conns = json.loads(sys.argv[1]).get("connections", [])
if not conns:
    print("No paired tunnels to diagnose.")
    raise SystemExit(0)

for c in conns:
    iface = c.get("interface")
    negotiated = c.get("peer_endpoint") or ""
    print("[{}]".format(iface))
    print("  Endpoint negotiated at pairing: {}".format(negotiated or "unknown"))
    try:
        dump = subprocess.run(["wg", "show", iface, "dump"],
                              capture_output=True, text=True, check=True).stdout
    except Exception as exc:
        print("  Could not read live state: {}\n".format(exc))
        continue

    for row in [line.split("\t") for line in dump.strip().splitlines()][1:]:
        live, handshake = row[2], int(row[4] or 0)
        rx, tx = int(row[5] or 0), int(row[6] or 0)
        age = int(time.time()) - handshake if handshake else None
        print("  Live endpoint in use:           {}".format(live))
        print("  Last handshake:                 {}".format(
            "{}s ago".format(age) if age is not None else "never"))
        print("  Transfer:                       rx {} B / tx {} B".format(rx, tx))
        if handshake == 0:
            verdict = "NO HANDSHAKE - traversal has not succeeded (check UDP reachability / port forwarding)"
        elif age > 180:
            verdict = "STALE - handshake is old, the path may have dropped"
        elif live != negotiated:
            verdict = "DIRECT via a re-learned endpoint - hole punching worked and the path roamed"
        else:
            verdict = "DIRECT on the originally negotiated endpoint - traversal is working"
        print("  Verdict:                        {}\n".format(verdict))
PY

read -r -d '' PY_STATUS <<'PY' || true
import json, sys
session = json.loads(sys.argv[1])
conns = json.loads(sys.argv[2]).get("connections", []) if len(sys.argv) > 2 else []
match = next((c for c in conns if c.get("network_id") == session.get("network_id")), {})
session["interface"] = match.get("interface", "")
session["listen_port"] = match.get("listen_port", "")
session["subnet"] = match.get("subnet", "")
if sys.argv[3] == "json":
    print(json.dumps(session))
else:
    print("Code:          {}".format(session.get("code")))
    print("Role:          {}".format(session.get("role")))
    print("Status:        {}".format(session.get("status")))
    print("Interface:     {}".format(session.get("interface") or "-"))
    print("Peer endpoint: {}".format(session.get("peer_endpoint") or "-"))
    if session.get("error"):
        print("Error:         {}".format(session["error"]))
PY

field() { python3 -c "$PY_FIELD" "$1" "$2"; }

check_response() {
    [[ $RESP_CODE =~ ^2 ]] && return 0
    local detail
    detail=$(field "$RESP_BODY" detail)
    die "${detail:-API request failed (HTTP $RESP_CODE)}"
}

cmd_check() {
    [[ ${1:-} == --json ]] && JSON_OUT=1
    request GET /pairing/preflight
    check_response
    if [[ $JSON_OUT -eq 1 ]]; then
        echo "$RESP_BODY"
    else
        python3 -c "$PY_CHECK" "$RESP_BODY"
    fi
}

cmd_create() {
    local name="" ttl=15 port="" endpoint="" wait=1
    while [[ $# -gt 0 ]]; do
        case $1 in
            --name) name=${2:?--name requires a value}; shift 2 ;;
            --ttl) ttl=${2:?--ttl requires a value}; shift 2 ;;
            --port) port=${2:?--port requires a value}; shift 2 ;;
            --endpoint) endpoint=${2:?--endpoint requires a value}; shift 2 ;;
            --no-wait) wait=0; shift ;;
            --json) JSON_OUT=1; shift ;;
            *) die "unknown option for create: $1" ;;
        esac
    done

    local body
    body=$(python3 -c "$PY_BODY" "name=$name" "ttl_minutes=$ttl" "listen_port=$port" "endpoint_override=$endpoint")
    request POST /pairing/initiate "$body"
    check_response

    local code
    code=$(field "$RESP_BODY" code)

    if [[ $JSON_OUT -eq 1 ]]; then
        echo "$RESP_BODY"
    else
        cat <<EOF

  Pairing code: $code

  Tunnel name:  $(field "$RESP_BODY" interface)
  Listen port:  UDP $(field "$RESP_BODY" listen_port)
  Expires at:   $(field "$RESP_BODY" expires_at)

  On the other server, run:
    sudo wireshield-pair join $code

EOF
    fi

    [[ $wait -eq 1 ]] || return 0
    wait_for_peer "$code"
}

wait_for_peer() {
    local code=$1 status detail
    [[ $JSON_OUT -eq 1 ]] || echo "Waiting for the other server to enter the code (Ctrl-C to stop waiting)..."
    while true; do
        sleep 3
        request GET "/pairing/session/${code}"
        check_response
        status=$(field "$RESP_BODY" status)
        case $status in
            complete)
                if [[ $JSON_OUT -eq 1 ]]; then
                    echo "$RESP_BODY"
                else
                    echo "Tunnel established (peer endpoint: $(field "$RESP_BODY" peer_endpoint))"
                fi
                return 0
                ;;
            expired|failed)
                detail=$(field "$RESP_BODY" error)
                die "pairing ${status}: ${detail:-no further detail}"
                ;;
        esac
    done
}

cmd_join() {
    local code=${1:-}
    [[ -n $code ]] || die "join requires a pairing code"
    shift
    local name="" port="" endpoint=""
    while [[ $# -gt 0 ]]; do
        case $1 in
            --name) name=${2:?--name requires a value}; shift 2 ;;
            --port) port=${2:?--port requires a value}; shift 2 ;;
            --endpoint) endpoint=${2:?--endpoint requires a value}; shift 2 ;;
            --json) JSON_OUT=1; shift ;;
            *) die "unknown option for join: $1" ;;
        esac
    done

    local body
    body=$(python3 -c "$PY_BODY" "code=${code^^}" "name=$name" "listen_port=$port" "endpoint_override=$endpoint")
    request POST /pairing/join "$body"
    check_response

    if [[ $JSON_OUT -eq 1 ]]; then
        echo "$RESP_BODY"
    else
        python3 -c "$PY_JOINED" "$RESP_BODY"
    fi
}

cmd_list() {
    [[ ${1:-} == --json ]] && JSON_OUT=1
    request GET /pairing/connections
    check_response
    if [[ $JSON_OUT -eq 1 ]]; then
        echo "$RESP_BODY"
    else
        python3 -c "$PY_LIST" "$RESP_BODY"
    fi
}

cmd_status() {
    local code=${1:-}
    [[ -n $code ]] || die "status requires a pairing code"
    [[ ${2:-} == --json ]] && JSON_OUT=1

    request GET "/pairing/session/${code^^}"
    check_response
    local session=$RESP_BODY

    request GET /pairing/connections
    check_response

    python3 -c "$PY_STATUS" "$session" "$RESP_BODY" "$([[ $JSON_OUT -eq 1 ]] && echo json || echo text)"
}

cmd_remove() {
    local target=${1:-} assume_yes=0
    [[ -n $target ]] || die "remove requires a tunnel name (see: wireshield-pair list)"
    shift
    while [[ $# -gt 0 ]]; do
        case $1 in
            -y|--yes) assume_yes=1; shift ;;
            *) die "unknown option for remove: $1" ;;
        esac
    done

    request GET /pairing/connections
    check_response

    local network_id
    network_id=$(python3 -c "$PY_RESOLVE" "$RESP_BODY" "$target")
    [[ -n $network_id ]] || die "no paired tunnel named '$target' (see: wireshield-pair list)"

    if [[ $assume_yes -eq 0 ]]; then
        local reply
        read -r -p "Remove tunnel '$target' from this server? [y/N] " reply
        [[ ${reply,,} == y* ]] || { echo "Aborted."; return 0; }
    fi

    request DELETE "/networks/${network_id}"
    check_response
    echo "Removed tunnel $target"
    echo "The other server still has its half of this tunnel; run the same command there."
}

cmd_cancel() {
    local code=${1:-}
    [[ -n $code ]] || die "cancel requires a pairing code"
    local path="/pairing/session/${code^^}"
    [[ ${2:-} == --force ]] && path="${path}?force=true"
    request DELETE "$path"
    check_response
    echo "Cancelled pairing code ${code^^}"
}

cmd_diag() {
    command -v wg >/dev/null || die "wg command not found"

    request GET /pairing/preflight
    check_response
    echo "Public IP:            $(field "$RESP_BODY" public_ip)"
    echo "NAT endpoint (STUN):  $(field "$RESP_BODY" nat_endpoint)"
    echo "UPnP port mapping:    $(field "$RESP_BODY" upnp_available)"
    echo

    request GET /pairing/connections
    check_response
    python3 -c "$PY_DIAG" "$RESP_BODY"
}

main() {
    local cmd=${1:-}
    [[ -n $cmd ]] || { usage; exit 1; }
    shift || true

    case $cmd in
        -h|--help|help) usage; exit 0 ;;
    esac

    load_key

    case $cmd in
        check) cmd_check "$@" ;;
        create) cmd_create "$@" ;;
        join) cmd_join "$@" ;;
        status) cmd_status "$@" ;;
        list) cmd_list "$@" ;;
        remove) cmd_remove "$@" ;;
        cancel) cmd_cancel "$@" ;;
        diag) cmd_diag ;;
        *) usage; die "unknown command: $cmd" ;;
    esac
}

main "$@"
