42 lines
1.2 KiB
Python
42 lines
1.2 KiB
Python
|
"""NL2SQL 查询结果脱敏。"""
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import hashlib
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
|
||
|
|
def _mask_partial(value: Any) -> Any:
|
||
|
|
if value is None:
|
||
|
|
return None
|
||
|
|
text = str(value)
|
||
|
|
if len(text) == 1:
|
||
|
|
return "*"
|
||
|
|
if len(text) <= 4:
|
||
|
|
return f"{text[0]}{'*' * (len(text) - 2)}{text[-1]}"
|
||
|
|
return f"{text[:3]}{'*' * max(1, len(text) - 5)}{text[-2:]}"
|
||
|
|
|
||
|
|
|
||
|
|
def _mask_value(value: Any, mask_type: str) -> Any:
|
||
|
|
if value is None:
|
||
|
|
return None
|
||
|
|
if mask_type == "hash":
|
||
|
|
return hashlib.sha256(str(value).encode("utf-8")).hexdigest()
|
||
|
|
return _mask_partial(value)
|
||
|
|
|
||
|
|
|
||
|
|
def mask_rows(
|
||
|
|
rows: list[dict[str, Any]],
|
||
|
|
masks: dict[tuple[str, str], str],
|
||
|
|
*,
|
||
|
|
table_name: str | None = None,
|
||
|
|
) -> list[dict[str, Any]]:
|
||
|
|
"""复制并脱敏结果行，不修改数据库查询返回的原始对象。"""
|
||
|
|
output = [dict(row) for row in rows]
|
||
|
|
for (table, column), mask_type in masks.items():
|
||
|
|
if table_name is not None and table != table_name:
|
||
|
|
continue
|
||
|
|
for row in output:
|
||
|
|
if column in row:
|
||
|
|
row[column] = _mask_value(row[column], mask_type)
|
||
|
|
return output
|