gitscope
← marzneshin / v0.3.0-alpha4 / app/models/user.py
https://dev.kronoz.win/r/marzneshin.git
app/models/user.py · Python
Blame Raw
  1import secrets
  2from datetime import datetime
  3from enum import Enum
  4from typing import List, Optional, Union, Annotated, Literal
  5
  6from pydantic import (
  7    field_validator,
  8    ConfigDict,
  9    BaseModel,
 10    Field,
 11    computed_field,
 12    StringConstraints,
 13    model_validator,
 14)
 15
 16from config import XRAY_SUBSCRIPTION_URL_PREFIX
 17
 18USERNAME_REGEXP = r"^\w{3,32}$"
 19
 20
 21class ReminderType(str, Enum):
 22    expiration_date = "expiration_date"
 23    data_usage = "data_usage"
 24
 25
 26class UserStatus(str, Enum):
 27    active = "active"
 28    disabled = "disabled"
 29    limited = "limited"
 30    expired = "expired"
 31    on_hold = "on_hold"
 32
 33
 34class UserStatusModify(str, Enum):
 35    active = "active"
 36    disabled = "disabled"
 37    on_hold = "on_hold"
 38
 39
 40class UserStatusCreate(str, Enum):
 41    active = "active"
 42    on_hold = "on_hold"
 43
 44
 45class UserDataLimitResetStrategy(str, Enum):
 46    no_reset = "no_reset"
 47    day = "day"
 48    week = "week"
 49    month = "month"
 50    year = "year"
 51
 52
 53class UserBase(BaseModel):
 54    id: Optional[int] = None
 55    username: Annotated[str, StringConstraints(to_lower=True, pattern=USERNAME_REGEXP)]
 56    expire: Union[datetime, None, Literal[0]] = Field(None)
 57    key: str = Field(default_factory=lambda: secrets.token_hex(16))
 58    data_limit: Optional[int] = Field(
 59        ge=0, default=None, description="data_limit can be 0 or greater"
 60    )
 61    data_limit_reset_strategy: UserDataLimitResetStrategy = (
 62        UserDataLimitResetStrategy.no_reset
 63    )
 64    note: Optional[Annotated[str, Field(max_length=500)]] = None
 65    sub_updated_at: Optional[datetime] = Field(None)
 66    sub_last_user_agent: Optional[str] = Field(None)
 67    online_at: Optional[datetime] = Field(None)
 68    on_hold_expire_duration: Optional[int] = Field(None)
 69    on_hold_timeout: Optional[datetime] = Field(None)
 70
 71    model_config = ConfigDict(from_attributes=True)
 72
 73
 74from app.models.service import ServiceBase
 75from app.models.proxy import InboundBase
 76
 77
 78class User(UserBase):
 79    pass
 80
 81
 82class UserCreate(User):
 83    status: UserStatusCreate = Field(UserStatusCreate.active)
 84    services: List[int] = Field([])
 85    model_config = ConfigDict(
 86        json_schema_extra={
 87            "example": {
 88                "username": "user1234",
 89                "services": [1, 2, 3],
 90                "expire": "2023-11-03T20:30:00",
 91                "data_limit": 0,
 92                "data_limit_reset_strategy": "no_reset",
 93                "status": "active",
 94                "note": "",
 95                "on_hold_timeout": "2023-11-03T20:30:00",
 96                "on_hold_expire_duration": 0,
 97            }
 98        }
 99    )
100
101    @model_validator(mode="after")
102    def validate_expiry(self):
103        if self.status == UserStatusCreate.on_hold:
104            if not self.on_hold_expire_duration:
105                raise ValueError(
106                    "User cannot be on hold without a valid on_hold_expire_duration."
107                )
108            if self.expire:
109                raise ValueError("User cannot be on hold with specified expire.")
110        elif self.status == UserStatusCreate.active:
111            if not self.expire:
112                raise ValueError("expire field should be set when status is active.")
113            if self.on_hold_expire_duration or self.on_hold_timeout:
114                raise ValueError(
115                    "on hold parametrs set when user status isn't on_hold."
116                )
117        return self
118
119    @field_validator("expire")
120    @classmethod
121    def validate_expire(cls, v: Union[datetime, None, Literal[0]]):
122        if isinstance(v, datetime) and (v.tzinfo is not None):
123            raise ValueError(
124                "Expire date should be offset naive, and preferably in utc timezone."
125            )
126        if isinstance(v, datetime) and v < datetime.utcnow():
127            raise ValueError("Expire date should be in the future, not in the past.")
128        return v
129
130
131class UserModify(User):
132    status: Optional[UserStatusModify] = Field(None)
133    services: List[int] = Field([])
134    data_limit_reset_strategy: Optional[UserDataLimitResetStrategy] = Field(None)
135    model_config = ConfigDict(
136        json_schema_extra={
137            "example": {
138                "username": "mammad1234",
139                "services": [1, 2, 3],
140                "expire": 0,
141                "data_limit": 0,
142                "data_limit_reset_strategy": "no_reset",
143                "status": "active",
144                "note": "",
145                "on_hold_timeout": "2023-11-03T20:30:00",
146                "on_hold_expire_duration": 0,
147            }
148        }
149    )
150
151    @model_validator(mode="after")
152    def validate_expiry(self):
153        if self.on_hold_expire_duration and self.expire:
154            raise ValueError(
155                "can't set both expire and on hold expire at the same time."
156            )
157        return self
158
159    @field_validator("expire")
160    @classmethod
161    def validate_expire(cls, v: Union[datetime, None, Literal[0]]):
162        if isinstance(v, datetime) and v.tzinfo is not None:
163            raise ValueError(
164                "Expire date should be offset naive, and preferably in utc timezone."
165            )
166        if isinstance(v, datetime) and v < datetime.utcnow():
167            raise ValueError("Expire date should be in the future, not in the past.")
168        return v
169
170
171class UserResponse(User):
172    id: int
173    status: UserStatus
174    used_traffic: int
175    lifetime_used_traffic: int = 0
176    created_at: datetime
177    inbounds: List[InboundBase] = Field([])
178    services: List[ServiceBase] = Field([])
179    links: List[str] | None = []  # Field(None)
180
181    # subscription_url: str = ""
182    model_config = ConfigDict(from_attributes=True)
183
184    @computed_field
185    @property
186    def service_ids(self) -> List[int]:
187        return [s.id for s in self.services]
188
189    @computed_field
190    @property
191    def subscription_url(self) -> str:
192        salt = secrets.token_hex(8)
193        url_prefix = XRAY_SUBSCRIPTION_URL_PREFIX.replace("*", salt)
194        return f"{url_prefix}/sub/{self.username}/{self.key}"
195
196
197class UserUsageResponse(BaseModel):
198    node_id: Union[int, None] = None
199    node_name: str
200    used_traffic: int
201
202
203class UserUsagesResponse(BaseModel):
204    username: str
205    usages: List[UserUsageResponse]