zhanghongyu_0626
0fb7d34d7a
- Enhanced the `AnalystAgent` class to include an `_audit_terminal` method for logging query denials, clarifications, and errors, ensuring compliance and traceability. - Updated error handling paths to call the new audit method, capturing relevant details such as question, user authentication, and SQL context. - Introduced new validation checks in `sql_guard.py` to enforce ownership filters for sensitive queries, improving security measures. - Added unit tests to verify the correct logging behavior and ownership filter enforcement, ensuring robust functionality. This update significantly strengthens the auditing capabilities of the analyst agent, enhancing security and compliance in query handling.
97 lines
3.8 KiB
Python
97 lines
3.8 KiB
Python
"""sql_guard 单元测试（Wave 6）。"""
|
|
import unittest
|
|
|
|
from app.service.sql_guard import SqlGuardError, extract_tables, inject_ownership, validate
|
|
|
|
|
|
class TestSqlGuard(unittest.TestCase):
|
|
def test_select_allowed(self):
|
|
r = validate("SELECT risk_code, COUNT(*) AS c FROM core_customer GROUP BY risk_code", "full")
|
|
self.assertTrue(r.allowed)
|
|
|
|
def test_customer_self_in_scope(self):
|
|
r = validate(
|
|
"SELECT COUNT(*) FROM core_holding WHERE customer_id='CUST-9527'",
|
|
"self",
|
|
["CUST-9527"],
|
|
)
|
|
self.assertTrue(r.allowed)
|
|
|
|
def test_customer_self_out_of_scope(self):
|
|
with self.assertRaises(SqlGuardError) as cm:
|
|
validate(
|
|
"SELECT * FROM core_holding WHERE customer_id='CUST-1001'",
|
|
"self",
|
|
["CUST-9527"],
|
|
)
|
|
self.assertEqual(cm.exception.error_code, "AUTH_403_NOT_OWNER")
|
|
|
|
def test_insert_rejected(self):
|
|
with self.assertRaises(SqlGuardError) as cm:
|
|
validate("INSERT INTO core_customer VALUES (1)", "full")
|
|
self.assertEqual(cm.exception.error_code, "SQL_NOT_SELECT")
|
|
|
|
def test_advisor_out_of_scope_rejected(self):
|
|
with self.assertRaises(SqlGuardError) as cm:
|
|
validate(
|
|
"SELECT * FROM core_holding WHERE customer_id = 'CUST-1004'",
|
|
"assigned",
|
|
["CUST-1001"],
|
|
)
|
|
self.assertEqual(cm.exception.error_code, "AUTH_403_NOT_ASSIGNED")
|
|
|
|
def test_ops_aggregate_allowed(self):
|
|
r = validate("SELECT COUNT(DISTINCT customer_id) AS cnt FROM core_holding", "aggregate")
|
|
self.assertTrue(r.allowed)
|
|
|
|
def test_created_at_in_select_allowed(self):
|
|
"""Q17：列名 created_at 不应触发 create 关键字误杀。"""
|
|
r = validate(
|
|
"SELECT alert_id, created_at FROM jinrong_agent.risk_alert WHERE status='pending_review'",
|
|
"full",
|
|
)
|
|
self.assertTrue(r.allowed)
|
|
|
|
def test_gap_a_select_customer_id_column_without_where_denied_self(self):
|
|
with self.assertRaises(SqlGuardError) as cm:
|
|
validate(
|
|
"SELECT customer_id, product_id, market_value FROM core_holding",
|
|
"self",
|
|
["CUST-9527"],
|
|
)
|
|
self.assertEqual(cm.exception.error_code, "AUTH_403_SCOPE")
|
|
|
|
def test_gap_a_select_customer_id_column_without_where_denied_assigned(self):
|
|
with self.assertRaises(SqlGuardError) as cm:
|
|
validate(
|
|
"SELECT customer_id, product_id, market_value FROM core_holding",
|
|
"assigned",
|
|
["CUST-9527"],
|
|
)
|
|
self.assertEqual(cm.exception.error_code, "AUTH_403_SCOPE")
|
|
|
|
def test_gap_b_ops_risk_alert_denied(self):
|
|
with self.assertRaises(SqlGuardError) as cm:
|
|
validate("SELECT * FROM jinrong_agent.risk_alert", "aggregate")
|
|
self.assertEqual(cm.exception.error_code, "AUTH_403_SCOPE")
|
|
|
|
def test_gap_b_advisor_risk_alert_without_filter_denied(self):
|
|
with self.assertRaises(SqlGuardError) as cm:
|
|
validate("SELECT * FROM jinrong_agent.risk_alert", "assigned", ["CUST-9527"])
|
|
self.assertEqual(cm.exception.error_code, "AUTH_403_SCOPE")
|
|
|
|
def test_create_table_still_rejected(self):
|
|
with self.assertRaises(SqlGuardError) as cm:
|
|
validate("CREATE TABLE evil (id INT)", "full")
|
|
self.assertIn(cm.exception.error_code, ("SQL_FORBIDDEN", "SQL_NOT_SELECT"))
|
|
|
|
def test_inject_ownership(self):
|
|
out = inject_ownership("SELECT * FROM core_holding", ["CUST-1", "CUST-2"])
|
|
self.assertIn("CUST-1", out)
|
|
|
|
def test_extract_tables(self):
|
|
self.assertEqual(
|
|
extract_tables("SELECT * FROM core_customer c JOIN core_holding h ON h.customer_id=c.customer_id"),
|
|
["core_customer", "core_holding"],
|
|
)
|