yilunjiang/AgentLab
1
0
Fork 0
mirror of https://github.com/ThearchyHelios/AgentLab.git synced 2026-11-15 06:11:58 +00:00
Code Issues Packages Projects Releases Wiki Activity
Files
f5e6e77ea9fd06ba2f3ed2c21a5fbde730c2dccf
AgentLab/backend/tests/test_masking_opaque_sql.py
T

531 lines
34 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""认不全的 SQL 写法整份遮（engine/masking.py 的 _tainted），不再逐个形状补规则（K3 第 3 轮复验留下的几处）：
- (表达式).* 展开：复合类型的列 (p.profile).*、函数交回的记录 (jsonb_each_text(pj)).*、(jsonb_populate_record(…)).*、
子查询里的匿名 ROW (q.r).*。展开成几列、各叫什么 SQL 里看不出，表达式读到遮罩列的，整份结果都遮——最外层、子查询、
CTE、LATERAL 里都一样。以前改写成一个对不上来源的 *，最外层、没写列名表的派生来源里按「照原名带出」认，字段名
（tel、key/value、f1/f2）和遮罩列对不上，号码原样交给模型、落库。(s).*、s 是 FROM 里的来源时照旧按 s.* 认
- 表名、SQL 写在字符串里交给函数执行、整表导出的（query_to_xml、table_to_xml、cursor_to_xml、schema_to_xml、
database_to_xml、dblink，Oracle 的 DBMS_XMLGEN），统计信息里存着的值（pg_stats）：整份遮，解析不了的也一样
- LATERAL 看不见它自己：里面的整行引用 s 不因为 LATERAL 自己的结果列叫 s 就被当成列
- 分析不过来、分析出了意外：整份遮。派生来源对应的那块查询、一项的值取自哪些列在一次分析里记下来，一长串 CTE 一段套
一段地整行引用（250 段，2 万字的上限里写得下）以前要二十几秒，现在零点几秒；超出工作量上限（ANALYSIS_STEPS）的整份遮
这里没有真的 Postgres：结果按它交回的样子造（复合类型、记录展开成几列）。数据都是虚构的：手机号、地址、区域。
"""
from __future__ import annotations
import json
import time
import uuid
import pytest
from httpx import ASGITransport, AsyncClient
from app.core import artifact_store
from app.data.engine import engines
from app.db.base import SessionLocal
from app.db.models import DataSource
from app.engine import masking
from app.main import app
from tests.test_verify_agent import Script, q
from tests.test_verify_agent_masking import ask, clean_spend # noqa: F401 - clean_spend 是自动用上的夹具
MASKED_WORD = "已遮罩"
#: people.profile 是复合类型（tel、addr），people2.pj 是 JSON（tel、addr），sales 是订单：设了遮罩的是 profile、pj、buyer_phone
SCHEMA = {"tables": {
"people": {"qualified": "people", "schema": None, "columns": [{"name": c} for c in ["id", "profile", "address"]]},
"people2": {"qualified": "people2", "schema": None, "columns": [{"name": c} for c in ["id", "pj", "city"]]},
"sales": {"qualified": "sales", "schema": None,
"columns": [{"name": c} for c in ["id", "region", "buyer_phone", "amount"]]},
}}
MASKED = ["profile", "pj", "buyer_phone"]
PHONE = "13900000001"
def _leaks(done: masking.QueryMask) -> bool:
return "1390000000" in json.dumps(done.payload, ensure_ascii=False)
def masked(sql: str, columns: list[str], rows: list[list], kind: str = "postgresql") -> masking.QueryMask:
payload = {"columns": columns, "rows": rows, "sql": sql}
return masking.mask_query_result(payload, MASKED, kind, schema=SCHEMA)
# --------------------------------------------------------------------------
# (表达式).* 展开
# --------------------------------------------------------------------------
#: 三种展开（复合类型的列、函数交回的记录、子查询里的匿名 ROW）× 四个位置（最外层、子查询、CTE、LATERAL）：
#: (SQL, 结果列, 一行)
EXPANSIONS = [
pytest.param("SELECT p.id AS n, (p.profile).* FROM people p", ["n", "tel", "addr"], [1, PHONE, "某路 1 号"],
id="复合类型-最外层"),
pytest.param("SELECT x.* FROM (SELECT p.id AS n, (p.profile).* FROM people p) x", ["n", "tel", "addr"],
[1, PHONE, "某路 1 号"], id="复合类型-子查询"),
pytest.param("WITH x AS (SELECT p.id AS n, (p.profile).* FROM people p) SELECT * FROM x", ["n", "tel", "addr"],
[1, PHONE, "某路 1 号"], id="复合类型-CTE"),
pytest.param("SELECT x.* FROM people p, LATERAL (SELECT p.id AS n, (p.profile).*) x", ["n", "tel", "addr"],
[1, PHONE, "某路 1 号"], id="复合类型-LATERAL"),
pytest.param("SELECT (jsonb_each_text(pj)).* FROM people2", ["key", "value"], ["tel", PHONE],
id="记录函数-最外层"),
pytest.param("SELECT kv.value AS tel FROM (SELECT (jsonb_each_text(pj)).* FROM people2) kv WHERE kv.key = 'tel'",
["tel"], [PHONE], id="记录函数-子查询"),
pytest.param("WITH kv AS (SELECT (jsonb_each_text(pj)).* FROM people2) SELECT value AS 号码 FROM kv "
"WHERE key = 'tel'", ["号码"], [PHONE], id="记录函数-CTE"),
pytest.param("SELECT x.value FROM people2 p, LATERAL (SELECT (jsonb_each_text(p.pj)).*) x", ["value"], [PHONE],
id="记录函数-LATERAL"),
pytest.param("SELECT (q.r).* FROM (SELECT ROW(buyer_phone, region) AS r FROM sales) q", ["f1", "f2"],
[PHONE, "华东"], id="匿名 ROW-最外层"),
pytest.param("SELECT y.* FROM (SELECT (q.r).* FROM (SELECT ROW(buyer_phone, region) AS r FROM sales) q) y",
["f1", "f2"], [PHONE, "华东"], id="匿名 ROW-子查询"),
pytest.param("WITH q AS (SELECT ROW(buyer_phone, region) AS r FROM sales) SELECT (q.r).* FROM q", ["f1", "f2"],
[PHONE, "华东"], id="匿名 ROW-CTE"),
pytest.param("SELECT x.* FROM (SELECT ROW(buyer_phone, region) AS r FROM sales) q, LATERAL (SELECT (q.r).*) x",
["f1", "f2"], [PHONE, "华东"], id="匿名 ROW-LATERAL"),
pytest.param("SELECT p.id AS n, (jsonb_populate_record(NULL::contact, p.pj)).* FROM people2 p",
["n", "tel", "addr"], [1, PHONE, "某路 1 号"], id="按类型展开 JSON"),
]
@pytest.mark.parametrize("sql, columns, row", EXPANSIONS)
def test_an_expression_star_reading_a_masked_column_masks_the_whole_result(sql, columns, row):
"""(表达式).* 读到遮罩列：展开成几列、叫什么说不清，整份结果的每一列都遮（同一份里的编号也遮，多遮）。"""
done = masked(sql, columns, [list(row)])
assert done.columns == columns
assert done.payload["rows"] == [[MASKED_WORD] * len(columns)] and not _leaks(done)
@pytest.mark.parametrize("sql, columns, row", [
pytest.param("SELECT p.id AS n, (p.address).* FROM people p", ["n", "street", "city"], [1, "某路", "某市"],
id="没设遮罩的复合类型"),
pytest.param("SELECT (jsonb_each_text(to_jsonb(c))).* FROM (SELECT id, region FROM sales) c", ["key", "value"],
["region", "华东"], id="记录里没有遮罩列"),
])
def test_an_expression_star_that_reads_no_masked_column_stays_visible(sql, columns, row):
done = masked(sql, columns, [list(row)])
assert done.columns == [] and done.payload["rows"] == [list(row)]
def test_a_parenthesised_source_star_is_still_read_as_its_columns():
"""(s).*、s 是 FROM 里的来源：照旧当 s.* 认，只遮号码那一列，区域、金额照常。"""
done = masked("SELECT (s).* FROM sales s", ["id", "region", "buyer_phone", "amount"], [[1, "华东", PHONE, 1200]])
assert done.columns == ["buyer_phone"] and done.payload["rows"] == [[1, "华东", MASKED_WORD, 1200]]
#: 造出来的 Postgres 数据源：表结构同上，查询交给换掉的 run_query（不连库）
@pytest.fixture
async def composite_shop():
made: list[int] = []
async def make() -> str:
name = f"pg_{uuid.uuid4().hex[:8]}"
async with SessionLocal() as session:
row = DataSource(name=name, kind="postgresql", host="127.0.0.1", port=5432, username="probe",
database="shop", readonly=True, enabled=True,
options={"mask_columns": ", ".join(MASKED)}, schema_cache=SCHEMA)
session.add(row)
await session.commit()
made.append(row.id)
return name
yield make
for source_id in made:
await engines.invalidate(source_id)
@pytest.mark.parametrize("sql, columns, rows", [
("SELECT p.id AS n, (p.profile).* FROM people p", ["n", "tel", "addr"], [[1, PHONE, "某路 1 号"]]),
("SELECT x.* FROM (SELECT p.id AS n, (p.profile).* FROM people p) x", ["n", "tel", "addr"],
[[1, PHONE, "某路 1 号"]]),
("SELECT p.id AS n, (jsonb_populate_record(NULL::contact, p.pj)).* FROM people2 p", ["n", "tel", "addr"],
[[1, PHONE, "某路 1 号"]]),
("SELECT p.id AS n, ((p.profile).*) FROM people p", ["n", "tel", "addr"], [[1, PHONE, "某路 1 号"]]),
], ids=["复合类型", "子查询里的复合类型", "按类型展开 JSON", "外面套了括号"])
async def test_an_expression_star_never_reaches_the_model_or_the_store(monkeypatch, composite_shop, sql, columns,
rows):
"""端到端（核查 agent 的真实查询工具）：交给模型的、落库的、通用工件接口给的都只有「已遮罩」，masked_at_query 是
全部的列，快照记着整份遮（masked_whole），格子引用照实说写法认不全。以前前三条的 masked_at_query 都是空的，套了外括号
的只遮了第一列，号码原样交给模型、落库。"""
from app.data.engine import QueryResult
from app.tools import datasource
async def driver(source, statement, **kwargs):
assert statement == sql
return QueryResult(columns=list(columns), rows=[list(r) for r in rows], row_count=len(rows), truncated=False,
elapsed_ms=1, sql=statement)
monkeypatch.setattr(datasource, "run_query", driver)
name = await composite_shop()
script = Script(monkeypatch, [q(name, sql, "c1")], {"value": {"alias": "Q1", "row": 0, "column": "n"},
"rationale": "x"})
out, _ = await ask(name)
seen = script.tool_results()["c1"]
assert MASKED_WORD in seen and "1390000000" not in seen, seen
snapshot = artifact_store.load(out.queries[0]["artifact"])
assert snapshot["masked_at_query"] == columns and "1390000000" not in json.dumps(snapshot, ensure_ascii=False)
assert snapshot["masked_whole"] is True and "masked_whole" not in seen
assert out.cells[0]["reason"] == "Q1 的写法认不全，整份结果查询时已遮罩，不采信"
async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client:
got = await client.get(f"/api/artifacts/{out.queries[0]['artifact']}")
assert got.status_code == 200 and "1390000000" not in got.text
# --------------------------------------------------------------------------
# LATERAL 看不见它自己
# --------------------------------------------------------------------------
@pytest.mark.parametrize("sql, columns, row", [
pytest.param("SELECT y.* FROM sales t2, LATERAL (SELECT s FROM sales s LIMIT 1) y", ["s"],
[f"(1,华东,{PHONE},1200)"], id="选出整行"),
pytest.param("SELECT y.j FROM sales t2 CROSS JOIN LATERAL (SELECT 1 AS s, row_to_json(s) AS j FROM sales s "
"LIMIT 1) y", ["j"], [json.dumps({"id": 1, "buyer_phone": PHONE})], id="结果列也叫 s"),
pytest.param("SELECT y.j FROM sales t2, LATERAL (SELECT s.id AS s, row_to_json(s) AS j FROM sales s LIMIT 1) y",
["j"], [json.dumps({"id": 1, "buyer_phone": PHONE})], id="编号起名叫 s"),
])
def test_a_whole_row_inside_a_lateral_is_not_taken_for_its_own_column(sql, columns, row):
"""LATERAL 里的 s 是它自己 FROM 里的整行：LATERAL 的结果列碰巧叫 s，也不把它当成列（Postgres 里 LATERAL 看不见
自己）。以前当成了列，整行原样交了出去。"""
done = masked(sql, columns, [list(row)])
assert done.columns == columns and not _leaks(done)
def test_a_lateral_still_sees_the_sources_before_it():
"""排在 LATERAL 前面的来源照样看得见：前面的表有一列叫 s 时，s 是那一列（没设遮罩），不是整行。"""
schema = {"tables": {**SCHEMA["tables"], "codes": {"qualified": "codes", "schema": None,
"columns": [{"name": "id"}, {"name": "s"}]}}}
payload = {"columns": ["v"], "rows": [["A1"]],
"sql": "SELECT y.v FROM codes c, LATERAL (SELECT s AS v FROM sales x WHERE x.id = c.id) y"}
done = masking.mask_query_result(payload, MASKED, "postgresql", schema=schema)
assert done.columns == [] and done.payload["rows"] == [["A1"]]
# --------------------------------------------------------------------------
# 表名、SQL 写在字符串里交给函数，统计信息里的值
# --------------------------------------------------------------------------
@pytest.mark.parametrize("sql", [
"SELECT query_to_xml('select * from sales', true, false, '') AS x",
"SELECT table_to_xml('sales', true, false, '') AS x",
"SELECT query_to_xml_and_xmlschema('select buyer_phone from sales', true, false, '') AS x",
"SELECT cursor_to_xml('c1', 10, false, true, '') AS x",
"SELECT schema_to_xml('public', true, false, '') AS x",
"SELECT database_to_xml(true, false, '') AS x",
"SELECT x FROM dblink('dbname=shop', 'select buyer_phone from sales') AS t(x text)",
"SELECT DBMS_XMLGEN.GETXML('select * from sales') AS x FROM dual",
"SELECT array_to_string(most_common_vals, ',') AS x FROM pg_stats WHERE tablename = 'sales'",
])
def test_tables_and_sql_handed_to_functions_as_text_mask_the_whole_result(sql):
"""表名、SQL 写在字符串里交给函数执行、整表导出，统计信息里存着的值：读的是哪一列认不出，整份遮。以前 masked=[]，
号码原样。"""
done = masked(sql, ["x"], [[f"<row><buyer_phone>{PHONE}</buyer_phone></row>"]], kind="oracle"
if "DBMS" in sql else "postgresql")
assert done.columns == ["x"] and not _leaks(done)
@pytest.mark.parametrize("sql", [
"SELECT query AS x FROM pg_stat_monitor",
"SELECT q.query AS x FROM pg_stat_monitor_internal(true) q",
"SELECT quals::text AS x FROM pg_qualstats",
"SELECT constvalue AS x FROM pg_qualstats()",
"SELECT plan AS x FROM pg_store_plans",
])
def test_extensions_that_keep_other_sessions_sql_mask_the_whole_result(sql):
"""Postgres 扩展里记着别的会话执行过的语句、条件里的常量的视图和函数（pg_stat_monitor 的 query、pg_qualstats 的
constvalue、pg_store_plans 的计划）：交回的就是带号码的原文，和 pg_stat_statements 同一类，整份遮；可验证报告用的是同
一份名单，也只给哈希。"""
from app.bundle import build
done = masked(sql, ["x"], [[f"SELECT id FROM sales WHERE buyer_phone = '{PHONE}'"]])
assert done.columns == ["x"] and not _leaks(done)
assert masking._INDIRECT.search(sql) and build._INDIRECT_READS.search(sql)
def test_a_table_handed_to_a_function_masks_the_whole_result_even_when_the_sql_cannot_be_read():
sql = "SELECT table_to_xml('sales', true, false, '') AS x WHERE ("
hidden, _ = masking._derived_hidden({"sql": sql, "columns": ["x"], "rows": []}, {"buyer_phone"}, "postgresql")
assert "x" in hidden
# --------------------------------------------------------------------------
# 分析不过来、分析出了意外：整份遮；一长串 CTE 不再卡住
# --------------------------------------------------------------------------
def _chain(k: int, tail: str) -> str:
"""一长串 CTE，一段套一段地整行引用、展开上一段（复验时 250 段要二十几秒）。"""
parts = ["c0 AS (SELECT * FROM sales)"]
parts += [f"c{i} AS (SELECT row_to_json(c{i - 1}) AS j, c{i - 1}.* FROM c{i - 1})" for i in range(1, k)]
return "WITH " + ", ".join(parts) + f" {tail.format(last=f'c{k - 1}')}"
def _plain_chain(k: int) -> str:
parts = ["c0 AS (SELECT id, region, amount FROM sales)"]
parts += [f"c{i} AS (SELECT id, region, amount + {i} AS amount FROM c{i - 1})" for i in range(1, k)]
return "WITH " + ", ".join(parts) + f" SELECT region, SUM(amount) AS total FROM c{k - 1} GROUP BY region"
def test_a_long_chain_of_whole_rows_is_read_quickly_and_completely():
"""250 段一段套一段地整行引用：零点几秒认完（以前二十几秒），不靠工作量上限兜底——只数个数的结果照常不遮。"""
columns = masking._schema_columns(SCHEMA)
sql = _chain(250, "SELECT COUNT(*) AS n FROM {last}")
assert len(sql) < 20_000
started = time.perf_counter()
hidden, _ = masking._derived_hidden({"sql": sql, "columns": ["n"], "rows": []}, {"buyer_phone"}, "postgresql",
columns)
assert time.perf_counter() - started < 5
assert "n" not in hidden
hidden, _ = masking._derived_hidden({"sql": _chain(250, "SELECT j FROM {last}"), "columns": ["j"], "rows": []},
{"buyer_phone"}, "postgresql", columns)
assert "j" in hidden
def test_an_analysis_over_the_work_limit_masks_the_whole_result(monkeypatch):
"""超出工作量上限：不再分析，整份遮（结果本来没碰遮罩列也一样）。没超的照常。"""
sql = _plain_chain(40)
snapshot = {"sql": sql, "columns": ["region", "total"], "rows": []}
assert masking._derived_hidden(snapshot, {"buyer_phone"}, "postgresql")[0] == {"buyer_phone"}
monkeypatch.setattr(masking, "ANALYSIS_STEPS", 50)
assert masking._derived_hidden(snapshot, {"buyer_phone"}, "postgresql")[0] == {"buyer_phone", "region", "total"}
payload = {"sql": sql, "columns": ["region", "total"], "rows": [["华东", 1200]]}
done = masking.mask_query_result(payload, ["buyer_phone"], "postgresql")
assert done.columns == ["region", "total"] and done.payload["rows"] == [[MASKED_WORD, MASKED_WORD]]
def test_a_failing_analysis_masks_the_whole_result_instead_of_reading_the_text(monkeypatch):
"""解析之后的分析出了意外：整份遮。以前按「解析不了」只看文字，没写遮罩列、没有 * 的整行引用（SELECT s）就放过去了。"""
def broken(*args, **kwargs):
raise RuntimeError("分析里的意外")
monkeypatch.setattr(masking, "_tainted_in", broken)
snapshot = {"sql": "SELECT s AS r FROM sales s", "columns": ["r"], "rows": []}
assert "r" in masking._derived_hidden(snapshot, {"buyer_phone"}, "postgresql")[0]
# --------------------------------------------------------------------------
# (表达式).* 套了外括号、起了别名
# --------------------------------------------------------------------------
#: Postgres 实测：((p.card).*)、(((p.card)).*)、(((p.card).*))、(p.card).* AS x 都照样展开成好几列（别名不起作用）
WRAPPED = [
pytest.param("SELECT ((p.profile).*) FROM people p", ["tel", "addr"], [PHONE, "某路 1 号"], id="复合类型-外括号"),
pytest.param("SELECT (((p.profile)).*) FROM people p", ["tel", "addr"], [PHONE, "某路 1 号"],
id="复合类型-两层外括号"),
pytest.param("SELECT (((p.profile).*)) FROM people p", ["tel", "addr"], [PHONE, "某路 1 号"],
id="复合类型-套两层"),
pytest.param("SELECT (p.profile).* AS x FROM people p", ["tel", "addr"], [PHONE, "某路 1 号"], id="复合类型-别名"),
pytest.param("SELECT ((jsonb_each_text(pj)).*) FROM people2", ["key", "value"], ["tel", PHONE],
id="记录函数-外括号"),
pytest.param("SELECT ((q.r).*) FROM (SELECT ROW(region, buyer_phone) AS r FROM sales) q", ["f1", "f2"],
["华东", PHONE], id="匿名 ROW-外括号"),
pytest.param("SELECT ((jsonb_each_text(pj)).*) FROM people2 UNION ALL SELECT 'a', 'b'", ["key", "value"],
["tel", PHONE], id="UNION 第一支"),
pytest.param("SELECT 'a' AS k, 'b' AS v UNION ALL SELECT ((jsonb_each_text(pj)).*) FROM people2", ["k", "v"],
["tel", PHONE], id="UNION 后一支"),
pytest.param("SELECT 'a' AS k, 'b' AS v UNION ALL SELECT ((p.profile).*) FROM people p", ["k", "v"],
[PHONE, "某路 1 号"], id="UNION 后一支-复合类型"),
pytest.param("SELECT x.* FROM (SELECT ((q.r).*) FROM (SELECT ROW(buyer_phone, region) AS r FROM sales) q) x",
["f1", "f2"], [PHONE, "华东"], id="子查询里-外括号"),
# 和 * 写在同一支里：选出项和结果的列数本来就对不上（* 展开成几列），只能靠认出外括号里的 (表达式).*
pytest.param("SELECT p.*, ((p.profile).*) FROM people p", ["id", "profile", "address", "tel", "addr"],
[1, f"({PHONE},某路 1 号)", "某路 1 号", PHONE, "某路 1 号"], id="和 * 写在一起-外括号"),
pytest.param("SELECT p.*, (p.profile).* AS x FROM people p", ["id", "profile", "address", "tel", "addr"],
[1, f"({PHONE},某路 1 号)", "某路 1 号", PHONE, "某路 1 号"], id="和 * 写在一起-别名"),
]
@pytest.mark.parametrize("sql, columns, row", WRAPPED)
def test_a_wrapped_expression_star_reading_a_masked_column_masks_the_whole_result(sql, columns, row):
"""(表达式).* 外面套了括号、起了别名，Postgres 照样展开成好几列：整份遮。以前只认不带外括号的写法，((p.card).*) 按一项
数，只遮了第一列，号码原样交给模型、落库。"""
done = masked(sql, columns, [list(row)])
assert done.whole and done.columns == columns
assert done.payload["rows"] == [[MASKED_WORD] * len(columns)] and not _leaks(done)
def test_a_wrapped_expression_star_with_a_numeric_phone_is_masked_whole():
"""号码存成整数（bigint）、前面的查询已经认出它：数值格不按值换，(表达式).* 认不出就全靠整份遮挡住。"""
known = masking._Secrets()
known.add({"columns": ["p"], "rows": [[13900000001]]}, {"p"})
payload = {"columns": ["f1", "f2"], "rows": [[1200, 13900000001]],
"sql": "SELECT ((q.r).*) FROM (SELECT ROW(amount, buyer_phone) AS r FROM sales) q"}
done = masking.mask_query_result(payload, MASKED, "postgresql", known, schema=SCHEMA)
assert done.payload["rows"] == [[MASKED_WORD, MASKED_WORD]]
@pytest.mark.parametrize("sql", ["SELECT ((s).*) FROM sales s", "SELECT (s.*) AS x FROM sales s"])
def test_a_wrapped_source_star_is_still_read_as_its_columns(sql):
"""((s).*)、(s.*) AS x、s 是 FROM 里的来源：照旧当 s.* 认，只遮号码那一列。"""
done = masked(sql, ["id", "region", "buyer_phone", "amount"], [[1, "华东", PHONE, 1200]])
assert not done.whole and done.columns == ["buyer_phone"]
assert done.payload["rows"] == [[1, "华东", MASKED_WORD, 1200]]
# --------------------------------------------------------------------------
# 和数据源里的表同名的 CTE
# --------------------------------------------------------------------------
@pytest.mark.parametrize("sql, columns, row", [
pytest.param("WITH sales AS (SELECT * FROM sales WHERE region = '华东') SELECT row_to_json(sales) AS j FROM sales",
["j"], [json.dumps({"id": 1, "buyer_phone": PHONE})], id="整行引用"),
pytest.param("WITH sales AS (SELECT * FROM sales WHERE region = '华东') SELECT s AS r FROM sales s", ["r"],
[f"(1,华东,{PHONE},1200)"], id="选出整行"),
pytest.param("WITH orders AS (SELECT * FROM sales), sales AS (SELECT * FROM sales WHERE amount > 0) "
"SELECT to_jsonb(x) AS j FROM sales x", ["j"], [json.dumps({"buyer_phone": PHONE})],
id="排在后面的同名段"),
pytest.param("WITH sales AS (SELECT * FROM sales) SELECT * FROM sales t(a, b, c, d)", ["a", "b", "c", "d"],
[1, "华东", PHONE, 1200], id="列名表"),
])
def test_a_cte_named_like_a_table_reads_the_table(sql, columns, row):
"""WITH sales AS (SELECT * FROM sales WHERE …) 里的 sales 是数据源里的表，不是这一段自己（不带 RECURSIVE 的 WITH 里，
一段只看得见排在它前面的段）。以前认成它自己，* 展开绕回自己什么列都不带出，号码原样交了出去。"""
done = masked(sql, columns, [list(row)])
assert done.columns and not _leaks(done)
assert done.payload["rows"] == [[MASKED_WORD if c in done.columns else v for c, v in zip(columns, row)]]
def test_a_cte_named_like_a_table_masks_like_any_other_name():
"""对照：同一段 CTE 改叫 t，同样遮；列名表按表结构对上位置，只遮号码那一列。"""
for sql in ("WITH t AS (SELECT * FROM sales WHERE region = '华东') SELECT row_to_json(t) AS j FROM t",
"WITH sales AS (SELECT * FROM sales WHERE region = '华东') SELECT row_to_json(sales) AS j FROM sales"):
assert masked(sql, ["j"], [[json.dumps({"buyer_phone": PHONE})]]).columns == ["j"]
done = masked("WITH sales AS (SELECT * FROM sales) SELECT * FROM sales t(a, b, c, d)", ["a", "b", "c", "d"],
[[1, "华东", PHONE, 1200]])
assert done.columns == ["c"]
def test_a_recursive_cte_still_sees_itself():
"""带 RECURSIVE 的 WITH 里递归的那一支指向它自己：取自号码的那一列照遮，只数步数的照常。"""
body = ("WITH RECURSIVE t(n, p) AS (SELECT 1, buyer_phone FROM sales UNION ALL SELECT n + 1, p FROM t "
"WHERE n < 3) ")
assert masked(body + "SELECT n FROM t", ["n"], [[1]]).columns == []
assert masked(body + "SELECT p FROM t", ["p"], [[PHONE]]).columns == ["p"]
# --------------------------------------------------------------------------
# 语法树和实际执行的对不上、存着值的视图：和可验证报告共用一份名单
# --------------------------------------------------------------------------
@pytest.mark.parametrize("sql, kind, columns", [
pytest.param("SELECT id /*!, buyer_phone AS x */ FROM sales", "mysql", ["id", "x"], id="MySQL 会执行的注释"),
pytest.param("SELECT /*! buyer_phone AS */ region FROM sales", "mysql", ["region"],
id="MySQL 会执行的注释改了列的来历"),
pytest.param('SELECT U&"\\0062uyer_phone" AS x FROM sales', "postgresql", ["x"], id="U&转义的名字"),
pytest.param("SELECT id, U&\"buyer!005fphone\" UESCAPE '!' AS p FROM sales", "postgresql", ["id", "p"],
id="U&带 UESCAPE"),
pytest.param("SELECT query FROM pg_stat_get_activity(NULL)", "postgresql", ["query"], id="pg_stat_get_activity"),
pytest.param("SELECT data FROM pg_logical_slot_peek_changes('s', NULL, NULL)", "postgresql", ["data"],
id="逻辑解码"),
pytest.param("SELECT HISTOGRAM AS h FROM information_schema.COLUMN_STATISTICS WHERE TABLE_NAME = 'sales'", "mysql",
["h"], id="MySQL 直方图"),
pytest.param("SELECT sql_text AS v FROM performance_schema.events_statements_history_long", "mysql", ["v"],
id="MySQL 语句历史"),
pytest.param("SELECT info AS v FROM information_schema.PROCESSLIST", "mysql", ["v"], id="MySQL 会话列表"),
pytest.param("SELECT argument AS v FROM mysql.general_log", "mysql", ["v"], id="MySQL 通用日志"),
pytest.param("SELECT min_value AS v FROM mysql.column_stats", "mariadb", ["v"], id="MariaDB 列统计"),
pytest.param("SELECT sql_fulltext AS v FROM v$sql", "oracle", ["v"], id="Oracle V$SQL"),
pytest.param("SELECT value_string AS v FROM v$sql_bind_capture", "oracle", ["v"], id="Oracle 绑定变量"),
pytest.param("SELECT endpoint_actual_value AS v FROM all_tab_histograms WHERE table_name = 'SALES'", "oracle",
["v"], id="Oracle 直方图"),
pytest.param("SELECT * FROM all_tab_columns WHERE table_name = 'SALES'", "oracle", ["column_name", "low_value"],
id="Oracle 列的最小最大值"),
pytest.param("SELECT query_to_xml/**/('select * from sales', true, false, '') AS x", "postgresql", ["x"],
id="函数名后夹块注释"),
pytest.param("SELECT query_to_xml --c\n('select * from sales', true, false, '') AS x", "postgresql", ["x"],
id="函数名后夹行注释"),
pytest.param("SELECT id FROM sales", "mysql", ["id", "x"], id="选出项和结果列数对不上"),
])
def test_writings_the_analysis_cannot_follow_mask_the_whole_result(sql, kind, columns):
"""语法树和实际执行的对不上（MySQL 会执行的注释、U&"…"、选出项和结果的列数对不上），统计信息、会话记录里存着值的视图和
函数，函数名和括号之间夹着注释：整份遮。以前工具边界的名单比可验证报告的少一大截，这些写法号码原样交给模型、落库。"""
row = [PHONE] * len(columns)
done = masked(sql, columns, [row], kind=kind)
assert done.whole and done.payload["rows"] == [[MASKED_WORD] * len(columns)]
@pytest.mark.parametrize("sql, columns, row", [
pytest.param("WITH query(r, n) AS (SELECT region, COUNT(*) FROM sales GROUP BY region) SELECT r AS region, n "
"FROM query", ["region", "n"], ["华东", 2], id="叫 query 的 CTE"),
pytest.param("SELECT region, 'Query (x)' AS q FROM sales", ["region", "q"], ["华东", "Query (x)"],
id="字符串里的 Query ("),
pytest.param("SELECT column_name FROM information_schema.columns WHERE table_name = 'sales'", ["column_name"],
["buyer_phone"], id="information_schema 的列清单"),
])
def test_names_of_unsupported_databases_and_plain_schema_views_are_not_masked(sql, columns, row):
"""不支持的库（DuckDB 的 query()、SQL Server 的 OPENQUERY）的函数不再收：名叫 query 的 CTE、字符串里的「Query (x)」不再
整份遮；information_schema 里只有表结构的视图照常（核查 agent 查列名、类型常用）。"""
done = masked(sql, columns, [list(row)])
assert not done.whole and done.payload["rows"] == [list(row)]
def test_the_tool_boundary_and_the_bundle_share_one_list():
"""工具边界和可验证报告共用一份名单：报告那边另把系统 schema 整个挡（只决定附不附快照）。"""
from app.bundle import build
assert build._OPAQUE_SQL is masking.OPAQUE_SQL
for sql in ("SELECT HISTOGRAM FROM information_schema.COLUMN_STATISTICS", "SELECT * FROM v$sql",
"SELECT data FROM pg_logical_slot_peek_changes('s', NULL, NULL)"):
assert build._INDIRECT_READS.search(sql) and masking._INDIRECT.search(sql)
plain = "SELECT column_name FROM information_schema.columns"
assert build._INDIRECT_READS.search(plain) and not masking._INDIRECT.search(plain)
@pytest.mark.parametrize("sql", [
"SELECT row_to_json(s) AS j FROM sales s ORDER BY s.id USING >",
"SELECT s AS j FROM sales s WHERE s.region IS NOT NFC NORMALIZED OR true",
"SELECT to_jsonb(s) - 'id' AS j FROM sales s ORDER BY s.id USING <",
])
def test_postgres_sql_that_cannot_be_read_masks_the_whole_result(sql):
"""Postgres 合法、解析器解析不了的写法：整行引用不写列名、也没有 *，只看文字认不出，整份遮。以前 masked=[]，整行连同
号码交给模型、落库。"""
done = masked(sql, ["j"], [[json.dumps({"id": 1, "buyer_phone": PHONE})]])
assert done.whole and not _leaks(done)
# --------------------------------------------------------------------------
# 整份遮的结果不把每一格都收成遮罩值
# --------------------------------------------------------------------------
def test_a_whole_masked_result_does_not_turn_every_cell_into_a_masked_value():
"""统计视图整份遮之后，里面的列名、区域名不收成遮罩值：后一条查询的区域照常，封存的 SQL 照原样。以前每一格都收，
区域名、列名在同一次核查里处处被换掉，封存的 SQL 成了「WHERE 已遮罩 = '华东'」。"""
from app.engine.inquiry import _scrub_sql
stats = masked("SELECT attname, n_distinct FROM pg_stats WHERE tablename = 'sales'", ["attname", "n_distinct"],
[["amount", 5], ["buyer_phone", -1], ["region", 3]])
assert stats.whole and not stats.secrets.strong
payload = {"columns": ["id", "region", "amount"], "rows": [[1, "华东", 1200]],
"sql": "SELECT id, region, amount FROM sales ORDER BY id"}
later = masking.mask_query_result(payload, MASKED, "postgresql", stats.secrets, schema=SCHEMA)
assert later.payload["rows"] == [[1, "华东", 1200]]
sql = "SELECT region, COUNT(*) AS n FROM sales WHERE region = '华东' GROUP BY region"
assert _scrub_sql(sql, stats.secrets, {"buyer_phone"}, "postgresql") == sql
def test_a_whole_masked_result_still_collects_what_surely_came_from_a_masked_column():
"""整份遮时，认得出取自遮罩列的那几列照收（号码改了名字选出来、又夹着认不全的写法），别的列不收。"""
sql = "SELECT buyer_phone AS 手机号, region FROM sales UNION ALL SELECT query_to_xml('x', true, false, ''), 'y'"
done = masked(sql, ["手机号", "region"], [[PHONE, "华东"]])
assert done.whole and done.secrets.hit(PHONE) and not done.secrets.hit("华东")
def test_report_snapshots_that_were_masked_whole_give_only_their_sure_values():
"""报告自己的查询快照（报告取数时没遮）收遮罩值时同一条规则：写法认不全的只收确定取自遮罩列的那几列，区域名、列名不收
（审批质询、反方审稿把它们交给查询工具、拿来换回答和理由，收了就处处被换掉）。"""
from app.engine.inquiry import _add_based
secrets = masking._Secrets()
stats = {"source": "pg", "sql": "SELECT attname, n_distinct FROM pg_stats WHERE tablename = 'sales'",
"columns": ["attname", "n_distinct"], "rows": [["region", 3], ["buyer_phone", -1]]}
mixed = {"source": "pg", "columns": ["手机号", "region"], "rows": [[PHONE, "华东"]],
"sql": "SELECT buyer_phone AS 手机号, region FROM sales UNION ALL SELECT query_to_xml('x', true, false, ''), "
"'y'"}
_add_based(secrets, [(stats, set(MASKED)), (mixed, set(MASKED))], {"pg": "postgresql"})
assert secrets.hit(PHONE) and not secrets.hit("华东") and not secrets.hit("region")
Reference in New Issue View Git Blame Copy Permalink