#!/usr/bin/env python3
"""ShadowSense AI laptop agent.

Watches the network flows of THIS computer and sends them to your ShadowSense AI account for threat detection, in near real time.

What it sends: per-connection METADATA only: IP addresses, ports, protocol, duration, byte and packet counts.
What it never sends: packet payloads, web page contents, URLs, DNS names, passwords or files.
Run it only on computers you own or are authorised to monitor.

1. Sign in at <your ShadowSense AI site>/connect, click "Add this computer" and copy the token (it is shown once).
2. Run (needs administrator / root, because capturing packets is a privileged operation):

    Linux   : sudo python3 shadowsense_agent.py --server https://shadowsenseai.in --token ssd_XXXX      (no extra packages needed)
    macOS   : pip3 install scapy ; sudo python3 shadowsense_agent.py --server ... --token ...
    Windows : install Npcap (https://npcap.com) ; pip install scapy ; run as Administrator:
              python shadowsense_agent.py --server ... --token ...

Other options: --iface NAME, --list-interfaces, --pcap FILE (replay a classic .pcap, handy for testing), --dry-run (print flows, send nothing).
The token and server can also come from the SHADOWSENSE_TOKEN and SHADOWSENSE_URL environment variables.
"""
import argparse, collections, ipaddress, json, os, socket, struct, sys, threading, time, urllib.error, urllib.request, uuid

VERSION = "1.0"
TCP, UDP = 6, 17
FIN, SYN, RST, ACK = 0x01, 0x02, 0x04, 0x10


# ----------------------------------------------------------------------------- packet parsing (standard library only)
def parse_ip(b):
    """Parse an IPv4/IPv6 packet. Returns (src, dst, proto, sport, dport, payload_len, tcp_flags) for TCP/UDP, else None."""
    if not b: return None
    v = b[0] >> 4
    if v == 4:
        if len(b) < 20: return None
        ihl = (b[0] & 15) * 4; total = struct.unpack("!H", b[2:4])[0]; proto = b[9]
        if struct.unpack("!H", b[6:8])[0] & 0x1FFF: return None  # non-first fragment: no L4 header
        src, dst = socket.inet_ntop(socket.AF_INET, b[12:16]), socket.inet_ntop(socket.AF_INET, b[16:20])
        end = total if 0 < total <= len(b) else len(b)  # total==0 happens with segmentation offload; fall back to captured length
        l4 = b[ihl:end]
    elif v == 6:
        if len(b) < 40: return None
        proto = b[6]; plen = struct.unpack("!H", b[4:6])[0]; off = 40
        src, dst = socket.inet_ntop(socket.AF_INET6, b[8:24]), socket.inet_ntop(socket.AF_INET6, b[24:40])
        for _ in range(6):  # skip hop-by-hop, routing and destination-options headers
            if proto in (0, 43, 60) and len(b) >= off + 8: proto = b[off]; off += (b[off + 1] + 1) * 8
            else: break
        if proto == 44: return None  # fragment
        end = 40 + plen if 0 < plen and 40 + plen <= len(b) else len(b)
        l4 = b[off:end]
    else: return None
    if proto == TCP:
        if len(l4) < 20: return None
        sp, dp = struct.unpack("!HH", l4[:4]); hdr = (l4[12] >> 4) * 4
        return src, dst, TCP, sp, dp, max(0, len(l4) - hdr), l4[13]
    if proto == UDP:
        if len(l4) < 8: return None
        sp, dp = struct.unpack("!HH", l4[:4]); return src, dst, UDP, sp, dp, max(0, len(l4) - 8), 0
    return None

def parse_frame(b, linktype=1):
    """Link-layer frame -> parse_ip() result. Supports Ethernet(+VLAN), raw IP, Linux cooked (SLL) and loopback/null."""
    if linktype == 1:
        if len(b) < 14: return None
        et = struct.unpack("!H", b[12:14])[0]; off = 14
        while et in (0x8100, 0x88A8) and len(b) >= off + 4: et = struct.unpack("!H", b[off + 2:off + 4])[0]; off += 4
        return parse_ip(b[off:]) if et in (0x0800, 0x86DD) else None
    if linktype == 101: return parse_ip(b)
    if linktype == 113: return parse_ip(b[16:]) if len(b) > 16 and b[14:16] in (b"\x08\x00", b"\x86\xdd") else None
    if linktype in (0, 108): return parse_ip(b[4:])
    return None


