"""SQL-context analysis of HTTP parameters: does this value change a query's structure if an application puts it in one?

The database tier (shield/database.py, dbnet.py) understands queries; this module brings that understanding to the
edge. Every parameter (query string, body, cookies, the headers attackers like to stuff; keys too - payloads hide in
keys) is decoded as a server would (URL decoding up to three layers, HTML entities, %uXXXX, \\xNN, \\uNNNN, null
bytes), then placed in the contexts applications build queries with:

    numeric      ... WHERE id = <v>              string       ... WHERE s = '<v>'
    parenthesis  ... WHERE (id = <v>)  / ('<v>')  double-quote ... WHERE s = "<v>"      LIKE  ... LIKE '%<v>%'

and the resulting query is judged by the database tier's signatures (on the query with its string literals removed,
so a value that stays inside its quotes is just data) and, optionally, by DbNet on the query's structure. A value is
an injection only if it breaks out of at least one context *and* the result shows an attack's shape (a tautology, a
UNION extraction, a stacked statement, a time delay, an error-based extractor, a comment cut). O'Reilly, "rock 'n'
roll", Persian text, JSON, URLs stay data.

    hits = analyse(text)    # canonical request text ("method path | q:k=v | b:k=v | c:k=v | h:k=v")
    -> [{"field": "q:id", "value": "...", "context": "numeric", "families": ["injection"], "dbnet": 3.2}]
"""
from __future__ import annotations

import html
import re
from functools import lru_cache
from urllib.parse import unquote_plus

from dibaguard.shield.database import classify

CONTEXTS = {
    "numeric": "SELECT a FROM t WHERE id = {v}",
    "string": "SELECT a FROM t WHERE s = '{v}'",
    "dquote": "SELECT a FROM t WHERE s = \"{v}\"",
    "paren_num": "SELECT a FROM t WHERE (id = {v})",
    "paren_str": "SELECT a FROM t WHERE (s = '{v}')",
    "like": "SELECT a FROM t WHERE s LIKE '%{v}%'",
}
SERIOUS = ("injection", "rce", "privilege", "destructive")
_SWAP = str.maketrans({"'": '"', '"': "'"})
_LITERAL = ("string", "dquote", "paren_str", "like")         # contexts that put the value inside a string literal
# a value with none of these cannot change a query's structure
_PRE = re.compile(r"['\"`;()#=<>|&*\\]|--|/\*|\b(select|union|sleep|benchmark|waitfor|or|and|xor|having|order|group|exec|declare|"
                  r"extractvalue|updatexml|pg_sleep|information_schema|concat|char|chr|cast|convert|drop|insert|delete|update|like|rlike|regexp)\b", re.I)
# the gate a value must pass before any context is built: a character that can close a literal or cut a statement, a
# word only SQL uses, or a boolean joined to a number / quote / parenthesis (user agents and plain text never pass)
_GATE = re.compile(r"['\"`#]|--|/\*|\b(select|union|sleep|benchmark|waitfor|pg_sleep|having|exec|execute|declare|xp_\w+|information_schema|"
                   r"extractvalue|updatexml|dual|concat|char|chr|rlike|regexp)\b|\b(or|and|xor)\b\s*[\d'\"(@]|&&|\|\|", re.I)
