Coverage for product_risk_suite/threat_model/tests.py: 100%

1import json 

2 

3from django.conf import settings 

4from django.contrib.auth.models import User 

5from django.test import TestCase 

6from django.core.files.uploadedfile import SimpleUploadedFile 

7 

8from guardian.shortcuts import assign_perm 

9 

10from .models import * 

11from product.models import ProductRiskEntry, ProductRiskAnalysis 

12from risk_assessment.models import Asset, Origin, LifeCycle, Stride, Risk, RiskRating 

13 

14import tempfile 

15settings.MEDIA_ROOT = tempfile.mkdtemp("prs-tm") 

16 

17_SVG_BYTES = b'<svg xmlns="http://www.w3.org/2000/svg" style="color-scheme: light dark;"><g data-cell-id="cn-1"></g></svg>' 

18 

19class ThreatModelConnectionNameTest(TestCase): 

20 def setUp(self): 

21 Product.objects.create(id=0, title="My test project", description="A real cool test project") 

22 

23 def test_simple(self): 

24 td = [ 

25 {"tech_name": "0", "human_str": None}, 

26 {"tech_name": "1", "human_str": None}, 

27 {"tech_name": "DGsSiznoInNCpWqt2XE2-2", "human_str": None}, 

28 {"tech_name": "DGsSiznoInNCpWqt2XE2-1", "human_str": None}, 

29 {"tech_name": "DGsSiznoInNCpWqt2XE2-3", "human_str": None}, 

30 {"tech_name": "DGsSiznoInNCpWqt2XE2-4", "human_str": None}, 

31 {"tech_name": "DGsSiznoInNCpWqt2XE2-5", "human_str": None}, 

32 {"tech_name": "admin-connection", "human_str": None}, 

33 {"tech_name": "DGsSiznoInNCpWqt2XE2-11", "human_str": None}, 

34 {"tech_name": "rw-connection", "human_str": "Read Write Connection"}, 

35 {"tech_name": "DGsSiznoInNCpWqt2XE2-10", "human_str": None}, 

36 {"tech_name": "ro-connection", "human_str": None}, 

37 {"tech_name": "DGsSiznoInNCpWqt2XE2-9", "human_str": None}, 

38 ] 

39 

40 for tr in td: 

41 ThreatModelConnectionName.objects.create(tech_name=tr["tech_name"], human_name=tr["human_str"], product=Product.objects.get(id=0)) 

42 

43 all = ThreatModelConnectionName.objects.all() 

44 self.assertEqual(len(all), len(td)) 

45 

46 act_with_human_str = 0 

47 for tr in td: 

48 tmcn = ThreatModelConnectionName.objects.get(tech_name=tr["tech_name"]) 

49 if tr["human_str"]: 

50 self.assertEqual(str(tmcn), f"My test project: {tr['human_str']} ({tr['tech_name']})") 

51 act_with_human_str = act_with_human_str + 1 

52 else: 

53 self.assertEqual(str(tmcn), f"My test project: {tr['tech_name']}") 

54 self.assertEqual(act_with_human_str, 1) 

55 

56class ThreatModelTest(TestCase): 

57 def setUp(self): 

58 Product.objects.create(id=0, title="My test project", description="A real cool test project") 

59 

60 def test_simple(self): 

61 tm = ThreatModel.objects.create( 

62 id=0, 

63 product=Product.objects.get(id=0), 

64 title="Test TM", 

65 description="Description of TM" 

66 ) 

67 tm.svg = SimpleUploadedFile( 

68 name="test-threat-model.svg", 

69 content=b""" 

70 <svg> 

71 <g data-cell-id="0"> 

72 <g data-cell-id="1"> 

73 <g data-cell-id="rw connection"> 

74 </g> 

75 <g data-cell-id="usb"> 

76 </g> 

77 </g> 

78 </g> 

79 </svg> 

80 """, 

81 content_type="image/svg+xml" 

82 ) 

83 tm.save() 

84 

85 actual = ThreatModel.objects.get(id=0) 

86 self.assertEqual(actual.product.title, "My test project") 

87 self.assertEqual(actual.title, "Test TM") 

88 self.assertEqual(actual.description, "Description of TM") 

89 self.assertEqual(actual.svg, "uploads/test-threat-model.svg") 

90 act_connection_names = [cn.tech_name for cn in actual.connection_names.all()] 

91 self.assertEqual(len(act_connection_names), 4) 

92 self.assertTrue("0" in act_connection_names) 

93 self.assertTrue("1" in act_connection_names) 

94 self.assertTrue("rw connection" in act_connection_names) 

95 self.assertTrue("usb" in act_connection_names) 

96 

97 def test_str(self): 

98 tm = ThreatModel.objects.create( 

99 product=Product.objects.get(id=0), 

100 title="My Diagram", 

101 svg='', 

102 ) 

