search.py

"""
search.py - Real-time search & filter engine for NetSpy

Supported filter syntax (can be combined with spaces):
  tcp / udp / icmp / arp / dns / http    — protocol
  ip:192.168.1.1                         — match src or dst IP (partial ok)
  src:10.0.0.1                           — source IP
  dst:8.8.8.8                            — destination IP
  port:443                               — src or dst port
  sport:1234                             — source port
  dport:80                               — destination port
  flag:SYN / flag:RST / flag:ACK        — TCP flag
  len:>1000 / len:<100 / len:=64        — packet length
  info:google                            — search in info/payload text
  192.168.1.1                            — bare IP: match src or dst (partial)
  443                                    — bare number: match any port
  anything                               — match against any field (fuzzy)
"""

import re


class SearchFilter:
    """
    Parses a search query string and tests packets against it.
    All tokens must match (implicit AND).
    """

    PROTOCOLS = {"tcp", "udp", "icmp", "arp", "dns", "http", "https",
                 "ssh", "ftp", "smtp", "rdp", "ntp", "dhcp", "smb",
                 "ipv6", "other"}

    def __init__(self, query: str = ""):
        self.raw = query.strip()
        self.tokens = self._parse(self.raw)
        self.active = bool(self.raw)
        # Pre-compiled for highlight
        self._highlight_terms = self._extract_highlight_terms()

    def _parse(self, query: str):
        """Return list of (type, value) matcher tuples."""
        tokens = []
        parts = query.split()
        for part in parts:
            part_lower = part.lower()

            # Protocol shorthand
            if part_lower in self.PROTOCOLS:
                tokens.append(("proto", part_lower))
                continue

            # Keyed filters
            if ":" in part:
                key, _, val = part.partition(":")
                key = key.lower()
                val = val.strip()

                if key == "ip":
                    tokens.append(("ip_any", val))
                elif key == "src":
                    tokens.append(("src_ip", val))
                elif key == "dst":
                    tokens.append(("dst_ip", val))
                elif key == "port":
                    tokens.append(("port_any", val))
                elif key in ("sport", "sp"):
                    tokens.append(("src_port", val))
                elif key in ("dport", "dp"):
                    tokens.append(("dst_port", val))
                elif key == "flag":
                    tokens.append(("flag", val.upper()))
                elif key == "len":
                    tokens.append(("len", val))
                elif key == "info":
                    tokens.append(("info", val.lower()))
                elif key == "proto":
                    tokens.append(("proto", val.lower()))
                else:
                    # Unknown key → fuzzy
                    tokens.append(("fuzzy", part.lower()))
                continue

            # Bare number → port
            if part.isdigit():
                tokens.append(("port_any", part))
                continue

            # Bare IP-like
            if re.match(r"^\d{1,3}(\.\d{1,3}){0,3}$", part):
                tokens.append(("ip_any", part))
                continue

            # Fallback fuzzy
            tokens.append(("fuzzy", part.lower()))

        return tokens

    def _extract_highlight_terms(self):
        """Pull out plain text terms for UI highlighting."""
        terms = []
        for ttype, val in self.tokens:
            if ttype in ("fuzzy", "info", "ip_any", "src_ip", "dst_ip"):
                terms.append(val.lower())
            elif ttype in ("port_any", "src_port", "dst_port"):
                terms.append(val)
        return terms

    def _match_len(self, pkt_len: int, expr: str) -> bool:
        """Match len:>1000 / len:<100 / len:=64 / len:1000"""
        try:
            if expr.startswith(">"):
                return pkt_len > int(expr[1:])
            elif expr.startswith("<"):
                return pkt_len < int(expr[1:])
            elif expr.startswith("="):
                return pkt_len == int(expr[1:])
            else:
                return pkt_len == int(expr)
        except ValueError:
            return False

    def matches(self, pkt: dict) -> bool:
        """Return True if packet matches ALL tokens."""
        if not self.tokens:
            return True

        proto = pkt.get("proto", "").lower()
        src_ip = pkt.get("src_ip", "")
        dst_ip = pkt.get("dst_ip", "")
        src_port = str(pkt.get("src_port", "") or "")
        dst_port = str(pkt.get("dst_port", "") or "")
        flags = [f.upper() for f in pkt.get("flags", [])]
        info = pkt.get("info", "").lower()
        pkt_len = pkt.get("length", 0)

        # Build a single "all fields" string for fuzzy matching
        all_fields = " ".join(filter(None, [
            proto, src_ip, dst_ip, src_port, dst_port, info,
            ",".join(flags)
        ])).lower()

        for ttype, val in self.tokens:
            if ttype == "proto":
                if val not in proto:
                    return False

            elif ttype == "ip_any":
                if val not in src_ip and val not in dst_ip:
                    return False

            elif ttype == "src_ip":
                if val not in src_ip:
                    return False

            elif ttype == "dst_ip":
                if val not in dst_ip:
                    return False

            elif ttype == "port_any":
                if val != src_port and val != dst_port:
                    return False

            elif ttype == "src_port":
                if val != src_port:
                    return False

            elif ttype == "dst_port":
                if val != dst_port:
                    return False

            elif ttype == "flag":
                if val not in flags:
                    return False

            elif ttype == "len":
                if not self._match_len(pkt_len, val):
                    return False

            elif ttype == "info":
                if val not in info:
                    return False

            elif ttype == "fuzzy":
                if val not in all_fields:
                    return False

        return True

    def filter_packets(self, packets: list) -> list:
        """Return filtered list of packets."""
        if not self.active:
            return packets
        return [p for p in packets if self.matches(p)]

    @property
    def summary(self) -> str:
        """Human-readable summary of active filters."""
        if not self.tokens:
            return ""
        parts = []
        for ttype, val in self.tokens:
            if ttype == "proto":
                parts.append(f"proto={val.upper()}")
            elif ttype == "ip_any":
                parts.append(f"ip~{val}")
            elif ttype == "src_ip":
                parts.append(f"src~{val}")
            elif ttype == "dst_ip":
                parts.append(f"dst~{val}")
            elif ttype == "port_any":
                parts.append(f"port={val}")
            elif ttype == "src_port":
                parts.append(f"sport={val}")
            elif ttype == "dst_port":
                parts.append(f"dport={val}")
            elif ttype == "flag":
                parts.append(f"flag={val}")
            elif ttype == "len":
                parts.append(f"len{val}")
            elif ttype == "info":
                parts.append(f'info~"{val}"')
            elif ttype == "fuzzy":
                parts.append(f'"{val}"')
        return "  ".join(parts)