139 lines
4.2 KiB
Python
139 lines
4.2 KiB
Python
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", "正常"
|
|
DISABLED = "disabled", "禁用"
|
|
|
|
email = models.EmailField("邮箱", unique=True)
|
|
payment_user_id = models.CharField(
|
|
"支付系统用户 ID",
|
|
max_length=128,
|
|
blank=True,
|
|
help_text="外部支付系统用户标识,可为空。",
|
|
)
|
|
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"
|
|
verbose_name = "用户"
|
|
verbose_name_plural = "用户"
|
|
|
|
|
|
class UserWallet(models.Model):
|
|
user = models.OneToOneField(
|
|
settings.AUTH_USER_MODEL,
|
|
verbose_name="用户",
|
|
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"
|
|
verbose_name = "用户钱包"
|
|
verbose_name_plural = "用户钱包"
|
|
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", "有效"
|
|
REVOKED = "revoked", "已吊销"
|
|
|
|
user = models.ForeignKey(
|
|
settings.AUTH_USER_MODEL,
|
|
verbose_name="用户",
|
|
on_delete=models.PROTECT,
|
|
related_name="api_keys",
|
|
)
|
|
name = models.CharField("名称", max_length=80, blank=True)
|
|
key_hash = models.CharField("API Key 哈希", max_length=64, unique=True, editable=False)
|
|
key_prefix = models.CharField("Key 前缀", 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"
|
|
verbose_name = "API 密钥"
|
|
verbose_name_plural = "API 密钥"
|
|
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
|