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    @model_validator(mode="after")
 65    def validate_expiry(self):
 66        if self.expire_strategy == UserExpireStrategy.START_ON_FIRST_USE:
 67            if not self.usage_duration:
 68                raise ValueError(
 69                    "User expire_strategy cannot be start_on_first_use without a valid usage_duration."
 70                )
 71
 72            # Set expire_date to None if expire_strategy is START_ON_FIRST_USE
 73            self.expire_date = None
 74
 75        elif self.expire_strategy == UserExpireStrategy.FIXED_DATE:
 76            if not self.expire_date:
 77                raise ValueError(
 78                    "User expire_strategy cannot be fixed_date without a valid expire date."
 79                )
 80
 81            # Set usage_duration and activation_deadline to None if expire_strategy is START_ON_FIRST_USE
 82            self.usage_duration = None
 83            self.activation_deadline = None
 84
 85        elif self.expire_strategy == UserExpireStrategy.NEVER:
 86            # Set expire_date, usage_duration and activation_deadline to None if expire_strategy is NEVER
 87            self.expire_date = None
 88            self.usage_duration = None
 89            self.activation_deadline = None
 90
 91        return self
 92
 93
 94class UserCreate(User):
 95    service_ids: list[int] = Field([])
 96    model_config = ConfigDict(
 97        json_schema_extra={
 98            "example": {
 99                "username": "user1234",
100                "service_ids": [1, 2, 3],
101                "expire_strategy": "start_on_first_use",
102                "usage_duration": 86400 * 14,
103                "activation_deadline": "2024-11-03T20:30:00",
104                "data_limit": 0,
105                "data_limit_reset_strategy": "no_reset",
106                "note": "",
107            }
108        }
109    )
110
111
112class UserModify(UserCreate):
113    service_ids: list[int] | None = Field(None)
114    data_limit_reset_strategy: UserDataUsageResetStrategy | None = Field(None)
115    model_config = ConfigDict(
116        json_schema_extra={
117            "example": {
118                "username": "mammad1234",
119                "service_ids": [1, 2, 3],
120                "expire_strategy": "fixed_date",
121                "expire_date": "2023-11-03T20:30:00",
122                "data_limit": 0,
123                "data_limit_reset_strategy": "no_reset",
124                "note": "",
125            }
126        }
127    )
128
129
130class UserResponse(User):
131    id: int
132    activated: bool
133    is_active: bool
134    expired: bool
135    data_limit_reached: bool
136    enabled: bool
137    used_traffic: int
138    lifetime_used_traffic: int
139    sub_revoked_at: datetime | None
140    created_at: datetime
141    service_ids: list[int]
142    subscription_url: str
143    owner_username: str | None
144    traffic_reset_at: datetime | None
145
146    model_config = ConfigDict(from_attributes=True)
147
148
149class UserNodeUsageSeries(BaseModel):
150    node_id: int | None = None
151    node_name: str
152    usages: list[tuple[int, int]]  # list[UserUsage]
153
154
155class UserUsageSeriesResponse(BaseModel):
156    username: str
157    node_usages: list[UserNodeUsageSeries]
158    total: int