Blame for runtime/python/dibaguard/dbnet.py

"""DbNet: the learned tier of database protection - one small network for every query language DibaGuard reads
(SQL dialects, CQL, MongoDB filters / commands, Redis commands, Elasticsearch requests, Cypher).

Input is the query's *structure*, not its text: string literals become 's', numbers 0, hex blobs 0x, identifiers the
model has no business knowing (table / column / key names) become x, and the vocabulary of the languages themselves
(keywords, functions, operators, catalogs, commands) is kept. MySQL version-comments are unwrapped and inline comments
reduced to a marker, the usual obfuscations of injection payloads. So the network judges what a query *does* -
an injected tautology, a stacked statement, a COPY ... TO PROGRAM - whatever the schema, and a scary word inside a
properly quoted value is just 's'.

    structure -> byte n-grams (orders 2-6, hashed, the same hashing as the HTTP tier) -> EmbeddingBag(mean, d=48)
              -> Linear 48->64 -> ReLU -> Linear 64->7  (benign, injection, rce, privilege, destructive, recon, unbounded)

Inference is numpy only (~0.1 ms per query in batches); the weights are int8 with per-column scales.
"""
from __future__ import annotations

import re

import numpy as np

from dibaguard import hashing

CLASSES = ["benign", "injection", "rce", "privilege", "destructive", "recon", "unbounded"]
ORDERS = (2, 3, 4, 5, 6)
KIND = {"postgres": "sql", "mysql": "sql", "mssql": "sql", "oracle": "sql", "cockroach": "sql", "clickhouse": "sql", "sqlite": "sql",
        "cassandra": "cql", "mongodb": "doc", "couchdb": "doc", "redis": "kv", "elasticsearch": "es", "neo4j": "graph"}

VOCAB = set("""
select from where and or not in is null like ilike between exists any all some union intersect except distinct as on join inner left right
full outer cross natural using group by having order asc desc limit offset fetch first next rows row only top with recursive case when then
else end cast convert coalesce nullif count sum avg min max over partition window values insert into update set delete returning default
conflict do nothing merge matched create alter drop truncate rename table schema database keyspace index view sequence function procedure
trigger extension role user grant revoke to public privileges superuser createrole createdb login password replication bypassrls nologin
begin commit rollback savepoint transaction start isolation level read write committed serializable lock share mode for nowait skip locked
explain analyze vacuum reindex cluster copy program stdin stdout csv header delimiter format language plpgsql plpythonu plpython3u plperlu
c sql java javascript returns return declare execute exec prepare deallocate if else elsif loop while do raise notice perform call
true false unknown interval date time timestamp now current_date current_timestamp current_user session_user user version database
substring substr ascii char chr concat concat_ws group_concat string_agg length lower upper trim replace md5 sha1 sha2 hex unhex
sleep pg_sleep benchmark waitfor delay dbms_pipe receive_message dbms_lock randomblob extractvalue updatexml exp floor rand
load_file outfile dumpfile infile local data load lines terminated fields enclosed xp_cmdshell xp_dirtree xp_fileexist sp_oacreate
sp_oamethod sp_configure sp_addsrvrolemember sp_addlogin sp_executesql openrowset opendatasource openquery reconfigure sysadmin
lo_import lo_export lo_create pg_read_file pg_read_binary_file pg_ls_dir pg_write_file pg_stat_file dblink dblink_exec adminpack
sys_exec sys_eval utl_http utl_file utl_tcp utl_inaddr dbms_java dbms_scheduler dbms_xmlquery httpuritype
information_schema tables columns schemata table_name column_name table_schema user_privileges routines processlist
pg_catalog pg_shadow pg_authid pg_user pg_roles pg_settings pg_hba_file_rules pg_tables pg_class pg_namespace pg_attribute pg_proc
pg_stat_activity pg_stat_statements pg_database pg_stat_user_tables pg_locks pg_indexes
mysql sys sql_logins server_principals syslogins sysobjects syscolumns master sysdatabases authentication_string
system_schema system_auth system roles role_permissions keyspaces
dual rownum all_tables all_users dba_users v$version
allow filtering ttl writetime token batch apply using timestamp
db find findone aggregate insertone insertmany updateone updatemany replaceone deleteone deletemany remove drop dropdatabase
countdocuments distinct createindex mapreduce runcommand admincommand listdatabases listcollections getcmdlineopts serverstatus
createuser updateuser grantrolestouser dropuser getusers eval currentop killop
root useradminanydatabase readwriteanydatabase dbowner dbadmin clusteradmin backup restore readanydatabase
$where $function $accumulator $expr $ne $eq $gt $gte $lt $lte $in $nin $regex $exists $or $and $not $nor $set $inc $push $pull
$unset $lookup $match $group $project $sort $limit $skip $unwind $count $addfields $out $merge $size $elemmatch $options $text $search
get set del exists expire ttl incr decr hget hset hgetall hdel lpush rpush lpop rpop lrange sadd smembers srem zadd zrange zrem
mget mset setex getex publish subscribe multi exec discard watch ping auth select scan match keys flushall flushdb config dir
dbfilename save bgsave slaveof replicaof no one module eval evalsha script info client list kill setname acl setuser deluser
debug object shutdown monitor
_search _doc _bulk _count _mapping _settings _cat _cluster _nodes _security _snapshot _delete_by_query _update_by_query _reindex
_all indices health query bool must should filter term terms range match_all size scroll script source painless groovy aliases
call yield apoc dbms listconnections killconnection csv load merge detach node relationship
s url file s3 remote mysql postgresql hdfs
""".split())
_VCOMMENT = re.compile(r"/\*!\d*\s*(.*?)\*/", re.S)          # MySQL executable comments: the payload inside is real SQL
_COMMENT = re.compile(r"/\*.*?\*/", re.S)
_SQ = re.compile(r"'(?:[^'\\]|\\.|'')*'")
_DQ = re.compile(r"\"((?:[^\"\\]|\\.)*)\"(\s*:)?")
_HEX = re.compile(r"\b0x[0-9a-f]+\b")
_MIX = re.compile(r"(?<![\w$])\d[a-z0-9_-]*[a-z][a-z0-9_-]*|\b[0-9a-f]{8,}\b")
_NUM = re.compile(r"(?<![\w$])\d+(?:\.\d+)?(?:e[+-]?\d+)?\b")
_WORD = re.compile(r"\$?[a-z_][a-z0-9_$]*")
_WS = re.compile(r"\s+")
# secrets and the accounts they belong to: what credential dumping reads, whatever the schema calls them
_CRED = re.compile(r"(pass(word|wd)?|pwd|.*password.*|.*passwd.*|.*_hash|hash(ed)?_?\w*|.*token.*|.*secret.*|api_?keys?|credentials?"
                   r"|authentication_string|salted_hash|private_?key)")


