#!/usr/bin/python3
"""mite-tunnel: create/delete/inspect MITE kernel tunnels (module mite_tun.ko).

iproute2 has no "mite" link type plugin, so `ip link add ... type mite` cannot carry the
module's attributes (local, remote, mode, tid, allowed_src, ...). This helper speaks
rtnetlink directly (stdlib only) and uses `ip` for address/route/up.

  mite-tunnel add NAME --local A --remote B --mode A|B [--tid N] --allowed-src CIDR[,CIDR..]
                  [--deny-src CIDR,..|none] [--deny-dst CIDR,..|none]
                  [--unknown-proto pass|drop] [--mtu N]
  mite-tunnel del NAME
  mite-tunnel show [NAME]            config + counters of one tunnel (all if omitted)
  mite-tunnel counters NAME [-a]     non-zero counters (-a: all)
  mite-tunnel global | pairs         counters that cannot be attributed to a tunnel
  mite-tunnel up [FILE|DIR]          create tunnels from /etc/mite/tunnels.d/*.conf
  mite-tunnel down [FILE|DIR]        delete them

Config file (key=value, '#' comments, one tunnel per file; repeatable keys: route, address):
  name=mite0  local=198.18.1.2  remote=198.18.0.2  mode=B  tid=1
  allowed_src=10.202.0.0/24,10.202.1.0/24     # required
  deny_src= / deny_dst=                       # optional; absent = SPEC default list
  unknown_proto=pass   mtu=1490   address=10.255.0.1/32   route=10.202.0.0/24
  nft_fragdrop=yes    # also install nft prerouting -450 drop of proto-253 fragments
"""
import glob
import ipaddress
import os
import socket
import struct
import subprocess
import sys

DEBUGFS = "/sys/kernel/debug/mite"
CONFDIR = "/etc/mite/tunnels.d"

RTM_NEWLINK, RTM_DELLINK = 16, 17
NLM_F_REQUEST, NLM_F_ACK, NLM_F_EXCL, NLM_F_CREATE = 1, 4, 0x200, 0x400
IFLA_IFNAME, IFLA_MTU, IFLA_LINKINFO = 3, 4, 18
IFLA_INFO_KIND, IFLA_INFO_DATA = 1, 2
(M_LOCAL, M_REMOTE, M_MODE, M_TID, M_ALLOWED, M_DENY_SRC, M_DENY_DST, M_UNKNOWN) = range(1, 9)
NESTED = 0x8000


def attr(t, payload):
    n = 4 + len(payload)
    return struct.pack("=HH", n, t) + payload + b"\0" * (-n % 4)


def cidrs(s):
    out = b""
    for c in [x for x in s.split(",") if x]:
        n = ipaddress.ip_network(c.strip(), strict=False)
        if n.version != 4:
            raise SystemExit(f"not IPv4: {c}")
        out += n.network_address.packed + bytes([n.prefixlen, 0, 0, 0])
    return out


def rtnl(msg_type, flags, body):
    s = socket.socket(socket.AF_NETLINK, socket.SOCK_RAW, 0)
    try:
        s.bind((0, 0))
        s.send(struct.pack("=IHHII", 16 + len(body), msg_type, flags, 1, 0) + body)
        r = s.recv(65536)
    finally:
        s.close()
    ln, ty, fl, seq, pid = struct.unpack("=IHHII", r[:16])
    if ty == 2:
        err = struct.unpack("=i", r[16:20])[0]
        if err:
            ext = b""
            # extended ack message (NLMSGERR_ATTR_MSG) if present: scan attrs after echoed header
            p = r[20 + 16:] if len(r) > 36 else b""
            i = 0
            while i + 4 <= len(p):
                al, at = struct.unpack("=HH", p[i:i + 4])
                if al < 4:
                    break
                if at == 1:
                    ext = p[i + 4:i + al].rstrip(b"\0")
                i += (al + 3) & ~3
            raise SystemExit(f"netlink: {os.strerror(-err)}" + (f" ({ext.decode(errors='replace')})" if ext else ""))


def add(name, local, remote, mode, tid, allowed, deny_src=None, deny_dst=None, unknown="pass", mtu=None):
    d = attr(M_LOCAL, socket.inet_aton(local)) + attr(M_REMOTE, socket.inet_aton(remote))
    d += attr(M_MODE, bytes([0 if mode.upper() == "A" else 1]))
    if tid:
        d += attr(M_TID, bytes([int(tid)]))
    d += attr(M_ALLOWED, cidrs(allowed))
    if deny_src is not None:
        d += attr(M_DENY_SRC, b"" if deny_src == "none" else cidrs(deny_src))
    if deny_dst is not None:
        d += attr(M_DENY_DST, b"" if deny_dst == "none" else cidrs(deny_dst))
    d += attr(M_UNKNOWN, bytes([1 if unknown == "drop" else 0]))
    li = attr(IFLA_INFO_KIND, b"mite\0") + attr(IFLA_INFO_DATA | NESTED, d)
    body = struct.pack("=BxHiII", 0, 0, 0, 0, 0)
    body += attr(IFLA_IFNAME, name.encode() + b"\0")
    if mtu:
        body += attr(IFLA_MTU, struct.pack("=I", int(mtu)))
    body += attr(IFLA_LINKINFO | NESTED, li)
    rtnl(RTM_NEWLINK, NLM_F_REQUEST | NLM_F_ACK | NLM_F_CREATE | NLM_F_EXCL, body)
    # MITE is IPv4-only: keep IPv6 (RS/MLD/DAD) off the device so it never sees (and counts) non-IPv4 frames
    try:
        with open(f"/proc/sys/net/ipv6/conf/{name}/disable_ipv6", "w") as f:
            f.write("1")
    except OSError:
        pass


