Blame for runtime/python/dibaguard/shield/rules.py

"""Rules: the operator's own protection, written in plain language (Persian or English), active from the first
second - so there is no "unprotected while it learns" window.

A rule is interpreted into a precise, deterministic predicate over what the engine already measures (no model, no
latency). Interpretation is done by a tiny built-in parser for the common shapes, and - when an LLM is wired in -
by the LLM, which must return the same JSON schema; either way the layer shows its interpretation back so the
operator can confirm it before it acts.

Rule JSON:
  {"id": "...", "text": "<original>", "scope": "flow|http|source", "when": {<conditions>}, "action": "block|rate_limit|challenge|monitor|allow",
   "family": "<name>", "ttl_s": 1800, "note": "<how it was interpreted>"}
Conditions understood by the parser and the engine (all optional, AND-ed):
  port / ports (list), proto (tcp|udp|icmp), path_contains, ua_contains, country (if a geo hook is set), method,
  flows_per_min_over, new_ports_over, noreply_over, status, request_rate_over, service (name), src_in (cidr list)
"""
from __future__ import annotations

import ipaddress
import json
import re

ACTIONS = ("block", "rate_limit", "challenge", "monitor", "allow")
# a tiny bilingual lexicon so the built-in parser handles the common operator rules without an LLM
_ACT = {"block": "block", "\u0645\u0633\u062f\u0648\u062f": "block", "\u0628\u0628\u0646\u062f": "block", "deny": "block", "drop": "block", "\u0631\u062f": "block",
        "rate": "rate_limit", "\u0645\u062d\u062f\u0648\u062f": "rate_limit", "throttle": "rate_limit", "\u0646\u0631\u062e": "rate_limit",
        "challenge": "challenge", "\u0686\u0627\u0644\u0634": "challenge", "captcha": "challenge", "\u06a9\u067e\u0686\u0627": "challenge",
        "allow": "allow", "\u0627\u062c\u0627\u0632\u0647": "allow", "\u0645\u062c\u0627\u0632": "allow", "whitelist": "allow", "\u0633\u0641\u06cc\u062f": "allow",
        "monitor": "monitor", "\u0641\u0642\u0637 \u06af\u0632\u0627\u0631\u0634": "monitor", "\u067e\u0627\u06cc\u0634": "monitor", "\u06af\u0632\u0627\u0631\u0634": "monitor"}
_PORT = re.compile(r"(?:port|\u067e\u0648\u0631\u062a)\s*#?\s*(\d[\d,\s]*)", re.I)
_FILLER = r"(?:containing|contains|contain|with|=|:|\u062d\u0627\u0648\u06cc|\u0634\u0627\u0645\u0644|\u06a9\u0647|\u0628\u0631\u0627\u0628\u0631|is|to|\u0628\u0631\u0627\u06cc)\s+"
_PATH = re.compile(r"(?:path|url|\u0645\u0633\u06cc\u0631|\u0622\u062f\u0631\u0633)\s*(?:" + _FILLER + r")?([^\s,\u060c]+)", re.I)
_BARE_PATH = re.compile(r"(?<![\w])(/[\w./\-]+)")
_UA = re.compile(r"(?:user-?agent|ua|\u0639\u0627\u0645\u0644|\u0645\u0631\u0648\u0631\u06af\u0631)\s*(?:" + _FILLER + r")?([\w./\-]+)", re.I)
_RATE = re.compile(r"(?:more than|over|\u0628\u06cc\u0634 \u0627\u0632|\u0628\u06cc\u0634\u062a\u0631 \u0627\u0632)\s*(\d+)\s*(?:req|request|\u062f\u0631\u062e\u0648\u0627\u0633\u062a|/|per|\u062f\u0631|flows?|connections?|\u0627\u062a\u0635\u0627\u0644)", re.I)
_PORTS_N = re.compile(r"(?:more than|over|\u0628\u06cc\u0634 \u0627\u0632|\u0628\u06cc\u0634\u062a\u0631 \u0627\u0632)\s*(\d+)\s*(?:ports?|\u067e\u0648\u0631\u062a)", re.I)
_ATTEMPTS = re.compile(r"(?:more than|over|\u0628\u06cc\u0634 \u0627\u0632|\u0628\u06cc\u0634\u062a\u0631 \u0627\u0632)\s*(\d+)\s*(?:attempts?|tries|logins?|\u062a\u0644\u0627\u0634|\u0628\u0627\u0631|\u0648\u0631\u0648\u062f)", re.I)
_CIDR = re.compile(r"\b(\d{1,3}(?:\.\d{1,3}){3}(?:/\d{1,2})?)\b")
_SERVICES = {"ssh": 22, "rdp": 3389, "telnet": 23, "ftp": 21, "smb": 445, "redis": 6379, "mongodb": 27017,
             "elasticsearch": 9200, "mysql": 3306, "postgres": 5432, "vnc": 5900, "http": 80, "https": 443}