# ----------------------------------------------------------------------------- flow table
class Flow:
    __slots__ = ("o", "r", "proto", "t0", "t1", "ob", "rb", "op", "rp", "fin_o", "fin_r", "done")
    def __init__(self, o, r, proto, t): self.o, self.r, self.proto, self.t0, self.t1 = o, r, proto, t, t; self.ob = self.rb = self.op = self.rp = 0; self.fin_o = self.fin_r = self.done = False

class FlowTable:
    """Groups packets into bidirectional flows and emits Zeek conn.log style records (the format the ShadowSense backend ingests).
    A flow ends on TCP RST, after both sides sent FIN, after `idle` seconds of silence, or is cut every `active` seconds if it is long-lived.
    The table is bounded: under a flood the oldest flows are emitted early instead of exhausting memory."""
    def __init__(self, idle=15.0, active=60.0, max_flows=50000, skip=None):
        self.idle, self.active, self.max, self.skip = idle, active, max_flows, skip or (lambda ip: False)
        self.f = collections.OrderedDict(); self.out = []; self.packets = 0; self.raw = 0; self.ignored = 0; self.closed = {}  # closed: recently ended TCP flows, so the final ACK doesn't start a bogus 1-packet flow

    def add(self, ts, src, dst, proto, sp, dp, plen, flags=0):
        if self.skip(src) or self.skip(dst): self.ignored += 1; return
        a, b = (src, sp), (dst, dp); key = (proto, a, b) if a <= b else (proto, b, a); fl = self.f.get(key); self.packets += 1
        if fl is None and proto == TCP and ts - self.closed.get(key, -1e9) < 5 and not flags & SYN: return
        if fl is None:
            if len(self.f) >= self.max: self._emit(self.f.popitem(last=False)[1])
            # first packet decides who "originated" the flow; a lone SYN-ACK means we joined mid-handshake, so flip
            fl = Flow(b, a, proto, ts) if proto == TCP and flags & SYN and flags & ACK else Flow(a, b, proto, ts); self.f[key] = fl
        else: self.f.move_to_end(key)
        fwd = (src, sp) == fl.o; fl.t1 = max(fl.t1, ts)
        if fwd: fl.op += 1; fl.ob += plen
        else: fl.rp += 1; fl.rb += plen
        if proto == TCP:
            if flags & RST: fl.done = True
            if flags & FIN:
                if fwd: fl.fin_o = True
                else: fl.fin_r = True
                if fl.fin_o and fl.fin_r: fl.done = True
        if fl.done: self._emit(self.f.pop(key)); self.closed[key] = ts

    def sweep(self, now):
        """Emit flows that went idle or have been running too long. Call periodically with the current clock."""
        if len(self.closed) > 5000: self.closed = {k: t for k, t in self.closed.items() if now - t < 5}
        for k in [k for k, fl in self.f.items() if now - fl.t1 > self.idle or now - fl.t0 > self.active]: self._emit(self.f.pop(k))
        return self.take()

    def flush(self):
        for k in list(self.f): self._emit(self.f.pop(k))
        return self.take()

    def take(self): o, self.out = self.out, []; return o

    def _emit(self, fl):
        self.out.append({"uid": uuid.uuid4().hex[:16], "ts": fl.t0, "id.orig_h": fl.o[0], "id.orig_p": fl.o[1], "id.resp_h": fl.r[0], "id.resp_p": fl.r[1],
                         "proto": "tcp" if fl.proto == TCP else "udp", "duration": round(fl.t1 - fl.t0, 6), "orig_bytes": fl.ob, "resp_bytes": fl.rb,
                         "orig_pkts": fl.op, "resp_pkts": fl.rp})


def noise_filter():
    """Skip multicast/broadcast chatter (mDNS, SSDP...). Returns a predicate on an IP string."""
    cache = {}
    def skip(ip):
        r = cache.get(ip)
        if r is None:
            try: a = ipaddress.ip_address(ip.split("%")[0]); r = a.is_multicast or ip == "255.255.255.255"
            except ValueError: r = True
            if len(cache) < 100000: cache[ip] = r
        return r
    return skip


