Skip to content

Commit da8944f

Browse files
Bornunique911Bornunique911
committed
Add specialized OWASP cheat sheet map analysis behavior
1 parent 014456f commit da8944f

2 files changed

Lines changed: 426 additions & 0 deletions

File tree

‎application/tests/web_main_test.py‎

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

3535

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

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

1373+
@patch.object(web_main.redis, "connect")
1374+
@patch.object(rq.job.Job, "fetch")
1375+
@patch.object(db, "Node_collection")
1376+
def test_ma_job_results_adds_specialized_cheatsheet_section(
1377+
self, db_mock, fetch_mock, redis_connect_mock
1378+
) -> None:
1379+
shared_cre = defs.CRE(id="161-451", name="Prompt boundary protection", description="")
1380+
llm = defs.Standard(
1381+
name="OWASP Top 10 for LLM and Gen AI Apps 2025",
1382+
sectionID="LLM01",
1383+
section="Prompt Injection",
1384+
)
1385+
llm.add_link(
1386+
defs.Link(ltype=defs.LinkTypes.LinkedTo, document=shared_cre.shallow_copy())
1387+
)
1388+
cheatsheet = defs.Standard(
1389+
name="OWASP Cheat Sheets",
1390+
section="LLM Prompt Injection Prevention Cheat Sheet",
1391+
)
1392+
cheatsheet.add_link(
1393+
defs.Link(ltype=defs.LinkTypes.LinkedTo, document=shared_cre.shallow_copy())
1394+
)
1395+
1396+
cached_gap = {
1397+
"result": {
1398+
llm.id: {
1399+
"start": llm.shallow_copy().todict(),
1400+
"paths": {
1401+
cheatsheet.id: {
1402+
"end": cheatsheet.shallow_copy().todict(),
1403+
"path": [],
1404+
"score": 0,
1405+
}
1406+
},
1407+
"extra": 0,
1408+
}
1409+
}
1410+
}
1411+
1412+
fetch_mock.return_value = MockFinishedJob()
1413+
redis_connect_mock.return_value.exists.return_value = True
1414+
db_mock.return_value.get_gap_analysis_result.return_value = json.dumps(cached_gap)
1415+
1416+
with self.app.test_client() as client:
1417+
response = client.get(
1418+
"/rest/v1/ma_job_results?id=ABC",
1419+
headers={"Content-Type": "application/json"},
1420+
)
1421+
1422+
payload = json.loads(response.data)
1423+
self.assertEqual(200, response.status_code)
1424+
self.assertEqual(
1425+
"AI / LLM Cheat Sheets",
1426+
payload["specialized_cheatsheet_section"]["category"],
1427+
)
1428+
11731429
@patch.object(redis, "from_url")
11741430
@patch.object(db, "Node_collection")
11751431
def test_standards_from_db(self, node_mock, redis_conn_mock) -> None:

0 commit comments

Comments
 (0)