Coverage for product_risk_suite / product / models.py: 98%

1from django.db import models 

2 

3from collections import OrderedDict 

4 

5from risk_assessment.models import Risk, RiskRating, RiskMitigation, Evidence, Risk5x5, Status 

6from threat_model.shared_models import ThreatModelConnectionName 

7 

8class ProductRiskEntry(models.Model): 

9 id = models.AutoField(primary_key=True) 

10 risk = models.ForeignKey(Risk, on_delete=models.PROTECT) 

11 risk_accepted = models.BooleanField(default=False) 

12 risk_rating_initial = models.ForeignKey(RiskRating, on_delete=models.PROTECT, related_name='initial') 

13 risk_mitigation = models.ForeignKey(RiskMitigation, on_delete=models.PROTECT, blank=True, null=True) 

14 evidences = models.ManyToManyField(Evidence, blank=True) 

15 risk_rating_after_mitigation = models.ForeignKey(RiskRating, on_delete=models.PROTECT, related_name='after_mitigation', blank=True, null=True) 

16 

17 svg_id = models.ManyToManyField(ThreatModelConnectionName, blank=True) 

18 

19 def __str__(self): 

20 if self.risk_accepted: 

21 return f"Accepted Risk: {self.risk.custom_id} - {self.risk.title}" 

22 if self.risk_mitigation: 

23 r = self.risk_rating_after_mitigation.risk if self.risk_rating_after_mitigation else "forgotten" 

24 return f"Mitigated Risk: {self.risk.custom_id} - {self.risk.title} Risk: {r}" 

25 return f"*Un*mitigated Risk: {self.risk.custom_id} - {self.risk.title} Open risk: {self.risk_rating_initial.risk}" 

26 

27 @property 

28 def list_svg_ids(self): 

29 return [id.name for id in self.svg_id.all()] 

30 

31class ProductRiskAnalysis(models.Model): 

32 id = models.AutoField(primary_key=True) 

33 name = models.CharField(max_length=200, unique=True) 

34 risk_entries = models.ManyToManyField(ProductRiskEntry) 

35 slug = models.SlugField(default="", null=False, unique=True) 

36 

37 def __str__(self): 

38 return self.name 

39 

40 @staticmethod 

41 def get_risk_summary(risk_entries: ProductRiskEntry): 

42 retval = { 

43 "risk_severities_before_mitigation" : { Risk5x5.LOW.name : 0, Risk5x5.MID.name: 0, Risk5x5.HIGH.name: 0}, 

44 "risk_severities_after_mitigation" : { Risk5x5.LOW.name : 0, Risk5x5.MID.name: 0, Risk5x5.HIGH.name: 0}, 

45 "risk_severity_colors": { Risk5x5.LOW.name : "#d0e6dc", Risk5x5.MID.name: "#fff3cc", Risk5x5.HIGH.name: "#f8d6d9"}, 

46 "n_risks": len(risk_entries.all()), 

47 "unmitigated_risks": [], 

48 "mitigation_stati": {}, 

49 "covered_sec_reqs" : [] 

50 } 

51 sec_reqs = {} 

52 mitigation_stati = {} 

53 for k,v in Status.STATUS_CHOICES.items(): 

54 mitigation_stati[Status.web_idx_name(k)] = 0 

55 for e in risk_entries.all(): 

56 rate = RiskRating.risk_level(e.risk_rating_initial.risk) 

57 retval["risk_severities_before_mitigation"][rate.name] = retval["risk_severities_before_mitigation"][rate.name] + 1 

58 if e.risk_mitigation: 

59 if e.risk_rating_after_mitigation: 

60 rate = RiskRating.risk_level(e.risk_rating_after_mitigation.risk) 

61 else: 

62 retval["unmitigated_risks"].append(e.risk.custom_id) 

63 retval["risk_severities_after_mitigation"][rate.name] = retval["risk_severities_after_mitigation"][rate.name] + 1 

64 if e.risk_mitigation: 

65 for sr in e.risk_mitigation.list_security_requirements: 

66 sec_reqs[sr["str"]] = sr["slug"] 

67 if e.evidences: 

68 for ev in e.evidences.all(): 

69 mitigation_stati[Status.web_idx_name(ev.status.status)] = mitigation_stati[Status.web_idx_name(ev.status.status)] + 1 

70 retval["covered_sec_reqs"] = OrderedDict(sorted(sec_reqs.items())) 

71 retval["mitigation_stati"] = mitigation_stati 

72 

73 return retval 

74 

75class Product(models.Model): 

76 id = models.AutoField(primary_key=True) 

77 title = models.CharField(max_length=200, unique=True) 

78 description = models.TextField() 

79 slug = models.SlugField(default="", null=False, unique=True) 

80 

81 analyzes = models.ManyToManyField(ProductRiskAnalysis, blank=True) 

82 

83 class Meta: 

84 permissions = () 

85 

86 def __str__(self): 

87 return self.title