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

1from django import forms 

2from django.contrib import admin, messages 

3from django.contrib.auth import get_user_model 

4from django.contrib.auth.models import Group 

5from django.forms import SelectMultiple 

6from django.db import models 

7from django.utils.translation import gettext_lazy as _ 

8 

9from .models import Product, ProductRiskAnalysis, ProductRiskEntry 

10from risk_assessment.models import Risk 

11from threat_model.shared_models import ThreatModelConnectionName 

12from guardian.admin import GuardedModelAdmin 

13 

14 

15class ObjPermUserSelectForm(forms.Form): 

16 user = forms.ModelChoiceField( 

17 label=_("User identification"), 

18 queryset=get_user_model()._default_manager.order_by(get_user_model().USERNAME_FIELD), 

19 ) 

20 

21 

22class ObjPermGroupSelectForm(forms.Form): 

23 group = forms.ModelChoiceField( 

24 label=_("Group"), 

25 queryset=Group.objects.order_by("name"), 

26 ) 

27 

28 

29class ProductAdmin(GuardedModelAdmin): 

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

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

32 

33 def get_obj_perms_user_select_form(self, request): 

34 return ObjPermUserSelectForm 

35 

36 def get_obj_perms_group_select_form(self, request): 

37 return ObjPermGroupSelectForm 

38 

39class ProductRiskAnalysisAdmin(admin.ModelAdmin): 

40 filter_horizontal = ['risk_entries'] 

41 

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

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

44 

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

46 def n_entries(self, obj): 

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

48 

49def upgrade_to_latest_risk_version(modeladmin, request, queryset): 

50 upgraded = 0 

51 already_current = 0 

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

53 if entry.risk.is_current: 

54 already_current += 1 

55 continue 

56 try: 

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

58 entry.risk = latest 

59 entry.save() 

60 upgraded += 1 

61 except Risk.DoesNotExist: 

62 modeladmin.message_user( 

63 request, 

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

65 messages.ERROR, 

66 ) 

67 

68 if upgraded: 

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

70 if already_current: 

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

72 

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

74 

75 

76class ProductRiskEntryAdmin(admin.ModelAdmin): 

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

78 ordering = ['risk__custom_id'] 

79 

80 def created_at_display(self, obj): 

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

82 created_at_display.short_description = 'Created at' 

83 

84 def last_change_display(self, obj): 

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

86 last_change_display.short_description = 'Last update' 

87 

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

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

90 if obj is None: 

91 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." 

92 else: 

93 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 

94 return form 

95 

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

97 if db_field.name == "svg_id": 

98 # Get the current ProductRiskEntry instance (if editing) 

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

100 if obj: 

101 try: 

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

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

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

105 connection_names = ThreatModelConnectionName.objects.filter( 

106 threatmodel__product=products 

107 ).distinct().order_by("tech_name") 

108 kwargs['queryset'] = connection_names 

109 except (ProductRiskEntry.DoesNotExist, AttributeError): 

110 pass 

111 except (ProductRiskAnalysis.DoesNotExist, AttributeError): 

112 pass 

113 except (Product.DoesNotExist, AttributeError): 

114 pass 

115 else: 

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

117 pass 

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

119 

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

121 actions = [upgrade_to_latest_risk_version] 

122 

123 list_display = [ 

124 "risk_Id", 

125 "risk_outdated", 

126 "risk_title", 

127 "risk_asset", 

128 "risk_origin", 

129 "risk_stride", 

130 "tm_linked", 

131 "risk_initial", 

132 "risk_accepted", 

133 "risk_mitigated", 

134 "risk_after_mitigation", 

135 ] 

136 

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

138 def risk_Id(self, obj): 

139 return obj.risk.custom_id 

140 

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

142 def risk_outdated(self, obj): 

143 return not obj.risk.has_newer_version 

144 

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

146 def risk_title(self, obj): 

147 return obj.risk.title 

148 

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

150 def risk_asset(self, obj): 

151 return obj.risk.asset 

152 

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

154 def risk_origin(self, obj): 

155 return obj.risk.origin 

156 

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

158 def risk_stride(self, obj): 

159 return obj.risk.stride_str 

160 

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

162 def risk_accepted(self, obj): 

163 return obj.risk_accepted 

164 

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

166 def tm_linked(self, obj): 

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

168 tm_linked.boolean = True 

169 

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

171 def risk_initial(self, obj): 

172 return obj.risk_rating_initial.risk 

173 

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

175 def risk_mitigated(self, obj): 

176 return obj.risk_mitigation != None 

177 

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

179 def risk_after_mitigation(self, obj): 

180 if obj.risk_rating_after_mitigation: 

181 return obj.risk_rating_after_mitigation.risk 

182 return None 

183 

184admin.site.register(Product, ProductAdmin) 

185admin.site.register(ProductRiskAnalysis, ProductRiskAnalysisAdmin) 

186admin.site.register(ProductRiskEntry, ProductRiskEntryAdmin)