103 self.assertEqual(str(tm), "My Diagram") 

104 

105 def test_list_connection_names_empty_when_none_assigned(self): 

106 tm = ThreatModel.objects.create( 

107 product=Product.objects.get(id=0), 

108 title="Empty TM", 

109 svg='', 

110 ) 

111 self.assertEqual(tm.list_connection_names, []) 

112 

113 def test_list_connection_names_returns_tech_names(self): 

114 product = Product.objects.get(id=0) 

115 cn1 = ThreatModelConnectionName.objects.create(product=product, tech_name="conn-alpha") 

116 cn2 = ThreatModelConnectionName.objects.create(product=product, tech_name="conn-beta") 

117 tm = ThreatModel.objects.create(product=product, title="Linked TM", svg='') 

118 tm.connection_names.set([cn1, cn2]) 

119 self.assertCountEqual(tm.list_connection_names, ["conn-alpha", "conn-beta"]) 

120 

121 

122# --------------------------------------------------------------------------- 

123# View tests 

124# --------------------------------------------------------------------------- 

125 

126def _make_risk(custom_id, asset, origin, lc, stride): 

127 risk = Risk(custom_id=custom_id, asset=asset, origin=origin, life_cycle=lc, title=f"Risk {custom_id}", description="d") 

128 risk.save() 

129 risk.stride.set([stride]) 

130 risk.save() 

131 return risk 

132 

133 

134class ThreatModelIndexViewTest(TestCase): 

135 def setUp(self): 

136 self.user = User.objects.create_user(username="viewer", password="pw") 

137 self.product = Product.objects.create(title="Accessible Product", description="d", slug="accessible-product") 

138 other_product = Product.objects.create(title="Other Product", description="d", slug="other-product") 

139 assign_perm("product.view_product", self.user, self.product) 

140 

141 self.tm_accessible = ThreatModel.objects.create(product=self.product, title="Accessible TM", svg='') 

142 self.tm_other = ThreatModel.objects.create(product=other_product, title="Other TM", svg='') 

143 

144 def test_unauthenticated_redirects_to_login(self): 

145 response = self.client.get('/threat-models/') 

146 self.assertRedirects(response, '/login?next=/threat-models/', fetch_redirect_response=False) 

147 

148 def test_authenticated_returns_200(self): 

149 self.client.force_login(self.user) 

150 self.assertEqual(self.client.get('/threat-models/').status_code, 200) 

151 

152 def test_context_includes_accessible_tm(self): 

153 self.client.force_login(self.user) 

154 response = self.client.get('/threat-models/') 

155 self.assertIn(self.tm_accessible, response.context['tm_list']) 

156 

157 def test_context_excludes_inaccessible_tm(self): 

158 self.client.force_login(self.user) 

159 response = self.client.get('/threat-models/') 

160 self.assertNotIn(self.tm_other, response.context['tm_list']) 

161 

162 def test_context_empty_for_user_without_any_permission(self): 

163 no_perm_user = User.objects.create_user(username="noperm", password="pw") 

164 self.client.force_login(no_perm_user) 

165 response = self.client.get('/threat-models/') 

166 self.assertEqual(list(response.context['tm_list']), []) 

167 

168 

169class ThreatModelDetailViewTest(TestCase): 

170 def setUp(self): 

171 self.user = User.objects.create_user(username="viewer", password="pw") 

172 self.product = Product.objects.create(title="Detail Product", description="d", slug="detail-product") 

173 assign_perm("product.view_product", self.user, self.product) 

174 

175 tm = ThreatModel.objects.create(product=self.product, title="Detail TM") 

176 tm.svg = SimpleUploadedFile(name="detail-view-tm.svg", content=_SVG_BYTES, content_type="image/svg+xml") 

177 tm.save() 

178 self.tm = tm 

179 self.url = f'/threat-model/{self.tm.slug}/' 

180 

181 def test_unauthenticated_redirects_to_login(self): 

182 response = self.client.get(self.url) 

183 self.assertRedirects(response, f'/login?next={self.url}', fetch_redirect_response=False) 

184 

185 def test_authenticated_returns_200(self): 

186 self.client.force_login(self.user) 

187 self.assertEqual(self.client.get(self.url).status_code, 200) 

188 

189 def test_nonexistent_slug_returns_404(self): 

190 self.client.force_login(self.user) 

191 self.assertEqual(self.client.get('/threat-model/does-not-exist/').status_code, 404) 

192 

193 def test_inaccessible_product_returns_404(self): 

194 no_perm_user = User.objects.create_user(username="noperm", password="pw") 

195 self.client.force_login(no_perm_user) 

196 self.assertEqual(self.client.get(self.url).status_code, 404) 

