1import re
  2import secrets
  3from datetime import datetime
  4from enum import Enum
  5from typing import List, Optional, Union, Annotated
  6
  7from pydantic import field_validator, ConfigDict, BaseModel, Field, validator, computed_field, ValidationInfo, constr, \
  8    StringConstraints
  9
 10from config import XRAY_SUBSCRIPTION_URL_PREFIX
 11
 12USERNAME_REGEXP = re.compile(r"^\w{3,32}$")
 13
 14
 15class ReminderType(str, Enum):
 16    expiration_date = "expiration_date"
 17    data_usage = "data_usage"
 18
 19
 20class UserStatus(str, Enum):
 21    active = "active"
 22    disabled = "disabled"
 23    limited = "limited"
 24    expired = "expired"
 25    on_hold = "on_hold"
 26
 27
 28# from app.utils.share import generate_v2ray_links
 29
 30
 31class UserStatusModify(str, Enum):
 32    active = "active"
 33    disabled = "disabled"
 34    on_hold = "on_hold"
 35
 36
 37class UserStatusCreate(str, Enum):
 38    active = "active"
 39    on_hold = "on_hold"
 40
 41
 42class UserDataLimitResetStrategy(str, Enum):
 43    no_reset = "no_reset"
 44    day = "day"
 45    week = "week"
 46    month = "month"
 47    year = "year"
 48
 49
 50class UserBase(BaseModel):
 51    # proxies: Dict[ProxyTypes, ProxySettings] = {}
 52    id: Optional[int] = None
 53    username: Annotated[str, StringConstraints(to_lower=True)]
 54    expire: Optional[int] = Field(None, nullable=True)
 55    key: str = Field(default_factory=lambda: secrets.token_hex(16))
 56    data_limit: Optional[int] = Field(
 57        ge=0, default=None, description="data_limit can be 0 or greater"
 58    )
 59    data_limit_reset_strategy: UserDataLimitResetStrategy = (
 60        UserDataLimitResetStrategy.no_reset
 61    )
 62    # inbounds: Dict[ProxyTypes, List[str]] = {}
 63    note: Optional[str] = Field(None, nullable=True)
 64    sub_updated_at: Optional[datetime] = Field(None, nullable=True)
 65    sub_last_user_agent: Optional[str] = Field(None, nullable=True)
 66    online_at: Optional[datetime] = Field(None, nullable=True)
 67    on_hold_expire_duration: Optional[int] = Field(None, nullable=True)
 68    on_hold_timeout: Optional[Union[datetime, None]] = Field(None, nullable=True)
 69
 70    @field_validator("username")
 71    @classmethod
 72    def validate_username(cls, v: str):
 73        if not USERNAME_REGEXP.match(v):
 74            raise ValueError(
 75                "Username should be 3 to 32 characters containing A-Z, a-z, 0-9 and underscores in between."
 76            )
 77        return v
 78
 79    @field_validator("note")
 80    @classmethod
 81    def validate_note(cls, v: str):
 82        if v and len(v) > 500:
 83            raise ValueError("User's note can be a maximum of 500 character")
 84        return v
 85
 86    @field_validator("on_hold_expire_duration", "on_hold_timeout", check_fields=False)
 87    @classmethod
 88    def validate_timeout(cls, v: Optional[int]):
 89        # Check if expire is 0 or None and timeout is not 0 or None
 90        if (v in (0, None)):
 91            return None
 92        return v
 93    model_config = ConfigDict(from_attributes=True)
 94
 95
 96from app.models.service import ServiceBase
 97from app.models.proxy import InboundBase
 98
 99
