| 1 | import datetime |
| 2 | import re |
| 3 | import secrets |
| 4 | import string |
| 5 | from enum import Enum |
| 6 | from typing import List |
| 7 | from typing import Optional |
| 8 | |
| 9 | import bcrypt |
| 10 | from pydantic import BaseModel |
| 11 | from pydantic import EmailStr |
| 12 | from pydantic import field_validator |
| 13 | from sqlmodel import Field |
| 14 | from sqlmodel import Relationship |
| 15 | from sqlmodel import SQLModel |
| 16 | |
| 17 | |
| 18 | class Role(SQLModel, table=True): |
| 19 | id: Optional[int] = Field(primary_key=True) |
| 20 | name: str = Field(max_length=256) |
| 21 | description: str = Field(max_length=256) |
| 22 | |
| 23 | user: Optional["User"] = Relationship(back_populates="role") |
| 24 | tag_access: List["RoleTagAccess"] = Relationship(back_populates="role") |
| 25 | |
| 26 | |
| 27 | class UserCustomerAccess(SQLModel, table=True): |
| 28 | __tablename__ = "user_customer_access" |
| 29 | id: Optional[int] = Field(primary_key=True) |
| 30 | user_id: int = Field(foreign_key="user.id") |
| 31 | customer_code: str = Field(foreign_key="customers.customer_code") |
| 32 | created_at: datetime.datetime = Field(default_factory=datetime.datetime.now) |
| 33 | |
| 34 | # Relationships |
| 35 | user: "User" = Relationship(back_populates="customer_access") |
| 36 | |
| 37 | |
| 38 | class UserTagAccess(SQLModel, table=True): |
| 39 | """Defines which tags a user can access (allow-list).""" |
| 40 | |
| 41 | __tablename__ = "user_tag_access" |
| 42 | id: Optional[int] = Field(primary_key=True) |
| 43 | user_id: int = Field(foreign_key="user.id") |
| 44 | tag_id: int = Field(foreign_key="incident_management_alerttag.id") |
| 45 | created_at: datetime.datetime = Field(default_factory=datetime.datetime.now) |
| 46 | |
| 47 | # Relationships |
| 48 | user: "User" = Relationship(back_populates="tag_access") |
| 49 | |
| 50 | |
| 51 | class RoleTagAccess(SQLModel, table=True): |
| 52 | """Defines which tags a role can access (allow-list).""" |
| 53 | |
| 54 | __tablename__ = "role_tag_access" |
| 55 | id: Optional[int] = Field(primary_key=True) |
| 56 | role_id: int = Field(foreign_key="role.id") |
| 57 | tag_id: int = Field(foreign_key="incident_management_alerttag.id") |
| 58 | created_at: datetime.datetime = Field(default_factory=datetime.datetime.now) |
| 59 | |
| 60 | # Relationships |
| 61 | role: "Role" = Relationship(back_populates="tag_access") |
| 62 | |
| 63 | |
| 64 | class User(SQLModel, table=True): |
| 65 | id: Optional[int] = Field(primary_key=True) |
| 66 | username: str = Field(index=True, max_length=256) |
| 67 | password: str = Field(max_length=256, min_length=6) |
| 68 | email: EmailStr |
| 69 | created_at: datetime.datetime = datetime.datetime.now() |
| 70 | role_id: Optional[int] = Field(foreign_key="role.id") |
| 71 | |
| 72 | role: Optional["Role"] = Relationship(back_populates="user") |
| 73 | customer_access: List["UserCustomerAccess"] = Relationship(back_populates="user") |
| 74 | tag_access: List["UserTagAccess"] = Relationship(back_populates="user") |
| 75 | |
| 76 | |
| 77 | # Enum class for role_id 1,2 |
| 78 | class RoleEnum(int, Enum): |
| 79 | admin = 1 |
| 80 | analyst = 2 |
| 81 | scheduler = 3 |
| 82 | customer_user = 4 |
| 83 | |
| 84 | |
| 85 | class UserInput(SQLModel): |
| 86 | username: str |
| 87 | # bcrypt's input is capped at 72 bytes; longer passwords would either be silently |
| 88 | # truncated (bcrypt 4 behaviour) or rejected with ValueError (bcrypt 5+). We enforce |
| 89 | # the limit here so users get a clean 422 ValidationError instead. |
| 90 | password: str = Field( |
| 91 | max_length=72, |
| 92 | min_length=8, |
| 93 | regex="^(?=.*[a-z])(?=.*[A-Z])(?=.*\\d)(?=.*[@$!%*?&#]).{8,}$", |
| 94 | description="Password must be 8-72 characters long and contain at least one uppercase letter, one lowercase letter, one number, and one special character", |
| 95 | ) |
| 96 | email: EmailStr |
| 97 | role_id: RoleEnum = Field( |
| 98 | RoleEnum.analyst, |
| 99 | description="Role ID 1: admin, 2: analyst", |
| 100 | foreign_key="role.id", |
| 101 | ) |
| 102 | |
| 103 | @field_validator("role_id") |
| 104 | @classmethod |
| 105 | def check_role_id(cls, value): |
| 106 | if value not in [e.value for e in RoleEnum]: |
| 107 | raise ValueError("Invalid role ID") |
| 108 | return value |
| 109 | |
| 110 | |
| 111 | class UserLogin(SQLModel): |
| 112 | username: str |
| 113 | password: str |
| 114 | |
| 115 | |
| 116 | class Password(BaseModel): |
| 117 | # Capped at 72 — bcrypt's input limit. The generated chars are ASCII so 72 chars = 72 bytes. |
| 118 | length: int = Field( |
| 119 | default=12, |
| 120 | ge=8, |
| 121 | le=72, |
| 122 | description="The length of the password (max 72 — bcrypt input limit)", |
| 123 | ) |
| 124 | hashed: str # Holds the hashed password |
| 125 | plain: str # Holds the plain password |
| 126 | |
| 127 | @field_validator("length") |
| 128 | @classmethod |
| 129 | def validate_length(cls, value): |
| 130 | if value < 8 or value > 72: |
| 131 | raise ValueError("Password length must be between 8 and 72 characters.") |
| 132 | return value |
| 133 | |
| 134 | @classmethod |
| 135 | def generate(cls, length: int = 12) -> "Password": |
| 136 | if length < 8 or length > 72: |
| 137 | raise ValueError("Password length should be between 8 and 72 characters.") |
| 138 | |
| 139 | # Define the characters that can be used in the password |
| 140 | lowercase = string.ascii_lowercase |
| 141 | uppercase = string.ascii_uppercase |
| 142 | digits = string.digits |
| 143 | special = "@$!%*?&#" |
| 144 | |
| 145 | # Ensure the password has at least one lowercase, one uppercase, one digit, and one special char |
| 146 | password_chars = [ |
| 147 | secrets.choice(lowercase), |
| 148 | secrets.choice(uppercase), |
| 149 | secrets.choice(digits), |
| 150 | secrets.choice(special), |
| 151 | ] |
| 152 | |
| 153 | # Fill the rest of the password length with a cryptographically secure random mix |
| 154 | alphabet = lowercase + uppercase + digits + special |
| 155 | if length > 4: |
| 156 | password_chars += [secrets.choice(alphabet) for _ in range(length - 4)] |
| 157 | |
| 158 | # Shuffle using secrets-backed SystemRandom to avoid predictable patterns |
| 159 | secrets.SystemRandom().shuffle(password_chars) |
| 160 | |
| 161 | # Convert the list of characters into a string |
| 162 | password = "".join(password_chars) |
| 163 | |
| 164 | # Hash the password |
| 165 | hashed_password = bcrypt.hashpw(password.encode("utf-8"), bcrypt.gensalt()) |
| 166 | |
| 167 | # Return the Password object with both the plain and hashed password |
| 168 | return cls( |
| 169 | length=length, |
| 170 | hashed=hashed_password.decode("utf-8"), |
| 171 | plain=password, |
| 172 | ) |
| 173 | |
| 174 | |
| 175 | # ! PASSWORD RESET TOKEN GENERATION NOT USING FOR NOW! # |
| 176 | class PasswordResetRequest(BaseModel): |
| 177 | username: str |
| 178 | |
| 179 | |
| 180 | class PasswordResetToken(BaseModel): |
| 181 | username: str |
| 182 | reset_token: str |
| 183 | new_password: str |
| 184 | |
| 185 | @field_validator("new_password") |
| 186 | @classmethod |
| 187 | def validate_password(cls, password): |
| 188 | if len(password) < 8 or len(password) > 72: |
| 189 | raise ValueError("Password length must be between 8 and 72 characters (bcrypt input limit).") |
| 190 | if not re.search(r"[a-z]", password): |
| 191 | raise ValueError("Password must contain at least one lowercase letter.") |
| 192 | if not re.search(r"[A-Z]", password): |
| 193 | raise ValueError("Password must contain at least one uppercase letter.") |
| 194 | if not re.search(r"\d", password): |
| 195 | raise ValueError("Password must contain at least one digit.") |
| 196 | if not re.search(r"[@$!%*?&#]", password): |
| 197 | raise ValueError("Password must contain at least one special character.") |
| 198 | return password |
| 199 | |
| 200 | |
| 201 | class PasswordReset(BaseModel): |
| 202 | username: str |
| 203 | # Optional: when supplied (e.g. customer-portal self-service flow) the route verifies it |
| 204 | # against the stored hash before changing the password. Admin/analyst reset flows that |
| 205 | # already prove identity via the JWT may omit it, preserving backwards compatibility. |
| 206 | current_password: Optional[str] = None |
| 207 | # 8-72 chars — bcrypt's 72-byte input limit (matches UserInput). |
| 208 | new_password: str = Field( |
| 209 | max_length=72, |
| 210 | min_length=8, |
| 211 | regex="^(?=.*[a-z])(?=.*[A-Z])(?=.*\\d)(?=.*[@$!%*?&#])[A-Za-z\\d@$!%*?&#]{8,}$", |
| 212 | description="Password must be 8-72 characters long and contain at least one uppercase letter, one lowercase letter, one number, and one special character", |
| 213 | ) |