197 

198 def test_context_contains_tm(self): 

199 self.client.force_login(self.user) 

200 response = self.client.get(self.url) 

201 self.assertEqual(response.context['tm'], self.tm) 

202 

203 def test_context_svg_content_is_string(self): 

204 self.client.force_login(self.user) 

205 response = self.client.get(self.url) 

206 self.assertIsInstance(response.context['svg_content'], str) 

207 

208 def test_context_svg_removes_dark_mode_style(self): 

209 self.client.force_login(self.user) 

210 response = self.client.get(self.url) 

211 self.assertNotIn('color-scheme: light dark;', response.context['svg_content']) 

212 

213 def test_context_risk_entries_empty_when_no_entries_linked(self): 

214 self.client.force_login(self.user) 

215 response = self.client.get(self.url) 

216 self.assertEqual(json.loads(response.context['risk_entries']), {}) 

217 

218 def test_context_names_empty_when_no_entries_linked(self): 

219 self.client.force_login(self.user) 

220 response = self.client.get(self.url) 

221 self.assertEqual(response.context['names'], {}) 

222 

223 def _setup_risk_entry(self, custom_id="R-TM-001"): 

224 asset = Asset.objects.create(name="TMAsset") 

225 origin = Origin.objects.create(name="TMOrigin") 

226 lc = LifeCycle.objects.create(name="TMLC") 

227 stride = Stride.objects.create(name="S") 

228 rr = RiskRating.objects.create(likelihood_of_occurrence=2, severity_of_impact=3) 

229 risk = _make_risk(custom_id, asset, origin, lc, stride) 

230 cn = ThreatModelConnectionName.objects.get(product=self.product, tech_name="cn-1") 

231 entry = ProductRiskEntry.objects.create(risk=risk, risk_accepted=False, risk_rating_initial=rr) 

232 entry.svg_id.set([cn]) 

233 return entry 

234 

235 def test_context_risk_entries_populated_when_entry_linked(self): 

236 self._setup_risk_entry() 

237 self.client.force_login(self.user) 

238 response = self.client.get(self.url) 

239 entries = json.loads(response.context['risk_entries']) 

240 self.assertIn("cn-1", entries) 

241 self.assertEqual(len(entries["cn-1"]), 1) 

242 

243 def test_context_names_populated_when_entry_linked(self): 

244 self._setup_risk_entry() 

245 self.client.force_login(self.user) 

246 response = self.client.get(self.url) 

247 self.assertIn("cn-1", response.context['names']) 

248 

249 

250class ThreatModelSvgAccessTest(TestCase): 

251 def setUp(self): 

252 self.user = User.objects.create_user(username="viewer", password="pw") 

253 self.other = User.objects.create_user(username="other", password="pw") 

254 self.product = Product.objects.create(title="SVG Product", description="d", slug="svg-product") 

255 assign_perm("product.view_product", self.user, self.product) 

256 

257 tm = ThreatModel.objects.create(product=self.product, title="SVG TM") 

258 tm.svg = SimpleUploadedFile(name="access-test.svg", content=_SVG_BYTES, content_type="image/svg+xml") 

259 tm.save() 

260 self.tm = tm 

261 self.url = f'/media/{tm.svg.name}' 

262 

263 def test_unauthenticated_redirects_to_login(self): 

264 response = self.client.get(self.url) 

265 self.assertRedirects(response, f'/login?next={self.url}', fetch_redirect_response=False) 

266 

267 def test_authorized_user_gets_200(self): 

268 self.client.force_login(self.user) 

269 response = self.client.get(self.url) 

270 self.assertEqual(response.status_code, 200) 

271 self.assertEqual(response['Content-Type'], 'image/svg+xml') 

272 

273 def test_unauthorized_user_gets_403(self): 

274 self.client.force_login(self.other) 

275 response = self.client.get(self.url) 

276 self.assertEqual(response.status_code, 403) 

277 

278 def test_unknown_file_gets_404(self): 

279 self.client.force_login(self.user) 

280 response = self.client.get('/media/uploads/does-not-exist.svg') 

281 self.assertEqual(response.status_code, 404) 

282 

283 

284# --------------------------------------------------------------------------- 

285# Admin tests 

286# --------------------------------------------------------------------------- 

287 

288from django.contrib.admin.sites import AdminSite 

289from threat_model.admin import ThreatModelAdmin, ThreatModelConnectionNameAdmin 

290 

291 

292class ThreatModelConnectionNameAdminConfigTest(TestCase): 

293 def setUp(self): 

294 self.site = AdminSite() 

295 self.ma = ThreatModelConnectionNameAdmin(ThreatModelConnectionName, self.site) 

296 

297 def test_list_display(self): 

298 self.assertEqual(self.ma.list_display, ["tech_name", "product", "human_name"]) 

