68 lines
1.7 KiB
Python
68 lines
1.7 KiB
Python
import random
|
|
import string
|
|
from datetime import datetime, timedelta
|
|
|
|
from sqlalchemy.orm import Session
|
|
|
|
from app.config import settings
|
|
from app.models.user import User
|
|
from app.models.verification_token import TokenType, VerificationToken
|
|
|
|
|
|
def _generate_code(length: int = 6) -> str:
|
|
return "".join(random.choices(string.digits, k=length))
|
|
|
|
|
|
def _invalidate_tokens(db: Session, user_id: int, token_type: TokenType) -> None:
|
|
db.query(VerificationToken).filter(
|
|
VerificationToken.user_id == user_id,
|
|
VerificationToken.token_type == token_type,
|
|
).delete()
|
|
|
|
|
|
def create_verification_token(
|
|
db: Session, user: User, token_type: TokenType
|
|
) -> VerificationToken:
|
|
_invalidate_tokens(db, user.id, token_type)
|
|
|
|
if token_type == TokenType.activation:
|
|
hours = settings.activation_code_expire_hours
|
|
else:
|
|
hours = settings.password_reset_code_expire_hours
|
|
|
|
token = VerificationToken(
|
|
user_id=user.id,
|
|
code=_generate_code(),
|
|
token_type=token_type,
|
|
expires_at=datetime.utcnow() + timedelta(hours=hours),
|
|
)
|
|
db.add(token)
|
|
db.flush()
|
|
return token
|
|
|
|
|
|
def verify_code(
|
|
db: Session, email: str, code: str, token_type: TokenType
|
|
) -> User | None:
|
|
from app.auth import get_user_by_email
|
|
|
|
user = get_user_by_email(db, email)
|
|
if not user:
|
|
return None
|
|
|
|
token = (
|
|
db.query(VerificationToken)
|
|
.filter(
|
|
VerificationToken.user_id == user.id,
|
|
VerificationToken.token_type == token_type,
|
|
VerificationToken.code == code.strip(),
|
|
VerificationToken.expires_at > datetime.utcnow(),
|
|
)
|
|
.first()
|
|
)
|
|
if not token:
|
|
return None
|
|
|
|
db.delete(token)
|
|
return user
|