1import secrets
  2from datetime import datetime
  3from enum import Enum
  4from typing import Annotated
  5
  6from pydantic import (
  7    ConfigDict,
  8    BaseModel,
  9    Field,
 10    StringConstraints,
 11    model_validator,
 12)
 13
 14USERNAME_REGEXP = r"^\w{3,32}$"
 15
 16
 17class ReminderType(str, Enum):
 18    expiration_date = "expiration_date"
 19    data_usage = "data_usage"
 20
 21
 22class UserStatus(str, Enum):
 23    ACTIVE = "active"
 24    INACTIVE = "inactive"
 25
 26
 27class UserDataUsageResetStrategy(str, Enum):
 28    no_reset = "no_reset"
 29    day = "day"
 30    week = "week"
 31    month = "month"
 32    year = "year"
 33
 34
 35class UserExpireStrategy(str, Enum):
 36    NEVER = "never"
 37    FIXED_DATE = "fixed_date"
 38    START_ON_FIRST_USE = "start_on_first_use"
 39
 40
 41class User(BaseModel):
 42    id: int | None = None
 43    username: Annotated[
 44        str, StringConstraints(to_lower=True, pattern=USERNAME_REGEXP)
 45    ]
 46    expire_strategy: UserExpireStrategy
 47    expire_date: datetime | None = Field(None)
 48    usage_duration: int | None = Field(None)
 49    activation_deadline: datetime | None = Field(None)
 50    key: str = Field(default_factory=lambda: secrets.token_hex(16))
 51    data_limit: int | None = Field(
 52        ge=0, default=None, description="data_limit can be 0 or greater"
 53    )
 54    data_limit_reset_strategy: UserDataUsageResetStrategy = (
 55        UserDataUsageResetStrategy.no_reset
 56    )
 57    note: Annotated[str, Field(max_length=500)] | None = None
 58    sub_updated_at: datetime | None = Field(None)
 59    sub_last_user_agent: str | None = Field(None)
 60    online_at: datetime | None = Field(None)
 61
 62    model_config = ConfigDict(from_attributes=True)
 63
 64
 65class UserCreate(User):
 66    service_ids: list[int] = Field([])
 67    model_config = ConfigDict(
 68        json_schema_extra={
 69            "example": {
 70                "username": "user1234",
 71                "services": [1, 2, 3],
 72                "expire_strategy": "start_on_first_use",
 73                "usage_duration": 86400 * 14,
 74                "activation_deadline": "2024-11-03T20:30:00",
 75                "data_limit": 0,
 76                "data_limit_reset_strategy": "no_reset",
 77                "note": "",
 78            }
 79        }
 80    )
 81
 82    @model_validator(mode="after")
 83    def validate_expiry(self):
 84        if self.expire_strategy == UserExpireStrategy.START_ON_FIRST_USE:
 85            if not self.usage_duration:
 86                raise ValueError(
 87                    "User expire_strategy can not be start_on_first_use without a valid usage_duration."
 88                )
 89        elif (
 90            self.expire_strategy == UserExpireStrategy.FIXED_DATE
 91            and not self.expire_date
 92        ):
 93            raise ValueError(
 94                "User expire_strategy can not be fixed_date without a valid expire date."
 95            )
 96        return self
 97
 98
 99class UserModify(UserCreate):
100    service_ids: list[int] | None = Field(None)
101    data_limit_reset_strategy: UserDataUsageResetStrategy | None = Field(None)
102    model_config = ConfigDict(
103        json_schema_extra={
104            "example": {
105                "username": "mammad1234",
106                "service_ids": [1, 2, 3],
107                "expire_strategy": "fixed_date",
108                "expire_date": "2023-11-03T20:30:00",
109                "data_limit": 0,
110                "data_limit_reset_strategy": "no_reset",
111                "note": "",
112            }
113        }
114    )
115
116
117class UserResponse(User):
118    id: int
119    activated: bool
120    is_active: bool
121    expired: bool
122    data_limit_reached: bool
123    enabled: bool
124    used_traffic: int
125    lifetime_used_traffic: int
126    created_at: datetime
127    service_ids: list[int]
128    subscription_url: str
129    owner_username: str
130
131    model_config = ConfigDict(from_attributes=True)
132
133
134class UserUsageResponse(BaseModel):
135    node_id: int | None = None
136    node_name: str
137    used_traffic: int
138
139
140class UserUsagesResponse(BaseModel):
141    username: str
142    usages: list[UserUsageResponse]