def _word(m: re.Match) -> str:
    w = m.group(0)
    return w if w in VOCAB else "cred" if _CRED.fullmatch(w) else "x"


def _dq(m: re.Match) -> str:
    """a JSON key keeps its name when it is part of a language (a command, an operator) or a credential field;
    a JSON value - or a double-quoted SQL identifier / MySQL string - is just 's'"""
    v, key = m.group(1), m.group(2)
    if key:
        name = v if v in VOCAB else "cred" if _CRED.fullmatch(v) else "k"
        return f'"{name}":'
    return f'"{v}"' if v.startswith("$") and v in VOCAB else '"s"'


def structure(query: str, engine: str = "") -> str:
    """the query as its structure (see module doc); prefixed with the engine's language kind"""
    q = _VCOMMENT.sub(r" \1 ", str(query)[:2000])
    q = _COMMENT.sub(" /**/ ", q).lower()
    q = _SQ.sub("'s'", q)
    q = _DQ.sub(_dq, q)
    q = _HEX.sub("0x", q)
    q = _MIX.sub("h", q)                     # ids that mix digits and letters (hex keys, uuids, hashes) are values
    q = _NUM.sub("0", q)
    q = _WORD.sub(_word, q)
    return KIND.get(engine, "sql") + ": " + _WS.sub(" ", q).strip()[:600]


class DbNet:
    def __init__(self, d: dict):
        self.buckets = int(d["buckets"])
        self.emb = d["emb_q"].astype(np.int8)
        self.scale = d["emb_scale"].astype(np.float32)
        self.w1, self.b1 = d["w1"].astype(np.float32), d["b1"].astype(np.float32)
        self.w2, self.b2 = d["w2"].astype(np.float32), d["b2"].astype(np.float32)
        self.thr = float(d["thr_1e4"]) if "thr_1e4" in d else 0.0     # global threshold until a database has its own reference

    @classmethod
    def load(cls, path: str) -> "DbNet":
        with np.load(path) as z:
            return cls({k: z[k] for k in z.files})

    def logits(self, queries: list[str], engines: list[str] | None = None) -> np.ndarray:
        engines = engines or [""] * len(queries)
        s = [structure(q, e) for q, e in zip(queries, engines)]
        idx, doc = hashing.ngram_buckets(s, self.buckets, ORDERS)
        x = hashing.bag_mean(self.emb, idx, doc, len(s), self.scale)
        h = np.maximum(x @ self.w1.T + self.b1, 0.0)
        return h @ self.w2.T + self.b2

    def proba(self, queries: list[str], engines: list[str] | None = None) -> np.ndarray:
        z = self.logits(queries, engines)
        z = z - z.max(1, keepdims=True)
        p = np.exp(z)
        return p / p.sum(1, keepdims=True)

    def score(self, queries: list[str], engines: list[str] | None = None) -> tuple[np.ndarray, np.ndarray]:
        """-> (attack score = log-odds of 'not benign', most likely attack family per query)"""
        p = self.proba(queries, engines)
        pb = np.clip(p[:, 0], 1e-7, 1 - 1e-7)
        fam = np.asarray(CLASSES, dtype=object)[1 + p[:, 1:].argmax(1)]
        return np.log((1 - pb) / pb), fam