|
9 | 9 |
|
10 | 10 | from django import forms |
11 | 11 | from django.contrib import admin |
| 12 | +from django.contrib.admin.widgets import FilteredSelectMultiple |
| 13 | +from django.contrib.auth.admin import GroupAdmin as BasicGroupAdmin |
| 14 | +from django.contrib.auth.models import Group |
| 15 | +from django.contrib.auth.models import User |
12 | 16 | from django.core.validators import validate_email |
13 | 17 |
|
14 | 18 | from vulnerabilities.models import ApiUser |
|
17 | 21 | from vulnerabilities.models import VulnerabilityReference |
18 | 22 | from vulnerabilities.models import VulnerabilitySeverity |
19 | 23 |
|
| 24 | +admin.site.site_header = "VulnerableCode Administration" |
| 25 | +admin.site.site_title = "VulnerableCode Admin Portal" |
| 26 | +admin.site.index_title = "Welcome to VulnerableCode Management" |
| 27 | + |
20 | 28 |
|
21 | 29 | @admin.register(Vulnerability) |
22 | 30 | class VulnerabilityAdmin(admin.ModelAdmin): |
@@ -97,3 +105,50 @@ def get_form(self, request, obj=None, **kwargs): |
97 | 105 | defaults["form"] = self.add_form |
98 | 106 | defaults.update(kwargs) |
99 | 107 | return super().get_form(request, obj, **defaults) |
| 108 | + |
| 109 | + |
| 110 | +class GroupWithUsersForm(forms.ModelForm): |
| 111 | + users = forms.ModelMultipleChoiceField( |
| 112 | + queryset=User.objects.all(), |
| 113 | + required=False, |
| 114 | + widget=FilteredSelectMultiple("Users", is_stacked=False), |
| 115 | + label="Users", |
| 116 | + ) |
| 117 | + |
| 118 | + class Meta: |
| 119 | + model = Group |
| 120 | + fields = "__all__" |
| 121 | + |
| 122 | + def __init__(self, *args, **kwargs): |
| 123 | + super().__init__(*args, **kwargs) |
| 124 | + self.fields["users"].label_from_instance = lambda user: ( |
| 125 | + f"{user.username} | {user.email}" if user.email else user.username |
| 126 | + ) |
| 127 | + if self.instance.pk: |
| 128 | + self.fields["users"].initial = self.instance.user_set.all() |
| 129 | + |
| 130 | + def save(self, commit=True): |
| 131 | + group = super().save(commit=commit) |
| 132 | + group.save() |
| 133 | + self.save_m2m() |
| 134 | + group.user_set.set(self.cleaned_data["users"]) |
| 135 | + return group |
| 136 | + |
| 137 | + |
| 138 | +admin.site.unregister(Group) |
| 139 | + |
| 140 | + |
| 141 | +@admin.register(Group) |
| 142 | +class GroupAdmin(admin.ModelAdmin): |
| 143 | + form = GroupWithUsersForm |
| 144 | + search_fields = ("name",) |
| 145 | + ordering = ("name",) |
| 146 | + filter_horizontal = ("permissions",) |
| 147 | + |
| 148 | + def formfield_for_manytomany(self, db_field, request=None, **kwargs): |
| 149 | + if db_field.name == "permissions": |
| 150 | + qs = kwargs.get("queryset", db_field.remote_field.model.objects) |
| 151 | + # Avoid a major performance hit resolving permission names which |
| 152 | + # triggers a content_type load: |
| 153 | + kwargs["queryset"] = qs.select_related("content_type") |
| 154 | + return super().formfield_for_manytomany(db_field, request=request, **kwargs) |
0 commit comments