| 271 | |
| 272 | |
| 273 | class PackageV2ViewSet(viewsets.ReadOnlyModelViewSet): |
| 274 | queryset = Package.objects.all().prefetch_related( |
| 275 | Prefetch( |
| 276 | "affected_by_vulnerabilities", |
| 277 | queryset=Vulnerability.objects.prefetch_related("fixed_by_packages"), |
| 278 | to_attr="prefetched_affected_vulnerabilities", |
| 279 | ) |
| 280 | ) |
| 281 | serializer_class = PackageV2Serializer |
| 282 | filter_backends = (filters.DjangoFilterBackend,) |
| 283 | filterset_class = PackageV2FilterSet |
| 284 | throttle_classes = [AnonRateThrottle, PermissionBasedUserRateThrottle] |
| 285 | |
| 286 | def get_queryset(self): |
| 287 | queryset = super().get_queryset() |
| 288 | package_purls = self.request.query_params.getlist("purl") |
| 289 | affected_by_vulnerability = self.request.query_params.get("affected_by_vulnerability") |
| 290 | fixing_vulnerability = self.request.query_params.get("fixing_vulnerability") |
| 291 | |
| 292 | if package_purls: |
| 293 | queryset = queryset.filter(package_url__in=package_purls) |
| 294 | if affected_by_vulnerability: |
| 295 | queryset = queryset.filter( |
| 296 | affected_by_vulnerabilities__vulnerability_id=affected_by_vulnerability |
| 297 | ) |
| 298 | if fixing_vulnerability: |
| 299 | queryset = queryset.filter( |
| 300 | fixing_vulnerabilities__vulnerability_id=fixing_vulnerability |
| 301 | ) |
| 302 | return queryset.with_is_vulnerable() |
| 303 | |
| 304 | def list(self, request, *args, **kwargs): |
| 305 | queryset = self.get_queryset() |
| 306 | |
| 307 | # Apply pagination |
| 308 | page = self.paginate_queryset(queryset) |
| 309 | if page is not None: |
| 310 | # Collect only vulnerabilities for packages in the current page |
| 311 | vulnerabilities = set() |
| 312 | for package in page: |
| 313 | vulnerabilities.update(package.affected_by_vulnerabilities.all()) |
| 314 | vulnerabilities.update(package.fixing_vulnerabilities.all()) |
| 315 | |
| 316 | # Serialize the vulnerabilities with vulnerability_id as keys |
| 317 | vulnerability_data = { |
| 318 | vuln.vulnerability_id: VulnerabilityV2Serializer(vuln).data |
| 319 | for vuln in vulnerabilities |
| 320 | } |
| 321 | |
| 322 | # Serialize the current page of packages |
| 323 | serializer = self.get_serializer(page, many=True) |
| 324 | data = serializer.data |
| 325 | |
| 326 | # Use 'self.get_paginated_response' to include pagination data |
| 327 | return self.get_paginated_response( |
| 328 | {"vulnerabilities": vulnerability_data, "packages": data} |
| 329 | ) |
| 330 | |