From 1e6d8dfb18f3ae15ff9f21a238ab40dc971b7739 Mon Sep 17 00:00:00 2001
From: Daniel Sheppard <dans@dansheps.com>
Date: Thu, 22 Aug 2024 10:37:02 -0500
Subject: [PATCH] Fixes: #16292 - Properly restrict GraphQL queries for querys
 with pk set

---
 netbox/ipam/graphql/schema.py | 66 +++++++++--------------------------
 netbox/netbox/settings.py     |  1 +
 2 files changed, 18 insertions(+), 49 deletions(-)

diff --git a/netbox/ipam/graphql/schema.py b/netbox/ipam/graphql/schema.py
index c02788c3a..d11eadf8a 100644
--- a/netbox/ipam/graphql/schema.py
+++ b/netbox/ipam/graphql/schema.py
@@ -7,84 +7,52 @@ from ipam import models
 from .types import *
 
 
-@strawberry.type
+@strawberry.type(name="Query")
 class IPAMQuery:
-    @strawberry.field
-    def asn(self, id: int) -> ASNType:
-        return models.ASN.objects.get(pk=id)
+    asn: ASNType = strawberry_django.field()
     asn_list: List[ASNType] = strawberry_django.field()
 
-    @strawberry.field
-    def asn_range(self, id: int) -> ASNRangeType:
-        return models.ASNRange.objects.get(pk=id)
+    asn_range: ASNRangeType = strawberry_django.field()
     asn_range_list: List[ASNRangeType] = strawberry_django.field()
 
-    @strawberry.field
-    def aggregate(self, id: int) -> AggregateType:
-        return models.Aggregate.objects.get(pk=id)
+    aggregate: AggregateType = strawberry_django.field()
     aggregate_list: List[AggregateType] = strawberry_django.field()
 
-    @strawberry.field
-    def ip_address(self, id: int) -> IPAddressType:
-        return models.IPAddress.objects.get(pk=id)
+    ip_address: IPAddressType = strawberry_django.field()
     ip_address_list: List[IPAddressType] = strawberry_django.field()
 
-    @strawberry.field
-    def ip_range(self, id: int) -> IPRangeType:
-        return models.IPRange.objects.get(pk=id)
+    ip_range: IPRangeType = strawberry_django.field()
     ip_range_list: List[IPRangeType] = strawberry_django.field()
 
-    @strawberry.field
-    def prefix(self, id: int) -> PrefixType:
-        return models.Prefix.objects.get(pk=id)
+    prefix: PrefixType = strawberry_django.field()
     prefix_list: List[PrefixType] = strawberry_django.field()
 
-    @strawberry.field
-    def rir(self, id: int) -> RIRType:
-        return models.RIR.objects.get(pk=id)
+    rir: RIRType = strawberry_django.field()
     rir_list: List[RIRType] = strawberry_django.field()
 
-    @strawberry.field
-    def role(self, id: int) -> RoleType:
-        return models.Role.objects.get(pk=id)
+    role: RoleType = strawberry_django.field()
     role_list: List[RoleType] = strawberry_django.field()
 
-    @strawberry.field
-    def route_target(self, id: int) -> RouteTargetType:
-        return models.RouteTarget.objects.get(pk=id)
+    route_target: RouteTargetType = strawberry_django.field()
     route_target_list: List[RouteTargetType] = strawberry_django.field()
 
-    @strawberry.field
-    def service(self, id: int) -> ServiceType:
-        return models.Service.objects.get(pk=id)
+    service: ServiceType = strawberry_django.field()
     service_list: List[ServiceType] = strawberry_django.field()
 
-    @strawberry.field
-    def service_template(self, id: int) -> ServiceTemplateType:
-        return models.ServiceTemplate.objects.get(pk=id)
+    service_template: ServiceTemplateType = strawberry_django.field()
     service_template_list: List[ServiceTemplateType] = strawberry_django.field()
 
-    @strawberry.field
-    def fhrp_group(self, id: int) -> FHRPGroupType:
-        return models.FHRPGroup.objects.get(pk=id)
+    fhrp_group: FHRPGroupType = strawberry_django.field()
     fhrp_group_list: List[FHRPGroupType] = strawberry_django.field()
 
-    @strawberry.field
-    def fhrp_group_assignment(self, id: int) -> FHRPGroupAssignmentType:
-        return models.FHRPGroupAssignment.objects.get(pk=id)
+    fhrp_group_assignment: FHRPGroupAssignmentType = strawberry_django.field()
     fhrp_group_assignment_list: List[FHRPGroupAssignmentType] = strawberry_django.field()
 
-    @strawberry.field
-    def vlan(self, id: int) -> VLANType:
-        return models.VLAN.objects.get(pk=id)
+    vlan: VLANType = strawberry_django.field()
     vlan_list: List[VLANType] = strawberry_django.field()
 
-    @strawberry.field
-    def vlan_group(self, id: int) -> VLANGroupType:
-        return models.VLANGroup.objects.get(pk=id)
+    vlan_group: VLANGroupType = strawberry_django.field()
     vlan_group_list: List[VLANGroupType] = strawberry_django.field()
 
-    @strawberry.field
-    def vrf(self, id: int) -> VRFType:
-        return models.VRF.objects.get(pk=id)
+    vrf: VRFType = strawberry_django.field()
     vrf_list: List[VRFType] = strawberry_django.field()
diff --git a/netbox/netbox/settings.py b/netbox/netbox/settings.py
index 869b6be31..fa05e9929 100644
--- a/netbox/netbox/settings.py
+++ b/netbox/netbox/settings.py
@@ -764,6 +764,7 @@ LOCALE_PATHS = (
 #
 STRAWBERRY_DJANGO = {
     "TYPE_DESCRIPTION_FROM_MODEL_DOCSTRING": True,
+    "DEFAULT_PK_FIELD_NAME": "id",
     "USE_DEPRECATED_FILTERS": True,
 }