1import secrets
  2from datetime import datetime
  3from enum import StrEnum
  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 UserStatus(StrEnum):
 18    ACTIVE = "active"
 19    INACTIVE = "inactive"
 20
 21
 22class UserDataUsageResetStrategy(StrEnum):
 23    no_reset = "no_reset"
 24    day = "day"
 25    week = "week"
 26    month = "month"
 27    year = "year"
 28
 29
 30class UserExpireStrategy(StrEnum):
 31    NEVER = "never"
 32    FIXED_DATE = "fixed_date"
 33    START_ON_FIRST_USE = "start_on_first_use"
 34
 35
 36class User(BaseModel):
 37    id: int | None = None
 38    username: Annotated[
 39        str, StringConstraints(to_lower=True, pattern=USERNAME_REGEXP)
 40    ]
 41    expire_strategy: UserExpireStrategy
 42    expire_date: datetime | None = Field(None)
 43    usage_duration: (
 44        Annotated[
 45            int,
 46            Field(
 47                le=9999 * 365 * 24 * 60 * 60,
 48                description="usage_duration must be less than 9999 years",
 49            ),
 50        ]
 51        | None
 52    ) = (None,)
 53    activation_deadline: datetime | None = Field(None)
 54    key: str = Field(default_factory=lambda: secrets.token_hex(16))
 55    data_limit: int | None = Field(
 56        ge=0, default=None, description="data_limit can be 0 or greater"
 57    )
 58    data_limit_reset_strategy: UserDataUsageResetStrategy = (
 59        UserDataUsageResetStrategy.no_reset
 60    )
 61    note: Annotated[str, Field(max_length=500)] | None = None
 62    sub_updated_at: datetime | None = Field(None)
 63    sub_last_user_agent: str | None = Field(None)
 64    online_at: datetime | None = Field(None)
 65
 66    model_config = ConfigDict(from_attributes=True)
 67
 68    @model_validator(mode="after")
 69    def validate_expiry(self):
 70        if self.expire_strategy == UserExpireStrategy.START_ON_FIRST_USE:
 71            if not self.usage_duration:
 72                raise ValueError(
 73                    "User expire_strategy cannot be start_on_first_use without a valid usage_duration."
 74                )
 75
 76            # Set expire_date to None if expire_strategy is START_ON_FIRST_USE
 77            self.expire_date = None
 78
 79        elif self.expire_strategy == UserExpireStrategy.FIXED_DATE:
 80            if not self.expire_date:
 81                raise ValueError(
 82                    "User expire_strategy cannot be fixed_date without a valid expire date."
 83                )
 84
 85            # Set usage_duration and activation_deadline to None if expire_strategy is START_ON_FIRST_USE
 86            self.usage_duration = None
 87            self.activation_deadline = None
 88
 89        elif self.expire_strategy == UserExpireStrategy.NEVER:
 90            # Set expire_date, usage_duration and activation_deadline to None if expire_strategy is NEVER
 91            self.expire_date = None
 92            self.usage_duration = None
 93            self.activation_deadline = None
 94
 95        return self
 96
 97
 98class UserCreate(User):
 99    service_ids: list[int] = Field([])
100    model_config = ConfigDict(
101        json_schema_extra={
102            "example": {
103                "username": "user1234",
104                "service_ids": [1, 2, 3],
105                "expire_strategy": "start_on_first_use",
106                "usage_duration": 86400 * 14,
107                "activation_deadline": "2024-11-03T20:30:00",
108                "data_limit": 0,
109                "data_limit_reset_strategy": "no_reset",
110                "note": "",
111            }
112        }
113    )
114
115
116class UserModify(UserCreate):
117    service_ids: list[int] | None = Field(None)
118    data_limit_reset_strategy: UserDataUsageResetStrategy | None = Field(None)
119    expire_strategy: UserExpireStrategy | None = Field(None)
120    model_config = ConfigDict(
121        json_schema_extra={
122            "example": {
123                "username": "mammad1234",
124                "service_ids": [1, 2, 3],
125                "expire_strategy": "fixed_date",
126                "expire_date": "2023-11-03T20:30:00",
127                "data_limit": 0,
128                "data_limit_reset_strategy": "no_reset",
129                "note": "",
130            }
131        }
132    )
133
134
135class UserResponse(User):
136    id: int
137    activated: bool
138    is_active: bool
139    expired: bool
140    data_limit_reached: bool
141    enabled: bool
142    used_traffic: int
143    lifetime_used_traffic: int
144    sub_revoked_at: datetime | None
145    created_at: datetime
146    service_ids: list[int]
147    subscription_url: str
148    owner_username: str | None
149    traffic_reset_at: datetime | None
150
151    model_config = ConfigDict(from_attributes=True)
152
153
154class UserNodeUsageSeries(BaseModel):
155    node_id: int | None = None
156    node_name: str
157    usages: list[tuple[int, int]]  # list[UserUsage]
158
159
160class UserUsageSeriesResponse(BaseModel):
161    username: str
162    node_usages: list[UserNodeUsageSeries]
163    total: int