| """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) |
| _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+") |
| |
| _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) |
| 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 |
|
|
| @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 |
|
|
