"""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"(? 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