_SCAN_WORDS = ("scan", "\u0627\u0633\u06a9\u0646", "port scan", "\u0627\u0633\u06a9\u0646 \u067e\u0648\u0631\u062a")
_BRUTE_WORDS = ("brute", "brute force", "\u0628\u0631\u0648\u062a", "\u0641\u0648\u0631\u0633", "\u0631\u0645\u0632", "password")


class Rule:
    def __init__(self, d: dict):
        self.id = d.get("id", "")
        self.text = d.get("text", "")
        self.scope = d.get("scope", "source")
        self.when = d.get("when", {})
        self.action = d.get("action", "monitor")
        self.family = d.get("family", "policy_rule")
        self.ttl_s = float(d.get("ttl_s", 1800))
        self.note = d.get("note", "")
        if self.action not in ACTIONS:
            raise ValueError(f"rule action must be one of {ACTIONS}: {self.action!r}")

    def as_dict(self) -> dict:
        return {"id": self.id, "text": self.text, "scope": self.scope, "when": self.when, "action": self.action,
                "family": self.family, "ttl_s": self.ttl_s, "note": self.note}

    # -- evaluation (deterministic; no model). ctx is a dict the layer fills per flow / request.
    def matches(self, ctx: dict) -> bool:
        w = self.when
        if "ports" in w and int(ctx.get("port", -1)) not in w["ports"]:
            return False
        if "proto" in w and str(ctx.get("proto", "")).lower() != w["proto"]:
            return False
        if "path_contains" in w and w["path_contains"] not in str(ctx.get("path", "")).lower():
            return False
        if "ua_contains" in w and w["ua_contains"] not in str(ctx.get("ua", "")).lower():
            return False
        if "method" in w and str(ctx.get("method", "")).lower() != w["method"]:
            return False
        if "status" in w and int(ctx.get("status", 0)) != int(w["status"]):
            return False
        if "src_in" in w:
            try:
                a = ipaddress.ip_address(ctx.get("src", ""))
                if not any(a in ipaddress.ip_network(c, strict=False) for c in w["src_in"]):
                    return False
            except ValueError:
                return False
        if "flows_per_min_over" in w and ctx.get("flows_per_min", 0) <= w["flows_per_min_over"]:
            return False
        if "request_rate_over" in w and ctx.get("flows_per_min", 0) <= w["request_rate_over"]:
            return False
        if "new_ports_over" in w and ctx.get("new_ports", 0) <= w["new_ports_over"]:
            return False
        if "noreply_over" in w and ctx.get("noreply", 0) <= w["noreply_over"]:
            return False
        if "repeat_over" in w and ctx.get("repeat", 0) <= w["repeat_over"]:
            return False
        if "family" in w and str(ctx.get("family", "")) != w["family"]:
            return False
        return bool(w)   # an empty rule never matches (prevents an accidental block-all)


def parse(text: str) -> dict:
    """built-in interpreter: a plain-language rule -> rule dict (best effort; the LLM path returns the same shape).
    Examples it understands:
      '\u0645\u0633\u062f\u0648\u062f \u06a9\u0646 \u0647\u0631 \u06a9\u0633\u06cc \u06a9\u0647 \u067e\u0648\u0631\u062a \u06f2\u06f2 \u0631\u0627 brute force \u0645\u06cc\u200c\u06a9\u0646\u062f'       -> block, service ssh, scan/brute behaviour
      'block any source scanning more than 50 ports per minute' -> block, new_ports behaviour
      'allow 10.0.0.0/8'                                        -> allow, src_in
      'rate limit requests to /api over 100 per minute'        -> rate_limit, path /api, request_rate
      'challenge user-agent containing bot'                     -> challenge, ua bot
    """
    low = _norm(text)
    action = next((a for k, a in _ACT.items() if k in low), "monitor")
    when: dict = {}
    scope = "source"
    fam = "policy_rule"
    m = _PORT.search(text)
    if m:
        when["ports"] = [int(x) for x in re.findall(r"\d+", m.group(1))]
    for svc, port in _SERVICES.items():
        if re.search(rf"\b{svc}\b", low):
            when.setdefault("ports", []).append(port)
            when["ports"] = sorted(set(when["ports"]))
    mp = _PATH.search(text) or _BARE_PATH.search(text)
    if mp and ("/" in mp.group(1) or "." in mp.group(1)):
        when["path_contains"] = mp.group(1).lower().strip("'\"")
        scope = "http"
    mu = _UA.search(text)
    if mu:
        when["ua_contains"] = mu.group(1).lower().strip("'\"")
        scope = "http"
    mc = _CIDR.search(text)
    if mc:
        when["src_in"] = [mc.group(1)]
    mr = _RATE.search(low)
    if mr:
        key = "request_rate_over" if scope == "http" else "flows_per_min_over"
        when[key] = int(mr.group(1))
    mp2 = _PORTS_N.search(low)
    if any(w in low for w in _SCAN_WORDS) or mp2:
        when["new_ports_over"] = int(mp2.group(1)) if mp2 else when.get("new_ports_over", 20)
        fam = "portscan"
    if any(w in low for w in _BRUTE_WORDS):
        # brute force = the same service again and again (an SSH guess completes its handshake: replies are normal)
        when["repeat_over"] = when.get("repeat_over", 0.6)
        mn = _ATTEMPTS.search(low)
        when["flows_per_min_over"] = int(mn.group(1)) if mn else when.get("flows_per_min_over", 10)
        fam = "bruteforce"
    note = _describe(action, when, scope)
    return {"text": text, "scope": scope, "when": when, "action": action, "family": fam, "note": note}


