Merge pull request #34 from RandomProgramm3r/develop

refactor(business auth): Consolidate JWT token utilities
- Simplify views using `CreateAPIView` and remove redundant mixins
- Add `@transaction.atomic` for atomic database operations during company creation
- Extract JWT token generation into reusable `generate_company_tokens()` utility
- Extract JWT token version bump into reusable `bump_company_token_version()` utility
- Validate UUID format rigorously in token refresh flow
- Add documentation in views
- Remove redundant imports
- Add and update test

Author: RandomProgramm3r

promo_code/business/serializers.py

@@ -1,8 +1,8 @@
 import uuid
 
 import django.contrib.auth.password_validation
-import django.core.exceptions
 import django.core.validators
+import django.db.transaction
 import pycountry
 import rest_framework.exceptions
 import rest_framework.serializers
@@ -11,7 +11,10 @@
 import rest_framework_simplejwt.tokens
 
 import business.constants
+import business.models
 import business.models as business_models
+import business.utils.auth
+import business.utils.tokens
 import business.validators
 
 
@@ -21,9 +24,9 @@ class CompanySignUpSerializer(rest_framework.serializers.ModelSerializer):
         write_only=True,
         required=True,
         validators=[django.contrib.auth.password_validation.validate_password],
+        style={'input_type': 'password'},
         min_length=business.constants.COMPANY_PASSWORD_MIN_LENGTH,
         max_length=business.constants.COMPANY_PASSWORD_MAX_LENGTH,
