from __future__ import annotations import hashlib import hmac import secrets from django.conf import settings from django.contrib.auth.models import AbstractUser from django.db import models from django.db.models import Q class User(AbstractUser): class Status(models.TextChoices): ACTIVE = "active", "Active" DISABLED = "disabled", "Disabled" email = models.EmailField("email address", unique=True) payment_user_id = models.CharField( max_length=128, blank=True, help_text="External payment system user identifier, if available.", ) status = models.CharField( max_length=20, choices=Status.choices, default=Status.ACTIVE, ) created_at = models.DateTimeField(auto_now_add=True) @property def is_business_active(self) -> bool: return self.status == self.Status.ACTIVE and self.is_active class Meta: db_table = "user" class UserWallet(models.Model): user = models.OneToOneField( settings.AUTH_USER_MODEL, on_delete=models.PROTECT, related_name="wallet", ) points_balance = models.BigIntegerField(default=0) created_at = models.DateTimeField(auto_now_add=True) updated_at = models.DateTimeField(auto_now=True) class Meta: db_table = "user_wallet" ordering = ("user_id",) constraints = [ models.CheckConstraint( condition=Q(points_balance__gte=0), name="user_wallet_points_balance_non_negative", ), ] def __str__(self) -> str: return f"{self.user} wallet: {self.points_balance}" class ApiKey(models.Model): KEY_PREFIX_LENGTH = 16 class Status(models.TextChoices): ACTIVE = "active", "Active" REVOKED = "revoked", "Revoked" user = models.ForeignKey( settings.AUTH_USER_MODEL, on_delete=models.PROTECT, related_name="api_keys", ) name = models.CharField(max_length=80, blank=True) key_hash = models.CharField(max_length=64, unique=True, editable=False) key_prefix = models.CharField(max_length=32, editable=False) status = models.CharField( max_length=20, choices=Status.choices, default=Status.ACTIVE, ) last_used_at = models.DateTimeField(null=True, blank=True) created_at = models.DateTimeField(auto_now_add=True) updated_at = models.DateTimeField(auto_now=True) class Meta: db_table = "api_key" ordering = ("-created_at", "-id") indexes = [ models.Index(fields=("user", "status")), models.Index(fields=("key_prefix",)), ] def __str__(self) -> str: return f"{self.user} {self.key_prefix}" @classmethod def generate_plaintext_key(cls) -> str: return f"sk_cmhub_{secrets.token_urlsafe(32)}" @classmethod def hash_key(cls, raw_key: str) -> str: return hashlib.sha256(raw_key.encode("utf-8")).hexdigest() @classmethod def build_prefix(cls, raw_key: str) -> str: return raw_key[: cls.KEY_PREFIX_LENGTH] @classmethod def create_for_user(cls, user, *, name: str = "") -> tuple["ApiKey", str]: raw_key = cls.generate_plaintext_key() api_key = cls(user=user, name=name) api_key.set_key(raw_key) api_key.save() return api_key, raw_key def set_key(self, raw_key: str) -> None: self.key_hash = self.hash_key(raw_key) self.key_prefix = self.build_prefix(raw_key) def matches_key(self, raw_key: str) -> bool: return hmac.compare_digest(self.key_hash, self.hash_key(raw_key)) @property def is_active_key(self) -> bool: return self.status == self.Status.ACTIVE