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

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 created_at = models.DateTimeField(auto_now_add=True) 

20 last_change = models.DateTimeField(auto_now=True) 

21 

22 def __str__(self): 

23 if self.risk_accepted: 

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

25 if self.risk_mitigation: 

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

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

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

29 

30 @property 

31 def list_svg_ids(self): 

32 return [id.tech_name for id in self.svg_id.all()] 

33 

34class ProductRiskAnalysis(models.Model): 

35 id = models.AutoField(primary_key=True) 

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

37 risk_entries = models.ManyToManyField(ProductRiskEntry) 

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

39 

40 def __str__(self): 

41 return self.name 

42 

43 @staticmethod 

44 def get_risk_summary(risk_entries: ProductRiskEntry): 

45 retval = { 

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

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

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

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

50 "unmitigated_risks": [], 

51 "mitigation_stati": {}, 

52 "covered_sec_reqs" : [] 

53 } 

54 sec_reqs = {} 

55 mitigation_stati = {} 

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

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

58 for e in risk_entries.all(): 

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

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

61 if e.risk_mitigation: 

62 if e.risk_rating_after_mitigation: 

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

64 else: 

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

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

67 if e.risk_mitigation: 

68 for sr in e.risk_mitigation.list_security_requirements: 

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

70 if e.evidences: 

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

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

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

74 retval["mitigation_stati"] = mitigation_stati 

75 

76 return retval 

77 

78class Product(models.Model): 

79 id = models.AutoField(primary_key=True) 

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

81 description = models.TextField() 

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

83 

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

85 

86 class Meta: 

87 permissions = () 

88 

89 def __str__(self): 

90 return self.title