_STRONG = re.compile(r"['\"`#;]|--|/\*|\b(union|select|sleep|benchmark|or|and|xor|from|where|exec|waitfor|extractvalue|updatexml)\b", re.I)
_PCTU = re.compile(r"%u([0-9a-fA-F]{4})")
_HEXESC = re.compile(r"\\x([0-9a-fA-F]{2})")
_UESC = re.compile(r"\\u([0-9a-fA-F]{4})")
# headers whose values are protocol syntax, not data an application puts in a query (exact names; "sec-*" by prefix).
# Accept-Language is *not* among them: it is data (RT-09 put it into the text, V-07 made this tier read it)
SKIP_HEADERS = ("h:accept", "h:accept-encoding", "h:content-type", "h:cache-control", "h:content-length", "h:connection")
_SKIP_PREFIX = "h:sec-"
# a value is judged in windows of this many characters (overlapping): an injection is local (a break-out and its
# tail), so padding in front of it changes nothing (V-06); a window plus the longest context template stays within the
# 1,024 characters database.views judges whole
WINDOW = 960
WINDOW_OVERLAP = 256
# the work of one request: this many characters of gate-passing windows are judged, the values that look most like SQL
# first; a request with more is answered with a soft hit (alarm + challenge) on the field where the budget ran out -
# never a silent pass (V-06), never a block on size alone. Plain text (no quote, comment or SQL word) costs nothing.
BUDGET = 6144
# an apostrophe inside an English contraction (it's, don't, I'm, they'll, we're, I've, I'd) cannot end a literal
# usefully - the letters after it are no SQL - so it does not open the gate by itself (V-18; the classify fallback
# needs real SQL structure too). A value that also has a comment, another quote or a SQL word is judged as it is.
_CONTRACTION = re.compile(r"(?<=[A-Za-z])'(?=(?:s|t|m|d|ll|re|ve)(?![A-Za-z0-9_]))")
# the learned tier counts only on values that speak SQL: at least two of these words (or one and a quote / comment)
_KW = re.compile(r"\b(select|union|sleep|benchmark|waitfor|or|and|xor|from|where|having|exec|declare|cast|convert|char|concat|null|like|rlike|"
                 r"dual|information_schema|case|when|then|substring|ascii|if)\b", re.I)


def decode(v: str) -> str:
    """the value as the application finally sees it"""
    for _ in range(3):
        d = unquote_plus(v) if "%" in v or "+" in v else v
        d = _PCTU.sub(lambda m: chr(int(m.group(1), 16)), d)
        if d == v:
            break
        v = d
    if "&" in v:
        v = html.unescape(v)
    if "\\" in v:
        v = _HEXESC.sub(lambda m: chr(int(m.group(1), 16)), v)
        v = _UESC.sub(lambda m: chr(int(m.group(1), 16)), v)
    return v.replace("\x00", "")


def fields(text: str) -> list[tuple[str, str]]:
    """(field, decoded value) pairs of a canonical request text; a key that looks like a payload is a value too"""
    parts = text.split(" | ")
    out = []
    head = parts[0].split(" ", 1)
    if len(head) == 2 and "?" in head[1]:                    # an un-split query string on the request line
        for kv in head[1].split("?", 1)[1].split("&"):
            k, _, v = kv.partition("=")
            out.append((f"q:{k}", decode(v)))
    if len(head) == 2:
        for seg in head[1].split("?", 1)[0].split("/"):         # path segments (REST ids)
            if seg and _PRE.search(seg):
                out.append(("path", decode(seg)))
    for p in parts[1:]:
        if len(p) < 3 or p[1] != ":" or p.startswith(_SKIP_PREFIX) or p.partition("=")[0] in SKIP_HEADERS:
            continue
        k, _, v = p.partition("=")
        if v:
            out.append((k, decode(v)))
        name = decode(k[2:])
        if _PRE.search(name) and any(c in name for c in " '\"(#"):
            # a payload posted as a key (often with no value, and with '=' inside it): judge the key, and the whole
            # 'key=value' as one value - split at its first '=' a tautology like 5*5=25 would be cut in half
            out.append((k + "#key", name))
            out.append((k + "#part", decode(p[2:])))
    return out