-        style={'input_type': 'password'},
     )
     name = rest_framework.serializers.CharField(
         required=True,
@@ -44,30 +47,18 @@ class CompanySignUpSerializer(rest_framework.serializers.ModelSerializer):
 
     class Meta:
         model = business_models.Company
-        fields = (
-            'id',
-            'name',
-            'email',
-            'password',
-        )
+        fields = ('id', 'name', 'email', 'password')
 
+    @django.db.transaction.atomic
     def create(self, validated_data):
-        try:
-            company = business_models.Company.objects.create_company(
-                email=validated_data['email'],
-                name=validated_data['name'],
-                password=validated_data['password'],
-            )
-            company.token_version += 1
-            company.save()
-            return company
-        except django.core.exceptions.ValidationError as e:
-            raise rest_framework.serializers.ValidationError(e.messages)
+        company = business_models.Company.objects.create_company(
+            **validated_data,
+        )
 
+        return business.utils.auth.bump_company_token_version(company)
 
-class CompanySignInSerializer(
-    rest_framework.serializers.Serializer,
-):
+
+class CompanySignInSerializer(rest_framework.serializers.Serializer):
     email = rest_framework.serializers.EmailField(required=True)
     password = rest_framework.serializers.CharField(
         required=True,
@@ -80,16 +71,15 @@ def validate(self, attrs):
         password = attrs.get('password')
 
         if not email or not password:
-            raise rest_framework.exceptions.ValidationError(
-                {'detail': 'Both email and password are required'},
-                code='required',
+            raise rest_framework.serializers.ValidationError(
+                'Both email and password are required.',
             )
 
         try:
             company = business_models.Company.objects.get(email=email)
         except business_models.Company.DoesNotExist:
             raise rest_framework.serializers.ValidationError(
-                'Invalid credentials',
+                'Invalid credentials.',
             )
 
         if not company.is_active or not company.check_password(password):
@@ -98,53 +88,55 @@ def validate(self, attrs):
                 code='authentication_failed',
             )
 
+        attrs['company'] = company
         return attrs
 
 
 class CompanyTokenRefreshSerializer(
     rest_framework_simplejwt.serializers.TokenRefreshSerializer,
 ):
     def validate(self, attrs):
+        attrs = super().validate(attrs)
         refresh = rest_framework_simplejwt.tokens.RefreshToken(
             attrs['refresh'],
         )
-        user_type = refresh.payload.get('user_type', 'user')
+        company = self.get_active_company_from_token(refresh)
+
+        company = business.utils.auth.bump_company_token_version(company)
 
-        if user_type != 'company':
+        return business.utils.tokens.generate_company_tokens(company)
+
+    def get_active_company_from_token(self, token):
+        if token.payload.get('user_type') != 'company':
             raise rest_framework_simplejwt.exceptions.InvalidToken(
                 'This refresh endpoint is for company tokens only',
             )
 
-        company_id = refresh.payload.get('company_id')
-        if not company_id:
+        company_id = token.payload.get('company_id')
+        try:
+            company_uuid = uuid.UUID(company_id)
+        except (TypeError, ValueError):
             raise rest_framework_simplejwt.exceptions.InvalidToken(
-                'Company ID missing in token',
+                'Invalid or missing company_id in token',
             )
 
         try:
-            company = business_models.Company.objects.get(
-                id=uuid.UUID(company_id),
+            company = business.models.Company.objects.get(
+                id=company_uuid,
+                is_active=True,
             )
-        except business_models.Company.DoesNotExist:
+        except business.models.Company.DoesNotExist:
             raise rest_framework_simplejwt.exceptions.InvalidToken(
-                'Company not found',
+                'Company not found or inactive',
             )
 
-        token_version = refresh.payload.get('token_version', 0)
+        token_version = token.payload.get('token_version', 0)
         if company.token_version != token_version:
             raise rest_framework_simplejwt.exceptions.InvalidToken(
                 'Token is blacklisted',
             )
 
-        new_refresh = rest_framework_simplejwt.tokens.RefreshToken()
-        new_refresh['user_type'] = 'company'
-        new_refresh['company_id'] = str(company.id)
-        new_refresh['token_version'] = company.token_version
-
-        return {
-            'access': str(new_refresh.access_token),
-            'refresh': str(new_refresh),
-        }
+        return company
 
 
 class TargetSerializer(rest_framework.serializers.Serializer):

promo_code/business/tests/auth/test_tokens.py

@@ -187,7 +187,7 @@ def test_missing_company_id(self):
             rest_framework.status.HTTP_401_UNAUTHORIZED,
         )
         self.assertIn(
-            'Company ID missing in token',
+            'Invalid or missing company_id in token',
             str(response.content.decode()),
         )
 

promo_code/business/tests/promocodes/test_permissions.py

@@ -11,29 +11,50 @@
 class TestIsCompanyUserPermission(
     business.tests.promocodes.base.BasePromoTestCase,
 ):
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+
+        cls.unique_payload = {
+            'description': 'Complimentary Pudge Skin on Registration!',
+            'target': {},
+            'max_count': 1,
+            'mode': 'UNIQUE',
+            'active_from': '2030-08-08',
+            'promo_unique': ['dota-arena', 'coda-core', 'warcraft3'],
+        }
+
     def setUp(self):
         self.factory = rest_framework.test.APIRequestFactory()
         self.permission = business.permissions.IsCompanyUser()
         get_user_model = django.contrib.auth.get_user_model
         self.regular_user = get_user_model().objects.create_user(
             name='regular',
-            password='testpass123',
+            password='SecurePass123!',
             surname='adadioa',
             email='[email protected]',
         )
-        self.company_user = business.models.Company.objects.create_company(
-            password='testpass123',
-            name='Test Company',
-            email='[email protected]',
+
+    def create_promo(self, token, payload):
+        self.client.credentials(HTTP_AUTHORIZATION='Bearer ' + token)
+        response = self.client.post(
+            self.promo_create_url,
+            payload,
+            format='json',
+        )
+        self.assertEqual(
+            response.status_code,
+            rest_framework.status.HTTP_201_CREATED,
         )
+        return response.data['id']
 
     def tearDown(self):
         business.models.Company.objects.all().delete()
         user.models.User.objects.all().delete()
 
     def test_has_permission_for_company_user(self):
         request = self.factory.get(self.promo_create_url)
-        request.user = self.company_user
+        request.user = self.company1
         self.assertTrue(self.permission.has_permission(request, None))
 
     def test_has_permission_for_regular_user(self):
@@ -45,3 +66,16 @@ def test_has_permission_for_anonymous_user(self):
         request = self.factory.get(self.promo_create_url)
         request.user = None
         self.assertFalse(self.permission.has_permission(request, None))
+
+    def test_has_permission_to_foreign_promo(self):
+        promo_id = self.create_promo(self.company2_token, self.unique_payload)
+        self.client.credentials(
+            HTTP_AUTHORIZATION='Bearer ' + self.company1_token,
+        )
+        url = self.promo_detail_url(promo_id)
+        patch_payload = {'description': '100% Cashback'}
+        response = self.client.patch(url, patch_payload, format='json')
+        self.assertEqual(
+            response.status_code,
+            rest_framework.status.HTTP_403_FORBIDDEN,
+        )

promo_code/business/utils/auth.py

@@ -0,0 +1,13 @@
+import business.models
+
+
+def bump_company_token_version(company):
+    """
+    Increment token_version, save it, and return the fresh instance.
+    """
+    company = business.models.Company.objects.select_for_update().get(
+        id=company.id,
+    )
+    company.token_version += 1
+    company.save(update_fields=['token_version'])
+    return company

promo_code/business/utils/tokens.py

@@ -0,0 +1,20 @@
+import rest_framework_simplejwt.tokens
+
+
+def generate_company_tokens(company):
+    """
+    Generate JWT tokens for a company.
+    """
+    refresh = rest_framework_simplejwt.tokens.RefreshToken()
+    refresh['user_type'] = 'company'
+    refresh['company_id'] = str(company.id)
+    refresh['token_version'] = company.token_version
+
+    access = refresh.access_token
+    access['user_type'] = 'company'
+    access['company_id'] = str(company.id)
+
+    return {
+        'access': str(access),
+        'refresh': str(refresh),
+    }