From 43fc7fb58aebe7f23252c64114d8dd02572056a0 Mon Sep 17 00:00:00 2001
From: Jeremy Stretch <jstretch@netboxlabs.com>
Date: Thu, 2 Oct 2025 16:05:09 -0400
Subject: [PATCH] Add constraints to enforce v1/v2-dependent fields

---
 .../users/migrations/0014_users_token_v2.py   | 25 +++++++++++++++++++
 netbox/users/models/tokens.py                 | 22 ++++++++++++++++
 2 files changed, 47 insertions(+)

diff --git a/netbox/users/migrations/0014_users_token_v2.py b/netbox/users/migrations/0014_users_token_v2.py
index ab0301244..1c0bb5c1d 100644
--- a/netbox/users/migrations/0014_users_token_v2.py
+++ b/netbox/users/migrations/0014_users_token_v2.py
@@ -72,4 +72,29 @@ class Migration(migrations.Migration):
             name='hmac_digest',
             field=models.CharField(blank=True, max_length=64, null=True),
         ),
+
+        # Add constraints to enforce v1/v2-dependent fields
+        migrations.AddConstraint(
+            model_name='token',
+            constraint=models.CheckConstraint(
+                name='enforce_version_dependent_fields',
+                condition=models.Q(
+                    models.Q(
+                        ('hmac_digest__isnull', True),
+                        ('key__isnull', True),
+                        ('pepper_id__isnull', True),
+                        ('plaintext__isnull', False),
+                        ('version', 1)
+                    ),
+                    models.Q(
+                        ('hmac_digest__isnull', False),
+                        ('key__isnull', False),
+                        ('pepper_id__isnull', False),
+                        ('plaintext__isnull', True),
+                        ('version', 2)
+                    ),
+                    _connector='OR'
+                )
+            )
+        ),
     ]
diff --git a/netbox/users/models/tokens.py b/netbox/users/models/tokens.py
index fe112c1c1..3e8e0f108 100644
--- a/netbox/users/models/tokens.py
+++ b/netbox/users/models/tokens.py
@@ -7,6 +7,7 @@ from django.contrib.postgres.fields import ArrayField
 from django.core.exceptions import ValidationError
 from django.core.validators import MinLengthValidator
 from django.db import models
+from django.db.models import Q
 from django.urls import reverse
 from django.utils import timezone
 from django.utils.translation import gettext_lazy as _
@@ -110,6 +111,27 @@ class Token(models.Model):
         ordering = ('-created',)
         verbose_name = _('token')
         verbose_name_plural = _('tokens')
+        constraints = [
+            models.CheckConstraint(
+                name='enforce_version_dependent_fields',
+                condition=(
+                    Q(
+                        version=1,
+                        key__isnull=True,
+                        pepper_id__isnull=True,
+                        hmac_digest__isnull=True,
+                        plaintext__isnull=False
+                    ) |
+                    Q(
+                        version=2,
+                        key__isnull=False,
+                        pepper_id__isnull=False,
+                        hmac_digest__isnull=False,
+                        plaintext__isnull=True
+                    )
+                ),
+            ),
+        ]
 
     def __init__(self, *args, token=None, **kwargs):
         super().__init__(*args, **kwargs)