# ----------------------------------------------------------------------------- capture backends
def replay_pcap(path, table):
    """Read a classic libpcap file (not pcapng: convert with `editcap -F pcap in.pcapng out.pcap`). Uses packet time as the clock."""
    with open(path, "rb") as f:
        h = f.read(24)
        if len(h) < 24: sys.exit("not a pcap file")
        magic = h[:4]
        if magic in (b"\xa1\xb2\xc3\xd4", b"\xa1\xb2\x3c\x4d"): e = ">"
        elif magic in (b"\xd4\xc3\xb2\xa1", b"\x4d\x3c\xb2\xa1"): e = "<"
        else: sys.exit("not a classic pcap file (pcapng is not supported; convert it with: editcap -F pcap in.pcapng out.pcap)")
        nano = magic in (b"\xa1\xb2\x3c\x4d", b"\x4d\x3c\xb2\xa1"); link = struct.unpack(e + "I", h[20:24])[0]; n = 0
        while True:
            r = f.read(16)
            if len(r) < 16: break
            sec, frac, cap, _ = struct.unpack(e + "IIII", r); data = f.read(cap); ts = sec + frac / (1e9 if nano else 1e6); p = parse_frame(data, link)
            if p: table.add(ts, *p[:3], p[3], p[4], p[5], p[6]); n += 1
            if n % 2000 == 0: table.sweep(ts)
    return n

def default_iface_linux():
    try:
        for line in open("/proc/net/route").read().splitlines()[1:]:
            c = line.split()
            if c[1] == "00000000": return c[0]
    except OSError: pass

def capture_afpacket(table, lock, iface, stop):
    """Linux, standard library only. Needs root."""
    s = socket.socket(socket.AF_PACKET, socket.SOCK_RAW, socket.ntohs(3)); s.settimeout(1.0)
    if iface: s.bind((iface, 0))
    while not stop.is_set():
        try: data = s.recv(65535)
        except socket.timeout: continue
        table.raw += 1; p = parse_frame(data, 1)
        if p:
            with lock: table.add(time.time(), *p[:3], p[3], p[4], p[5], p[6])

def capture_scapy(table, lock, iface, stop):
    try: from scapy.all import sniff, conf, IP, IPv6, TCP as sTCP, UDP as sUDP
    except ImportError: sys.exit("Install scapy first (pip install scapy). On Windows also install Npcap from https://npcap.com")
    if iface is None and sys.platform.startswith("win"): iface = "all"  # on Windows the automatic adapter choice is often wrong (VPN/virtual adapters)
    if iface == "all":  # every adapter that has an IP address (useful when the automatic choice is wrong, e.g. VPN or virtual adapters)
        names = []
        for i in conf.ifaces.values():
            try:
                if getattr(i, "ip", "") not in ("", "0.0.0.0", "127.0.0.1"): names.append(getattr(i, "network_name", None) or i.name)
            except Exception: pass
        iface = names or None
    print("capturing on:", iface if iface else f"{conf.iface} (automatic)", flush=True)
    def cb(pkt):
        with lock: table.raw += 1
        ip = pkt.getlayer(IP) or pkt.getlayer(IPv6)
        if ip is None: return
        l4 = pkt.getlayer(sTCP) or pkt.getlayer(sUDP)
        if l4 is None: return
        is_tcp = pkt.haslayer(sTCP); plen = len(bytes(l4.payload))
        with lock: table.add(float(pkt.time), ip.src, ip.dst, TCP if is_tcp else UDP, int(l4.sport), int(l4.dport), plen, int(l4.flags) if is_tcp else 0)
    sniff(iface=iface, store=False, prn=cb, filter="tcp or udp", stop_filter=lambda _: stop.is_set())

def list_interfaces():
    if sys.platform.startswith("linux"):
        try: print("\n".join(sorted(os.listdir("/sys/class/net")))); return
        except OSError: pass
    try:
        from scapy.all import conf
        for i in conf.ifaces.values(): print(f"name={i.name!r}  description={getattr(i, 'description', '')!r}  ip={getattr(i, 'ip', '')!r}")
    except ImportError: print("Install scapy to list interfaces (pip install scapy).")


# ----------------------------------------------------------------------------- sending
class Sender:
    def __init__(self, server, token, dry=False):
        self.url = server.rstrip("/") + "/api/agent/flows"; self.token = token; self.dry = dry
        self.q = collections.deque(maxlen=20000); self.sent = self.failed = self.dropped = self.rejected = 0; self.wait = 1.0; self.next_try = 0.0
    def add(self, recs):
        for r in recs:
            if len(self.q) == self.q.maxlen: self.dropped += 1  # server unreachable for a long time: oldest flows are dropped, counted
            self.q.append(r)
    def push(self):
        if time.time() < self.next_try: return
        while self.q:
            batch = [self.q.popleft() for _ in range(min(500, len(self.q)))]
            if self.dry:
                for r in batch: print(json.dumps(r), flush=True)
                self.sent += len(batch); continue
            req = urllib.request.Request(self.url, json.dumps(batch).encode(), {"Authorization": "Bearer " + self.token, "Content-Type": "application/json", "User-Agent": "shadowsense-agent/" + VERSION})
            try:
                with urllib.request.urlopen(req, timeout=20) as resp: res = json.loads(resp.read())
                self.sent += len(batch); self.rejected += res.get("rejected", 0); self.wait = 1.0
            except urllib.error.HTTPError as e:
                if e.code == 401: sys.exit("The server rejected this device token (revoked or wrong). Create a new one at /connect.")
                self.q.extendleft(reversed(batch)); self.failed += 1; self.wait = min(60, self.wait * 2); self.next_try = time.time() + self.wait
                print(f"server answered HTTP {e.code}; retrying in {self.wait:.0f}s", flush=True); return
            except Exception as e:
                self.q.extendleft(reversed(batch)); self.failed += 1; self.wait = min(60, self.wait * 2); self.next_try = time.time() + self.wait
                print(f"could not reach the server ({type(e).__name__}); retrying in {self.wait:.0f}s", flush=True); return