299 

300 def test_exclude(self): 

301 self.assertIn('id', self.ma.exclude) 

302 self.assertIn('product_id', self.ma.exclude) 

303 

304 def test_ordering(self): 

305 self.assertEqual(self.ma.ordering, ['human_name']) 

306 

307 def test_product_method_returns_product_title(self): 

308 product = Product.objects.create(title="Admin Product", description="d") 

309 cn = ThreatModelConnectionName.objects.create(product=product, tech_name="conn-x") 

310 self.assertEqual(self.ma.product(cn), "Admin Product") 

311 

312 

313class ThreatModelAdminConfigTest(TestCase): 

314 def setUp(self): 

315 self.site = AdminSite() 

316 self.ma = ThreatModelAdmin(ThreatModel, self.site) 

317 self.product = Product.objects.create(title="TM Admin Product", description="d") 

318 

319 def test_list_display(self): 

320 self.assertEqual(self.ma.list_display, ["product", "title", "description", "available_connections"]) 

321 

322 def test_exclude(self): 

323 self.assertIn('slug', self.ma.exclude) 

324 self.assertIn('connection_names', self.ma.exclude) 

325 

326 def test_product_method_returns_product_title(self): 

327 tm = ThreatModel.objects.create(product=self.product, title="My TM", svg='') 

328 self.assertEqual(self.ma.product(tm), "TM Admin Product") 

329 

330 def test_available_connections_zero_when_none(self): 

331 tm = ThreatModel.objects.create(product=self.product, title="Empty TM", svg='') 

332 self.assertEqual(self.ma.available_connections(tm), 0) 

333 

334 def test_available_connections_counts_linked_names(self): 

335 cn1 = ThreatModelConnectionName.objects.create(product=self.product, tech_name="c1") 

336 cn2 = ThreatModelConnectionName.objects.create(product=self.product, tech_name="c2") 

337 tm = ThreatModel.objects.create(product=self.product, title="Linked TM", svg='') 

338 tm.connection_names.set([cn1, cn2]) 

339 self.assertEqual(self.ma.available_connections(tm), 2) 

340 

341 

342class ThreatModelConnectionNameAdminViewTest(TestCase): 

343 def setUp(self): 

344 self.superuser = User.objects.create_superuser(username="admin", password="pw") 

345 self.product = Product.objects.create(title="View Product", description="d") 

346 self.cn = ThreatModelConnectionName.objects.create(product=self.product, tech_name="view-conn") 

347 

348 def test_changelist_returns_200(self): 

349 self.client.force_login(self.superuser) 

350 response = self.client.get('/admin/threat_model/threatmodelconnectionname/') 

351 self.assertEqual(response.status_code, 200) 

352 

353 def test_changeform_returns_200(self): 

354 self.client.force_login(self.superuser) 

355 response = self.client.get(f'/admin/threat_model/threatmodelconnectionname/{self.cn.pk}/change/') 

356 self.assertEqual(response.status_code, 200) 

357 

358 def test_changelist_excludes_id_and_product_id_columns(self): 

359 self.client.force_login(self.superuser) 

360 response = self.client.get('/admin/threat_model/threatmodelconnectionname/') 

361 # id and product_id are excluded from the form, not the list — verify list loads correctly 

362 self.assertContains(response, "view-conn") 

363 

364 

365class ThreatModelAdminViewTest(TestCase): 

366 def setUp(self): 

367 self.superuser = User.objects.create_superuser(username="admin", password="pw") 

368 self.product = Product.objects.create(title="TM View Product", description="d") 

369 self.tm = ThreatModel.objects.create(product=self.product, title="Admin View TM", svg='') 

370 

371 def test_changelist_returns_200(self): 

372 self.client.force_login(self.superuser) 

373 response = self.client.get('/admin/threat_model/threatmodel/') 

374 self.assertEqual(response.status_code, 200) 

375 

376 def test_changelist_shows_available_connections_column(self): 

377 self.client.force_login(self.superuser) 

378 response = self.client.get('/admin/threat_model/threatmodel/') 

379 self.assertContains(response, "Admin View TM") 

380 

381 def test_changeform_returns_200(self): 

382 self.client.force_login(self.superuser) 

383 response = self.client.get(f'/admin/threat_model/threatmodel/{self.tm.pk}/change/') 

384 self.assertEqual(response.status_code, 200) 

385 

386 def test_changeform_excludes_slug_and_connection_names(self): 

387 self.client.force_login(self.superuser) 

388 response = self.client.get(f'/admin/threat_model/threatmodel/{self.tm.pk}/change/') 

389 # slug and connection_names inputs must not appear in the form 

390 self.assertNotContains(response, 'name="slug"') 

391 self.assertNotContains(response, 'name="connection_names"')