Coverage for product_risk_suite/product/admin.py: 93%

1from django.contrib import admin, messages 

2from django.forms import SelectMultiple 

3from django.db import models 

4 

5from .models import Product, ProductRiskAnalysis, ProductRiskEntry 

6from risk_assessment.models import Risk 

7from threat_model.shared_models import ThreatModelConnectionName 

8from guardian.admin import GuardedModelAdmin 

9 

10class ProductAdmin(GuardedModelAdmin): 

11 list_display = ["title", "description"] 

12 prepopulated_fields = {"slug": ("title",)} 

13 

14class ProductRiskAnalysisAdmin(admin.ModelAdmin): 

15 filter_horizontal = ['risk_entries'] 

16 

17 list_display = ["name", "n_entries"] 

18 prepopulated_fields = {"slug": ("name",)} 

19 

20 @admin.display(description='N Entries') 

21 def n_entries(self, obj): 

22 return len(obj.risk_entries.all()) 

23 

24def upgrade_to_latest_risk_version(modeladmin, request, queryset): 

25 upgraded = 0 

26 already_current = 0 

27 for entry in queryset.select_related('risk'): 

28 if entry.risk.is_current: 

29 already_current += 1 

30 continue 

31 try: 

32 latest = Risk.objects.get(custom_id=entry.risk.custom_id, is_current=True) 

33 entry.risk = latest 

34 entry.save() 

35 upgraded += 1 

36 except Risk.DoesNotExist: 

37 modeladmin.message_user( 

38 request, 

39 f"No current version found for {entry.risk.custom_id} — skipped.", 

40 messages.ERROR, 

41 ) 

42 

43 if upgraded: 

44 modeladmin.message_user(request, f"Upgraded {upgraded} entry/entries to the latest risk version.", messages.SUCCESS) 

45 if already_current: 

46 modeladmin.message_user(request, f"{already_current} entry/entries were already on the current version.", messages.INFO) 

47 

48upgrade_to_latest_risk_version.short_description = "Upgrade to latest risk version" 

49 

50 

51class ProductRiskEntryAdmin(admin.ModelAdmin): 

52 readonly_fields = ('created_at_display', 'last_change_display') 

53 ordering = ['risk__custom_id'] 

54 

55 def created_at_display(self, obj): 

56 return obj.created_at.strftime('%Y-%m-%d %H:%M:%S') if obj.created_at else '-' 

57 created_at_display.short_description = 'Created at' 

58 

59 def last_change_display(self, obj): 

60 return obj.last_change.strftime('%Y-%m-%d %H:%M:%S') if obj.last_change else '-' 

61 last_change_display.short_description = 'Last update' 

62 

63 def get_form(self, request, obj=None, **kwargs): 

64 form = super().get_form(request, obj, **kwargs) 

65 if obj is None: 

66 form.base_fields['svg_id'].help_text = "Here all available IDs are listed, if you do not know which matches, assign this entry to a product analysis and that to a product, then this list is reduced." 

67 else: 

68 form.base_fields['svg_id'].help_text = "Select one or more threat model connections this risk applies to." + "<br>" + form.base_fields['svg_id'].help_text 

69 return form 

70 

71 def formfield_for_manytomany(self, db_field, request, **kwargs): 

72 if db_field.name == "svg_id": 

73 # Get the current ProductRiskEntry instance (if editing) 

74 obj = request.resolver_match.kwargs.get('object_id') 

75 if obj: 

76 try: 

77 instance = ProductRiskEntry.objects.get(pk=obj) 

78 ras = ProductRiskAnalysis.objects.get(risk_entries=instance) 

79 products = Product.objects.get(analyzes=ras) 

80 connection_names = ThreatModelConnectionName.objects.filter( 

81 threatmodel__product=products 

82 ).distinct().order_by("tech_name") 

83 kwargs['queryset'] = connection_names 

84 except (ProductRiskEntry.DoesNotExist, AttributeError): 

85 pass 

86 except (ProductRiskAnalysis.DoesNotExist, AttributeError): 

87 pass 

88 except (Product.DoesNotExist, AttributeError): 

89 pass 

90 else: 

91 # for a new instance ALL available tm link names are listed 

92 pass 

93 return super().formfield_for_manytomany(db_field, request, **kwargs) 

94 

95 filter_horizontal = ['evidences','svg_id'] 

96 actions = [upgrade_to_latest_risk_version] 

97 

98 list_display = [ 

99 "risk_Id", 

100 "risk_outdated", 

101 "risk_title", 

102 "risk_asset", 

103 "risk_origin", 

104 "risk_stride", 

105 "tm_linked", 

106 "risk_initial", 

107 "risk_accepted", 

108 "risk_mitigated", 

109 "risk_after_mitigation", 

110 ] 

111 

112 @admin.display(description='Risk Id') 

113 def risk_Id(self, obj): 

114 return obj.risk.custom_id 

115 

116 @admin.display(description='Current version', boolean=True) 

117 def risk_outdated(self, obj): 

118 return not obj.risk.has_newer_version 

119 

120 @admin.display(description='risk title') 

121 def risk_title(self, obj): 

122 return obj.risk.title 

123 

124 @admin.display(description='Asset') 

125 def risk_asset(self, obj): 

126 return obj.risk.asset 

127 

128 @admin.display(description='Origin') 

129 def risk_origin(self, obj): 

130 return obj.risk.origin 

131 

132 @admin.display(description='Stride') 

133 def risk_stride(self, obj): 

134 return obj.risk.stride_str 

135 

136 @admin.display(description='Risk Accepted') 

137 def risk_accepted(self, obj): 

138 return obj.risk_accepted 

139 

140 @admin.display(description='TM Linked') 

141 def tm_linked(self, obj): 

142 return len(obj.svg_id.all()) > 0 

143 tm_linked.boolean = True 

144 

145 @admin.display(description='Risk Initial') 

146 def risk_initial(self, obj): 

147 return obj.risk_rating_initial.risk 

148 

149 @admin.display(description='Mitigated') 

150 def risk_mitigated(self, obj): 

151 return obj.risk_mitigation != None 

152 

153 @admin.display(description='Risk after Mitigation') 

154 def risk_after_mitigation(self, obj): 

155 if obj.risk_rating_after_mitigation: 

156 return obj.risk_rating_after_mitigation.risk 

157 return None 

158 

159admin.site.register(Product, ProductAdmin) 

160admin.site.register(ProductRiskAnalysis, ProductRiskAnalysisAdmin) 

161admin.site.register(ProductRiskEntry, ProductRiskEntryAdmin)