| 139 | ) |
| 140 | ) |
| 141 | class VulnerabilityV2ViewSet(viewsets.ReadOnlyModelViewSet): |
| 142 | queryset = Vulnerability.objects.all() |
| 143 | serializer_class = VulnerabilityV2Serializer |
| 144 | lookup_field = "vulnerability_id" |
| 145 | throttle_classes = [AnonRateThrottle, PermissionBasedUserRateThrottle] |
| 146 | |
| 147 | def get_queryset(self): |
| 148 | queryset = super().get_queryset() |
| 149 | vulnerability_ids = self.request.query_params.getlist("vulnerability_id") |
| 150 | aliases = self.request.query_params.getlist("alias") |
| 151 | |
| 152 | if vulnerability_ids: |
| 153 | queryset = queryset.filter(vulnerability_id__in=vulnerability_ids) |
| 154 | |
| 155 | if aliases: |
| 156 | queryset = queryset.filter(aliases__alias__in=aliases).distinct() |
| 157 | |
| 158 | return queryset |
| 159 | |
| 160 | def get_serializer_class(self): |
| 161 | if self.action == "list": |
| 162 | return VulnerabilityListSerializer |
| 163 | return super().get_serializer_class() |
| 164 | |
| 165 | def list(self, request, *args, **kwargs): |
| 166 | queryset = self.get_queryset() |
| 167 | vulnerability_ids = request.query_params.getlist("vulnerability_id") |
| 168 | |
| 169 | # If exactly one vulnerability_id is provided, return the serialized data |
| 170 | if len(vulnerability_ids) == 1: |
| 171 | try: |
| 172 | vulnerability = queryset.get(vulnerability_id=vulnerability_ids[0]) |
| 173 | serializer = self.get_serializer(vulnerability) |
| 174 | return Response(serializer.data) |
| 175 | except Vulnerability.DoesNotExist: |
| 176 | return Response({"detail": "Not found."}, status=404) |
| 177 | |
| 178 | # Otherwise, return a dictionary of vulnerabilities keyed by vulnerability_id |
| 179 | page = self.paginate_queryset(queryset) |
| 180 | if page is not None: |
| 181 | serializer = self.get_serializer(page, many=True) |
| 182 | data = serializer.data |
| 183 | vulnerabilities = {item["vulnerability_id"]: item for item in data} |
| 184 | return self.get_paginated_response({"vulnerabilities": vulnerabilities}) |
| 185 | |
| 186 | serializer = self.get_serializer(queryset, many=True) |
| 187 | data = serializer.data |
| 188 | vulnerabilities = {item["vulnerability_id"]: item for item in data} |
| 189 | return Response({"vulnerabilities": vulnerabilities}) |
| 190 | |
| 191 | |
| 192 | class PackageV2Serializer(serializers.ModelSerializer): |