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