def main(argv=None):
    ap = argparse.ArgumentParser(description="ShadowSense AI laptop agent")
    ap.add_argument("--server", default=os.getenv("SHADOWSENSE_URL", "https://shadowsenseai.in")); ap.add_argument("--token", default=os.getenv("SHADOWSENSE_TOKEN"))
    ap.add_argument("--iface"); ap.add_argument("--pcap"); ap.add_argument("--dry-run", action="store_true"); ap.add_argument("--list-interfaces", action="store_true")
    ap.add_argument("--backend", choices=["auto", "afpacket", "scapy"], default="auto"); ap.add_argument("--interval", type=float, default=3.0, help="seconds between uploads")
    ap.add_argument("--idle", type=float, default=15.0, help="seconds of silence that end a flow"); a = ap.parse_args(argv)
    if a.list_interfaces: return list_interfaces()
    if not a.dry_run and not a.token: sys.exit("Missing token. Create one at /connect and pass --token (or set SHADOWSENSE_TOKEN).")

    skip = noise_filter(); host = a.server.split("://")[-1].split("/")[0].split(":")[0]; sips = set()
    try: sips = {x[4][0] for x in socket.getaddrinfo(host, None)}
    except OSError: pass
    srv_port = int((a.server.split("://")[-1].split("/")[0].split(":") + [443 if a.server.startswith("https") else 80])[1])
    table = FlowTable(idle=a.idle, skip=skip); sender = Sender(a.server, a.token or "", a.dry_run)

    if a.pcap:
        n = replay_pcap(a.pcap, table); sender.add(table.flush()); sender.next_try = 0; sender.push()
        print(f"replayed {n} packets -> sent {sender.sent} flows (failed batches {sender.failed}, rejected {sender.rejected})"); return

    print(f"ShadowSense agent {VERSION}: capturing flow metadata (no payloads) and sending it to {a.server}. Press Ctrl+C to stop.", flush=True)
    lock = threading.Lock(); stop = threading.Event()
    # our own uploads to the server must not be analysed (would create a feedback loop)
    base_add = table.add
    def add(ts, src, dst, proto, sp, dp, plen, flags=0):
        if proto == TCP and ((src in sips and sp == srv_port) or (dst in sips and dp == srv_port)): return
        base_add(ts, src, dst, proto, sp, dp, plen, flags)
    table.add = add

    def worker():
        last = time.time()
        while not stop.is_set():
            time.sleep(a.interval)
            with lock: recs = table.sweep(time.time())
            sender.add(recs); sender.push()
            if time.time() - last > 30:
                if table.raw == 0: print("WARNING: no packets seen at all. Wrong adapter? Try:  --iface all   (or --list-interfaces to pick one)", flush=True)
                print(f"raw_seen={table.raw} packets={table.packets} sent={sender.sent} queued={len(sender.q)} dropped={sender.dropped} failed_pushes={sender.failed}", flush=True); last = time.time()
    threading.Thread(target=worker, daemon=True).start()
    backend = a.backend if a.backend != "auto" else ("afpacket" if sys.platform.startswith("linux") else "scapy")
    try:
        if backend == "afpacket": capture_afpacket(table, lock, a.iface or default_iface_linux(), stop)
        else: capture_scapy(table, lock, a.iface, stop)
    except PermissionError: sys.exit("Packet capture needs administrator/root rights. Re-run with sudo (Linux/macOS) or as Administrator (Windows).")
    except KeyboardInterrupt: pass
    finally:
        stop.set()
        with lock: sender.add(table.flush())
        sender.next_try = 0; sender.push()

if __name__ == "__main__": main()