def delete(name):
    idx = socket.if_nametoindex(name)
    rtnl(RTM_DELLINK, NLM_F_REQUEST | NLM_F_ACK, struct.pack("=BxHiII", 0, 0, idx, 0, 0))


def sh(*a, check=True):
    return subprocess.run(a, check=check)


def read_conf(path):
    c = {"route": [], "address": []}
    for line in open(path):
        line = line.split("#", 1)[0].strip()
        if not line:
            continue
        k, _, v = line.partition("=")
        k, v = k.strip(), v.strip()
        if k in ("route", "address"):
            c[k].append(v)
        else:
            c[k] = v
    return c


def conf_files(arg):
    arg = arg or CONFDIR
    return sorted(glob.glob(os.path.join(arg, "*.conf"))) if os.path.isdir(arg) else [arg]


NFT = """table ip mite_frag {
  chain pre { type filter hook prerouting priority -450; policy accept;
    ip protocol 253 ip frag-off & 0x3fff != 0 counter drop
  }
}
"""


def up(arg):
    want_nft = False
    for f in conf_files(arg):
        c = read_conf(f)
        name = c["name"]
        if os.path.exists(f"/sys/class/net/{name}"):
            print(f"{name}: exists, skipping")
            continue
        add(name, c["local"], c["remote"], c["mode"], c.get("tid"), c["allowed_src"],
            c.get("deny_src"), c.get("deny_dst"), c.get("unknown_proto", "pass"), c.get("mtu", 1490))
        for a in c["address"]:
            sh("ip", "addr", "add", a, "dev", name)
        sh("ip", "link", "set", name, "up")
        for r in c["route"]:
            sh("ip", "route", "replace", r, "dev", name)
        want_nft |= c.get("nft_fragdrop", "no") == "yes"
        print(f"{name}: up")
    if want_nft and subprocess.run(["nft", "-f", "-"], input=NFT.encode()).returncode:
        print("warning: nft rule not installed (module drops proto-253 fragments itself)", file=sys.stderr)


def down(arg):
    for f in conf_files(arg):
        name = read_conf(f)["name"]
        if os.path.exists(f"/sys/class/net/{name}"):
            delete(name)
            print(f"{name}: deleted")
    sh("nft", "delete", "table", "ip", "mite_frag", check=False)


def counters(name, all_=False):
    for line in open(f"{DEBUGFS}/{name}/counters"):
        k, v = line.split()
        if all_ or int(v):
            print(f"{k:24s} {v}")


def main():
    a = sys.argv[1:]
    if not a:
        print(__doc__)
        return
    cmd = a[0]
    if cmd == "add":
        import argparse
        p = argparse.ArgumentParser(prog="mite-tunnel add")
        p.add_argument("name")
        p.add_argument("--local", required=True)
        p.add_argument("--remote", required=True)
        p.add_argument("--mode", required=True, choices=["A", "B", "a", "b"])
        p.add_argument("--tid")
        p.add_argument("--allowed-src", required=True)
        p.add_argument("--deny-src")
        p.add_argument("--deny-dst")
        p.add_argument("--unknown-proto", default="pass", choices=["pass", "drop"])
        p.add_argument("--mtu", default="1490")
        o = p.parse_args(a[1:])
        add(o.name, o.local, o.remote, o.mode, o.tid, o.allowed_src, o.deny_src, o.deny_dst,
            o.unknown_proto, o.mtu)
    elif cmd == "del":
        delete(a[1])
    elif cmd == "show":
        names = a[1:] or sorted(os.listdir(DEBUGFS) if os.path.isdir(DEBUGFS) else [])
        for n in names:
            if n in ("global", "pairs"):
                continue
            print(f"== {n}")
            sys.stdout.write(open(f"{DEBUGFS}/{n}/config").read())
            counters(n)
    elif cmd == "counters":
        counters(a[1], "-a" in a)
    elif cmd in ("global", "pairs"):
        sys.stdout.write(open(f"{DEBUGFS}/{cmd}").read())
    elif cmd == "up":
        up(a[1] if len(a) > 1 else None)
    elif cmd == "down":
        down(a[1] if len(a) > 1 else None)
    else:
        print(__doc__)
        sys.exit(2)


if __name__ == "__main__":
    main()