@lru_cache(maxsize=65536)
def judge_value(v: str, use_net: bool = False) -> tuple:
    """-> (contexts that broke, families, DbNet score or None, DbNet-only verdict) for one decoded value (cached:
    user agents, referers and cookies repeat endlessly)"""
    found = []
    sq, dq = "'" in v or "\\" in v, '"' in v or "\\" in v
    for ctx, tpl in CONTEXTS.items():
        if (ctx.startswith("paren") and ")" not in v) or (ctx == "like" and "%" not in v):
            continue
        # a value with no quote (or backslash) of the context's kind stays inside that context's literal: the query is
        # the template's own, nothing to judge there (RAW signatures see the value in the numeric context anyway)
        if (ctx in ("string", "paren_str", "like") and not sq) or (ctx == "dquote" and not dq):
            continue
        # MySQL reads "..." as a string: the double-quote context is judged as a string context with the quotes swapped
        q = CONTEXTS["string"].format(v=v.translate(_SWAP)) if ctx == "dquote" else tpl.format(v=v)
        # the comment-cut fallback reads the string-like contexts as "the first quote is the application's" and does
        # not apply where the value stands outside any literal (database.classify `literal`)
        fam = [f for f in classify(q, "mysql", literal=ctx in _LITERAL) if f in SERIOUS]
        if fam:
            found.append((ctx, tuple(fam)))
    score, learned = None, False
    kw = len(_KW.findall(v))
    if use_net and _NET[0] is not None and (kw >= 2 or (kw >= 1 and re.search(r"['\"`]|--|/\*|#", v))):
        qs = [tpl.format(v=v) for tpl in CONTEXTS.values()]
        s, fam_n = _NET[0].score(qs, ["mysql"] * len(qs))
        score = float(s.max())
        if not found and score > _NET[1] and fam_n[int(s.argmax())] == "injection":
            found.append((list(CONTEXTS)[int(s.argmax())], ("injection",)))
            learned = True
    return tuple(found), score, learned


_NET: list = [None, 4.4]             # DbNet and its per-value threshold (set by use_dbnet)


def use_dbnet(net, threshold: float = 4.4) -> None:
    """enable the learned tier on parameter values (threshold: benign parameter values of another site, FPR 1e-3)"""
    _NET[0], _NET[1] = net, float(threshold)
    judge_value.cache_clear()


def windows(v: str) -> list[str]:
    """the overlapping WINDOW-character pieces of a value (a short value is its own only window)"""
    if len(v) <= WINDOW:
        return [v]
    step = WINDOW - WINDOW_OVERLAP
    return [v[i:i + WINDOW] for i in range(0, len(v) - WINDOW_OVERLAP, step)]


def gated(v: str) -> bool:
    """could this decoded value change a query's structure? (contraction apostrophes do not count)"""
    if "'" in v:
        v = _CONTRACTION.sub("", v)
    return bool(_GATE.search(v))


def analyse(text: str, net=None, net_thr: float | None = None, max_fields: int | None = None) -> list[dict]:
    """A22 on one canonical request text: every field (the 65th like the first) and every window of a long value,
    within BUDGET; a request with more gate-passing text than that, or one parse.model_text says was too large to
    judge (x:oversized), gets a soft hit instead of a silent pass. max_fields: kept for the old signature; ignored."""
    if net is not None and (_NET[0] is not net or (net_thr is not None and _NET[1] != net_thr)):
        use_dbnet(net, net_thr if net_thr is not None else _NET[1])
    hits = []
    cand = [(f, v) for f, v in fields(text) if len(v) >= 2 and gated(v)]
    if len(cand) > 1:                    # the values that look most like SQL first
        cand.sort(key=lambda fv: -len(_STRONG.findall(fv[1][:WINDOW])))
    budget, over = BUDGET, None
    for field, v in cand:
        for w in windows(v):
            if len(w) < 2 or not gated(w):
                continue
            if len(w) > budget:
                over = field
                break
            budget -= len(w)
            found, score, learned = judge_value(w, net is not None)
            if found:
                hits.append({"field": field, "value": w[:200], "context": found[0][0], "contexts": [c for c, _ in found],
                             "families": sorted({f for _, fs in found for f in fs}), "dbnet": score, "learned_only": learned})
                break
        if over is not None:
            break
    if over is not None:
        hits.append({"field": over, "value": "too much SQL-looking text to judge", "context": "budget", "contexts": ["budget"],
                     "families": ["injection"], "dbnet": None, "learned_only": True})
    if " | x:oversized=" in text:
        hits.append({"field": "x:oversized", "value": text[text.index(" | x:oversized=") + 3:][:200], "context": "budget",
                     "contexts": ["budget"], "families": ["injection"], "dbnet": None, "learned_only": True})
    return hits