Skip to content

Commit 0da3813

Browse files
Bornunique911Bornunique911
committed
Add specialized OWASP cheat sheet map analysis behavior
1 parent 646cf9c commit 0da3813

2 files changed

Lines changed: 425 additions & 1 deletion

File tree

‎application/tests/web_main_test.py‎

Lines changed: 256 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,32 @@ def get_status(self):
3434
return rq.job.JobStatus.STARTED
3535

3636

37+
class MockFinishedJobResult:
38+
class Type:
39+
SUCCESSFUL = "successful"
40+
FAILED = "failed"
41+
42+
type = Type.SUCCESSFUL
43+
44+
@property
45+
def return_value(self):
46+
return [
47+
[
48+
"OWASP Top 10 for LLM and Gen AI Apps 2025",
49+
"OWASP Cheat Sheets",
50+
],
51+
{"status": "done"},
52+
]
53+
54+
55+
class MockFinishedJob:
56+
def get_status(self):
57+
return rq.job.JobStatus.FINISHED
58+
59+
def latest_result(self):
60+
return MockFinishedJobResult()
61+
62+
3763
class TestMain(unittest.TestCase):
3864
def tearDown(self) -> None:
3965
sqla.session.remove()
@@ -878,6 +904,180 @@ def test_gap_analysis_returns_direct_cre_overlap_when_backends_fail(
878904
self.assertIn(compare.id, payload["result"][base.id]["paths"])
879905
db_mock.return_value.add_gap_analysis_result.assert_called_once()
880906

907+
@patch.object(web_main.cre_main, "fetch_upstream_json")
908+
@patch.object(web_main.gap_analysis, "schedule")
909+
@patch.object(db, "Node_collection")
910+
def test_gap_analysis_filters_generic_cheatsheets_for_llm_top10(
911+
self, db_mock, schedule_mock, upstream_fetch_mock
912+
) -> None:
913+
broad_cre = defs.CRE(id="064-808", name="Output encoding", description="")
914+
llm_specific_cre = defs.CRE(
915+
id="161-451", name="Prompt boundary protection", description=""
916+
)
917+
918+
llm = defs.Standard(
919+
name="OWASP Top 10 for LLM and Gen AI Apps 2025",
920+
sectionID="LLM05",
921+
section="Improper Output Handling",
922+
)
923+
llm.add_link(
924+
defs.Link(ltype=defs.LinkTypes.LinkedTo, document=broad_cre.shallow_copy())
925+
)
926+
llm.add_link(
927+
defs.Link(
928+
ltype=defs.LinkTypes.LinkedTo, document=llm_specific_cre.shallow_copy()
929+
)
930+
)
931+
932+
generic_cheatsheet = defs.Standard(
933+
name="OWASP Cheat Sheets",
934+
section="Cross Site Scripting Prevention Cheat Sheet",
935+
)
936+
generic_cheatsheet.add_link(
937+
defs.Link(ltype=defs.LinkTypes.LinkedTo, document=broad_cre.shallow_copy())
938+
)
939+
940+
llm_cheatsheet = defs.Standard(
941+
name="OWASP Cheat Sheets",
942+
section="LLM Prompt Injection Prevention Cheat Sheet",
943+
)
944+
llm_cheatsheet.add_link(
945+
defs.Link(
946+
ltype=defs.LinkTypes.LinkedTo, document=llm_specific_cre.shallow_copy()
947+
)
948+
)
949+
950+
db_mock.return_value.get_gap_analysis_result.return_value = None
951+
db_mock.return_value.gap_analysis_exists.return_value = False
952+
db_mock.return_value.get_nodes.side_effect = lambda name=None, **kwargs: (
953+
[llm]
954+
if name == "OWASP Top 10 for LLM and Gen AI Apps 2025"
955+
else (
956+
[generic_cheatsheet, llm_cheatsheet]
957+
if name == "OWASP Cheat Sheets"
958+
else []
959+
)
960+
)
961+
schedule_mock.side_effect = RuntimeError("redis unavailable")
962+
upstream_fetch_mock.side_effect = RuntimeError("upstream unavailable")
963+
964+
with self.app.test_client() as client:
965+
response = client.get(
966+
"/rest/v1/map_analysis?standard=OWASP%20Top%2010%20for%20LLM%20and%20Gen%20AI%20Apps%202025&standard=OWASP%20Cheat%20Sheets",
967+
headers={"Content-Type": "application/json"},
968+
)
969+
970+
payload = json.loads(response.data)
971+
self.assertEqual(200, response.status_code)
972+
self.assertIn("result", payload)
973+
self.assertEqual(
974+
"AI / LLM Cheat Sheets",
975+
payload["specialized_cheatsheet_section"]["category"],
976+
)
977+
self.assertIn(llm.id, payload["result"])
978+
self.assertNotIn(llm_cheatsheet.id, payload["result"][llm.id]["paths"])
979+
self.assertNotIn(generic_cheatsheet.id, payload["result"][llm.id]["paths"])
980+
981+
@patch.object(web_main.cre_main, "fetch_upstream_json")
982+
@patch.object(web_main.gap_analysis, "schedule")
983+
@patch.object(db, "Node_collection")
984+
def test_gap_analysis_keeps_prompt_injection_cheatsheet_for_llm01(
985+
self, db_mock, schedule_mock, upstream_fetch_mock
986+
) -> None:
987+
prompt_cre = defs.CRE(
988+
id="161-451", name="Prompt boundary protection", description=""
989+
)
990+
991+
llm = defs.Standard(
992+
name="OWASP Top 10 for LLM and Gen AI Apps 2025",
993+
sectionID="LLM01",
994+
section="Prompt Injection",
995+
)
996+
llm.add_link(
997+
defs.Link(ltype=defs.LinkTypes.LinkedTo, document=prompt_cre.shallow_copy())
998+
)
999+
1000+
llm_cheatsheet = defs.Standard(
1001+
name="OWASP Cheat Sheets",
1002+
section="LLM Prompt Injection Prevention Cheat Sheet",
1003+
)
1004+
llm_cheatsheet.add_link(
1005+
defs.Link(ltype=defs.LinkTypes.LinkedTo, document=prompt_cre.shallow_copy())
1006+
)
1007+
1008+
db_mock.return_value.get_gap_analysis_result.return_value = None
1009+
db_mock.return_value.gap_analysis_exists.return_value = False
1010+
db_mock.return_value.get_nodes.side_effect = lambda name=None, **kwargs: (
1011+
[llm]
1012+
if name == "OWASP Top 10 for LLM and Gen AI Apps 2025"
1013+
else [llm_cheatsheet]
1014+
if name == "OWASP Cheat Sheets"
1015+
else []
1016+
)
1017+
schedule_mock.side_effect = RuntimeError("redis unavailable")
1018+
upstream_fetch_mock.side_effect = RuntimeError("upstream unavailable")
1019+
1020+
with self.app.test_client() as client:
1021+
response = client.get(
1022+
"/rest/v1/map_analysis?standard=OWASP%20Top%2010%20for%20LLM%20and%20Gen%20AI%20Apps%202025&standard=OWASP%20Cheat%20Sheets",
1023+
headers={"Content-Type": "application/json"},
1024+
)
1025+
1026+
payload = json.loads(response.data)
1027+
self.assertEqual(200, response.status_code)
1028+
self.assertIn(llm.id, payload["result"])
1029+
self.assertIn(llm_cheatsheet.id, payload["result"][llm.id]["paths"])
1030+
1031+
@patch.object(web_main.cre_main, "fetch_upstream_json")
1032+
@patch.object(web_main.gap_analysis, "schedule")
1033+
@patch.object(db, "Node_collection")
1034+
def test_gap_analysis_adds_specialized_section_for_api_cheatsheets(
1035+
self, db_mock, schedule_mock, upstream_fetch_mock
1036+
) -> None:
1037+
authz_cre = defs.CRE(
1038+
id="117-371", name="Centralized access control", description=""
1039+
)
1040+
base = defs.Standard(
1041+
name="OWASP API Security Top 10 2023",
1042+
sectionID="API1",
1043+
section="Broken Object Level Authorization",
1044+
)
1045+
base.add_link(
1046+
defs.Link(ltype=defs.LinkTypes.LinkedTo, document=authz_cre.shallow_copy())
1047+
)
1048+
1049+
compare = defs.Standard(
1050+
name="OWASP Cheat Sheets",
1051+
section="Authorization Cheat Sheet",
1052+
)
1053+
compare.add_link(
1054+
defs.Link(ltype=defs.LinkTypes.LinkedTo, document=authz_cre.shallow_copy())
1055+
)
1056+
1057+
db_mock.return_value.get_gap_analysis_result.return_value = None
1058+
db_mock.return_value.gap_analysis_exists.return_value = False
1059+
db_mock.return_value.get_nodes.side_effect = lambda name=None, **kwargs: (
1060+
[base]
1061+
if name == "OWASP API Security Top 10 2023"
1062+
else [compare] if name == "OWASP Cheat Sheets" else []
1063+
)
1064+
schedule_mock.side_effect = RuntimeError("redis unavailable")
1065+
upstream_fetch_mock.side_effect = RuntimeError("upstream unavailable")
1066+
1067+
with self.app.test_client() as client:
1068+
response = client.get(
1069+
"/rest/v1/map_analysis?standard=OWASP%20API%20Security%20Top%2010%202023&standard=OWASP%20Cheat%20Sheets",
1070+
headers={"Content-Type": "application/json"},
1071+
)
1072+
1073+
payload = json.loads(response.data)
1074+
self.assertEqual(200, response.status_code)
1075+
self.assertEqual(
1076+
"API Cheat Sheets",
1077+
payload["specialized_cheatsheet_section"]["category"],
1078+
)
1079+
self.assertIn(base.id, payload["specialized_cheatsheet_section"]["result"])
1080+
8811081
@patch.object(web_main.gap_analysis, "schedule")
8821082
@patch.object(db, "Node_collection")
8831083
def test_gap_analysis_supports_opencre_as_standard(
@@ -1222,6 +1422,62 @@ def test_gap_analysis_heroku_cache_miss_returns_404(
12221422
self.assertEqual(404, response.status_code)
12231423
redis_conn_mock.assert_not_called()
12241424

1425+
@patch.object(web_main.redis, "connect")
1426+
@patch.object(rq.job.Job, "fetch")
1427+
@patch.object(db, "Node_collection")
1428+
def test_ma_job_results_adds_specialized_cheatsheet_section(
1429+
self, db_mock, fetch_mock, redis_connect_mock
1430+
) -> None:
1431+
shared_cre = defs.CRE(id="161-451", name="Prompt boundary protection", description="")
1432+
llm = defs.Standard(
1433+
name="OWASP Top 10 for LLM and Gen AI Apps 2025",
1434+
sectionID="LLM01",
1435+
section="Prompt Injection",
1436+
)
1437+
llm.add_link(
1438+
defs.Link(ltype=defs.LinkTypes.LinkedTo, document=shared_cre.shallow_copy())
1439+
)
1440+
cheatsheet = defs.Standard(
1441+
name="OWASP Cheat Sheets",
1442+
section="LLM Prompt Injection Prevention Cheat Sheet",
1443+
)
1444+
cheatsheet.add_link(
1445+
defs.Link(ltype=defs.LinkTypes.LinkedTo, document=shared_cre.shallow_copy())
1446+
)
1447+
1448+
cached_gap = {
1449+
"result": {
1450+
llm.id: {
1451+
"start": llm.shallow_copy().todict(),
1452+
"paths": {
1453+
cheatsheet.id: {
1454+
"end": cheatsheet.shallow_copy().todict(),
1455+
"path": [],
1456+
"score": 0,
1457+
}
1458+
},
1459+
"extra": 0,
1460+
}
1461+
}
1462+
}
1463+
1464+
fetch_mock.return_value = MockFinishedJob()
1465+
redis_connect_mock.return_value.exists.return_value = True
1466+
db_mock.return_value.get_gap_analysis_result.return_value = json.dumps(cached_gap)
1467+
1468+
with self.app.test_client() as client:
1469+
response = client.get(
1470+
"/rest/v1/ma_job_results?id=ABC",
1471+
headers={"Content-Type": "application/json"},
1472+
)
1473+
1474+
payload = json.loads(response.data)
1475+
self.assertEqual(200, response.status_code)
1476+
self.assertEqual(
1477+
"AI / LLM Cheat Sheets",
1478+
payload["specialized_cheatsheet_section"]["category"],
1479+
)
1480+
12251481
@patch.object(redis, "from_url")
12261482
@patch.object(db, "Node_collection")
12271483
def test_standards_from_db(self, node_mock, redis_conn_mock) -> None:

0 commit comments

Comments
 (0)