from __future__ import annotations from datetime import datetime, timezone from typing import List, Optional import pytz from flask_sqlalchemy import SQLAlchemy from sqlalchemy import Column, DateTime, String, TypeDecorator from sqlalchemy.orm import Mapped, mapped_column, relationship db : SQLAlchemy = SQLAlchemy() deck_password = db.Table( "deck_password", Column("deck_id", db.Integer, db.ForeignKey("deck.id"), primary_key=True), Column("password_id", db.Integer, db.ForeignKey("password.id"), primary_key=True) ) deck_token = db.Table( "deck_token", Column("deck_id", db.Integer, db.ForeignKey("deck.id"), primary_key=True), Column("token_id", db.Integer, db.ForeignKey("token.id"), primary_key=True) ) class UTCDateTime(TypeDecorator): impl = DateTime def process_bind_param(self, value, dialect): # Ensure datetime is stored in naive UTC if value is not None and value.tzinfo is not None: value = value.astimezone(pytz.utc).replace(tzinfo=None) return value def process_result_value(self, value, dialect): # Re-attach UTC tzinfo when reading if value is not None: value = value.replace(tzinfo=pytz.utc) return value class Password(db.Model): __tablename__ = "password" id : Mapped[int] = mapped_column(primary_key=True) value : Mapped[str] = mapped_column(String(60)) expires_at: Mapped[Optional[datetime]] = mapped_column(UTCDateTime()) # Optional expiration decks : Mapped[List[Deck]] = relationship("Deck", secondary=deck_password, back_populates='passwords') def __init__(self, value: str, expires_at: datetime|None) -> None: self.value = value self.expires_at = expires_at def is_expired(self) -> bool: """Check if the password is expired""" return self.expires_at is not None and datetime.now(timezone.utc) > self.expires_at def __repr__(self) -> str: return f"" class Deck(db.Model): __tablename__ = "deck" id : Mapped[int] = mapped_column(primary_key=True) name : Mapped[str] = mapped_column(String(60), unique=True, nullable=False) deck : Mapped[str] = mapped_column(String(60), unique=True, nullable=False) description : Mapped[str] = mapped_column(String(60)) passwords : Mapped[List[Password]] = relationship("Password", secondary=deck_password, back_populates='decks') tokens : Mapped[List[Token]] = relationship("Token", secondary=deck_token, back_populates='decks') def __init__(self, name: str, deck: str, description: str = "") -> None: self.name = name self.deck = deck self.description = description def __repr__(self) -> str: return f"" class Token(db.Model): __tablename__ = "token" id : Mapped[int] = mapped_column(primary_key=True) name: Mapped[str] = mapped_column(String(60), unique=True, nullable=False) decks : Mapped[List[Deck]]= relationship("Deck", secondary=deck_token, back_populates='tokens') def __init__(self, name: str) -> None: self.name = name def __repr__(self) -> str: return f"" def is_unlocked(self, deck: str) -> bool : return deck in [deck.deck for deck in self.decks]