def _norm(s: str) -> str:
    s = s.lower()
    for a, b in zip("\u06f0\u06f1\u06f2\u06f3\u06f4\u06f5\u06f6\u06f7\u06f8\u06f9\u0660\u0661\u0662\u0663\u0664\u0665\u0666\u0667\u0668\u0669", "01234567890123456789"):
        s = s.replace(a, b)
    return s


def _describe(action, when, scope) -> str:
    if not when:
        return "could not interpret any condition - this rule will be inactive until edited"
    parts = []
    if "ports" in when:
        parts.append(f"destination port in {when['ports']}")
    if "path_contains" in when:
        parts.append(f"path contains '{when['path_contains']}'")
    if "ua_contains" in when:
        parts.append(f"user-agent contains '{when['ua_contains']}'")
    if "src_in" in when:
        parts.append(f"source in {when['src_in']}")
    if "flows_per_min_over" in when:
        parts.append(f"more than {when['flows_per_min_over']} flows/min from the source")
    if "request_rate_over" in when:
        parts.append(f"more than {when['request_rate_over']} requests/min from the source")
    if "new_ports_over" in when:
        parts.append(f"more than {when['new_ports_over']} new ports touched (scanning)")
    if "noreply_over" in when:
        parts.append(f"over {when['noreply_over']:.0%} of flows unanswered (failed connects)")
    if "repeat_over" in when:
        parts.append(f"over {when['repeat_over']:.0%} of flows to the same service again (repeated attempts)")
    return f"{action.upper()} a {scope} when " + " and ".join(parts)


LLM_SCHEMA_PROMPT = """You translate one network/web security rule written by an operator (Persian or English) into
a JSON object, and nothing else. Schema:
{"scope":"flow|http|source","when":{...},"action":"block|rate_limit|challenge|monitor|allow","family":"<name>","note":"<one sentence back-translation of what you understood>"}
Allowed keys in "when": ports (list of ints), proto ("tcp"|"udp"|"icmp"), path_contains (str), ua_contains (str),
method (str), status (int), src_in (list of CIDR), flows_per_min_over (num), request_rate_over (num),
new_ports_over (num per minute), noreply_over (0..1), repeat_over (0..1, share of flows to an already-used host:port: brute force), family (str). Use only conditions the text states. Never invent a
block-all. Return ONLY the JSON."""


class RuleSet:
    def __init__(self, rules: list | None = None):
        self.rules: list[Rule] = [Rule(r) if isinstance(r, dict) else r for r in (rules or [])]

    def add(self, text: str, interpret, idx: int) -> Rule:
        d = interpret(text) if interpret else parse(text)
        d["id"] = d.get("id") or f"R{idx:03d}"
        taken = {r.id for r in self.rules}
        if d["id"] in taken:                         # ids stay unique after a removal
            d["id"] = f"R{1 + max((int(x[1:]) for x in taken if re.fullmatch(r'R\d+', x)), default=0):03d}"
        r = Rule(d)
        self.rules.append(r)
        return r

    def remove(self, rule_id: str) -> Rule | None:
        """take one rule out (by id); returns it, or None when there is no such rule"""
        for i, r in enumerate(self.rules):
            if r.id == rule_id:
                return self.rules.pop(i)
        return None

    def decide(self, ctx: dict) -> Rule | None:
        """first matching rule wins (operator order); 'allow' short-circuits to no action"""
        for r in self.rules:
            if r.scope in (ctx.get("scope"), "source") and r.matches(ctx):
                return r
        return None