88 lines
4.5 KiB
Python
88 lines
4.5 KiB
Python
from __future__ import annotations
|
|
|
|
import re
|
|
from datetime import datetime
|
|
from enum import StrEnum
|
|
|
|
from django.db import models, transaction
|
|
from django.utils import timezone
|
|
|
|
_DIGEST = re.compile(r"^[0-9a-f]{64}$")
|
|
|
|
|
|
class ChallengeOwnershipConflict(RuntimeError):
|
|
pass
|
|
|
|
|
|
class ChallengeExpired(RuntimeError):
|
|
pass
|
|
|
|
|
|
class ConsumeOutcome(StrEnum):
|
|
CONSUMED = "consumed"
|
|
ALREADY_CONSUMED = "already-consumed"
|
|
|
|
|
|
def validate_digest(name: str, value: str) -> str:
|
|
if not _DIGEST.fullmatch(value):
|
|
raise ValueError(f"{name} must be a lowercase SHA-256 hex digest")
|
|
return value
|
|
|
|
|
|
def validate_owner(*, account_digest: str, sender_digest: str, portal_user_digest: str, portal_session_digest: str) -> tuple[str, str, str, str]:
|
|
return (validate_digest("account_digest", account_digest), validate_digest("sender_digest", sender_digest), validate_digest("portal_user_digest", portal_user_digest), validate_digest("portal_session_digest", portal_session_digest))
|
|
|
|
|
|
class PersonalWeixinChallengeManager(models.Manager):
|
|
def create_or_get_owned(self, *, challenge_digest: str, account_digest: str, sender_digest: str, portal_user_digest: str, portal_session_digest: str, expires_at: datetime) -> tuple["PersonalWeixinChallenge", bool]:
|
|
challenge_digest = validate_digest("challenge_digest", challenge_digest)
|
|
owner = validate_owner(account_digest=account_digest, sender_digest=sender_digest, portal_user_digest=portal_user_digest, portal_session_digest=portal_session_digest)
|
|
if expires_at <= timezone.now():
|
|
raise ValueError("expires_at must be in the future")
|
|
defaults = {"account_digest": owner[0], "sender_digest": owner[1], "portal_user_digest": owner[2], "portal_session_digest": owner[3], "expires_at": expires_at}
|
|
with transaction.atomic():
|
|
record, created = self.get_or_create(challenge_digest=challenge_digest, defaults=defaults)
|
|
if record.owner_tuple != owner:
|
|
raise ChallengeOwnershipConflict("challenge digest belongs to a different owner")
|
|
return record, created
|
|
|
|
def consume_owned(self, *, challenge_digest: str, account_digest: str, sender_digest: str, portal_user_digest: str, portal_session_digest: str, consumed_at: datetime | None = None) -> ConsumeOutcome:
|
|
challenge_digest = validate_digest("challenge_digest", challenge_digest)
|
|
owner = validate_owner(account_digest=account_digest, sender_digest=sender_digest, portal_user_digest=portal_user_digest, portal_session_digest=portal_session_digest)
|
|
consumed_at = consumed_at or timezone.now()
|
|
with transaction.atomic():
|
|
record = self.select_for_update().get(challenge_digest=challenge_digest)
|
|
if record.owner_tuple != owner:
|
|
raise ChallengeOwnershipConflict("challenge digest belongs to a different owner")
|
|
if consumed_at >= record.expires_at:
|
|
raise ChallengeExpired("challenge has expired")
|
|
if record.consumed_at is not None:
|
|
return ConsumeOutcome.ALREADY_CONSUMED
|
|
updated = self.filter(pk=record.pk, consumed_at__isnull=True).update(consumed_at=consumed_at)
|
|
if updated == 1:
|
|
return ConsumeOutcome.CONSUMED
|
|
record.refresh_from_db(fields=("consumed_at",))
|
|
if record.consumed_at is not None:
|
|
return ConsumeOutcome.ALREADY_CONSUMED
|
|
raise RuntimeError("atomic challenge consumption did not converge")
|
|
|
|
|
|
class PersonalWeixinChallenge(models.Model):
|
|
challenge_digest = models.CharField(max_length=64, unique=True)
|
|
account_digest = models.CharField(max_length=64)
|
|
sender_digest = models.CharField(max_length=64)
|
|
portal_user_digest = models.CharField(max_length=64)
|
|
portal_session_digest = models.CharField(max_length=64)
|
|
created_at = models.DateTimeField(auto_now_add=True)
|
|
expires_at = models.DateTimeField()
|
|
consumed_at = models.DateTimeField(null=True, blank=True)
|
|
|
|
objects = PersonalWeixinChallengeManager()
|
|
|
|
class Meta:
|
|
app_label = "personal_weixin"
|
|
indexes = [models.Index(fields=["account_digest", "sender_digest", "portal_user_digest", "portal_session_digest"], name="pwx_owner_lookup_idx"), models.Index(fields=["expires_at"], name="pwx_expiry_idx"), models.Index(fields=["consumed_at"], name="pwx_consumed_idx")]
|
|
|
|
@property
|
|
def owner_tuple(self) -> tuple[str, str, str, str]:
|
|
return (self.account_digest, self.sender_digest, self.portal_user_digest, self.portal_session_digest) |