100class User(UserBase):
101    pass
102
103
104class UserCreate(User):
105    status: UserStatusCreate = Field(None, validate_default=True)
106    services: List[int] = Field([])
107    model_config = ConfigDict(json_schema_extra={
108        "example": {
109            "username": "user1234",
110            "services": [1, 2, 3],
111            "expire": 0,
112            "data_limit": 0,
113            "data_limit_reset_strategy": "no_reset",
114            "status": "active",
115            "note": "",
116            "on_hold_timeout": "2023-11-03T20:30:00",
117            "on_hold_expire_duration": 0,
118        }
119    })
120
121    # TODO[pydantic]: We couldn't refactor the `validator`, please replace it by `field_validator` manually.
122    # Check https://docs.pydantic.dev/dev-v2/migration/#changes-to-validators for more information.
123    #@validator("status", pre=True, always=True)
124    #def validate_status(cls, value):
125        
126    # TODO[pydantic]: We couldn't refactor the `validator`, please replace it by `field_validator` manually.
127    # Check https://docs.pydantic.dev/dev-v2/migration/#changes-to-validators for more information.
128    @field_validator("status", mode="before")
129    def validate_status(cls, status: str, info: ValidationInfo) -> str:
130        if not status or status not in UserStatusCreate.__members__:
131            return UserStatusCreate.active  # Set to the default if not valid
132        on_hold_expire = info.data.get("on_hold_expire_duration")
133        expire = info.data.get("expire")
134        if status == UserStatusCreate.on_hold:
135            if on_hold_expire == 0 or on_hold_expire is None:
136                raise ValueError("User cannot be on hold without a valid on_hold_expire_duration.")
137            if expire:
138                raise ValueError("User cannot be on hold with specified expire.")
139        return status
140
141
142class UserModify(User):
143    status: Optional[UserStatusModify] = Field(None)
144    services: List[int] = Field([])
145    data_limit_reset_strategy: Optional[UserDataLimitResetStrategy] = Field(None)
146    model_config = ConfigDict(json_schema_extra={
147        "example": {
148            "username": "mammad1234",
149            "services": [1, 2, 3],
150            "expire": 0,
151            "data_limit": 0,
152            "data_limit_reset_strategy": "no_reset",
153            "status": "active",
154            "note": "",
155            "on_hold_timeout": "2023-11-03T20:30:00",
156            "on_hold_expire_duration": 0,
157        }
158    })
159    
160
161    # TODO[pydantic]: We couldn't refactor the `validator`, please replace it by `field_validator` manually.
162    # Check https://docs.pydantic.dev/dev-v2/migration/#changes-to-validators for more information.
163    @validator("status", pre=True, always=True)
164    def validate_status(cls, status, values):
165        on_hold_expire = values.get("on_hold_expire_duration")
166        expire = values.get("expire")
167        if status == UserStatusCreate.on_hold:
168            if not on_hold_expire:
169                raise ValueError("User cannot be on hold without a valid on_hold_expire_duration.")
170            if expire:
171                raise ValueError("User cannot be on hold with specified expire.")
172        return status
173
174
175class UserResponse(User):
176    id: int
177    status: UserStatus
178    used_traffic: int
179    lifetime_used_traffic: int = 0
180    created_at: datetime
181    inbounds: List[InboundBase] = Field([])
182    services: List[ServiceBase] = Field([])
183    links: List[str] | None = []  # Field(None)
184
185    # subscription_url: str = ""
186    model_config = ConfigDict(from_attributes=True)
187
188    @computed_field
189    @property
190    def service_ids(self) -> List[int]:
191        return [s.id for s in self.services]
192
193    @computed_field
194    @property
195    def subscription_url(self) -> str:
196        salt = secrets.token_hex(8)
197        url_prefix = (XRAY_SUBSCRIPTION_URL_PREFIX).replace('*', salt)
198        return f"{url_prefix}/sub/{self.username}/{self.key}"
199    
200
201class UsersResponse(BaseModel): 
202    total: int
203    users: List[UserResponse]
204
205
206class UserUsageResponse(BaseModel):
207    node_id: Union[int, None] = None
208    node_name: str
209    used_traffic: int
210
211
212class UserUsagesResponse(BaseModel):
213    username: str
214    usages: List[UserUsageResponse]