# Conflicts:
#	.DS_Store
This commit is contained in:
xds
2026-02-08 17:53:19 +03:00
33 changed files with 527 additions and 16 deletions

41
.vscode/launch.json vendored
View File

@@ -1,5 +1,44 @@
{
"version": "0.2.0",
"configurations": [
{"name":"Python Debugger: FastAPI","type":"debugpy","request":"launch","module":"uvicorn","args":["main:app","--reload", "--port", "8090"],"jinja":true}
{
"name": "Python Debugger: FastAPI",
"type": "debugpy",
"request": "launch",
"module": "uvicorn",
"args": [
"main:app",
"--reload",
"--port",
"8090"
],
"jinja": true,
"justMyCode": true
},
{
"name": "Python: Current File",
"type": "debugpy",
"request": "launch",
"program": "${file}",
"console": "integratedTerminal",
"justMyCode": true,
"env": {
"PYTHONPATH": "${workspaceFolder}"
}
},
{
"name": "Debug Tests: Current File",
"type": "debugpy",
"request": "launch",
"module": "pytest",
"args": [
"${file}"
],
"console": "integratedTerminal",
"justMyCode": true,
"env": {
"PYTHONPATH": "${workspaceFolder}"
}
}
]
}

Binary file not shown.

BIN
api/.DS_Store vendored Normal file

Binary file not shown.

BIN
api/endpoints/.DS_Store vendored Normal file

Binary file not shown.

Binary file not shown.

Binary file not shown.

96
api/endpoints/admin.py Normal file
View File

@@ -0,0 +1,96 @@
from typing import Annotated, List
from fastapi import APIRouter, Depends, HTTPException, status
from fastapi.security import OAuth2PasswordBearer
from pydantic import BaseModel
from repos.user_repo import UsersRepo, UserStatus
from utils.security import verify_password, create_access_token, ACCESS_TOKEN_EXPIRE_MINUTES, ALGORITHM, SECRET_KEY
from jose import JWTError, jwt
from starlette.requests import Request
router = APIRouter(prefix="/api/admin", tags=["admin"])
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="/auth/token")
from api.endpoints.auth import get_users_repo
async def get_current_user(token: Annotated[str, Depends(oauth2_scheme)], repo: Annotated[UsersRepo, Depends(get_users_repo)]):
credentials_exception = HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="Could not validate credentials",
headers={"WWW-Authenticate": "Bearer"},
)
try:
payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
username: str = payload.get("sub")
if username is None:
raise credentials_exception
except JWTError:
raise credentials_exception
user = await repo.get_user_by_username(username)
if user is None:
raise credentials_exception
return user
async def get_current_admin(user: Annotated[dict, Depends(get_current_user)]):
if not user.get("is_admin"):
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail="Not enough permissions",
)
return user
class UserResponse(BaseModel):
username: str
full_name: str | None = None
status: str
created_at: str | None = None
is_admin: bool
class Config:
from_attributes = True
@router.get("/approvals", response_model=List[UserResponse])
async def list_pending_users(
admin: Annotated[dict, Depends(get_current_admin)],
repo: Annotated[UsersRepo, Depends(get_users_repo)]
):
users = await repo.get_pending_users()
# Pydantic conversion handles the list of dicts
return [
UserResponse(
username=u["username"],
full_name=u.get("full_name"),
status=u["status"],
created_at=str(u.get("created_at")),
is_admin=u.get("is_admin", False)
) for u in users
]
@router.post("/approve/{username}")
async def approve_user(
username: str,
admin: Annotated[dict, Depends(get_current_admin)],
repo: Annotated[UsersRepo, Depends(get_users_repo)]
):
user = await repo.get_user_by_username(username)
if not user:
raise HTTPException(status_code=404, detail="User not found")
await repo.approve_user(username)
return {"message": f"User {username} approved"}
@router.post("/deny/{username}")
async def deny_user(
username: str,
admin: Annotated[dict, Depends(get_current_admin)],
repo: Annotated[UsersRepo, Depends(get_users_repo)]
):
user = await repo.get_user_by_username(username)
if not user:
raise HTTPException(status_code=404, detail="User not found")
await repo.deny_user(username)
return {"message": f"User {username} denied"}

View File

@@ -18,6 +18,8 @@ import logging
logger = logging.getLogger(__name__)
from api.endpoints.auth import get_current_user
router = APIRouter(prefix="/api/assets", tags=["Assets"])
@@ -49,7 +51,7 @@ async def get_asset(
return Response(content=content, media_type=media_type, headers=headers)
@router.delete("/{asset_id}", status_code=status.HTTP_204_NO_CONTENT)
@router.delete("/{asset_id}", status_code=status.HTTP_204_NO_CONTENT, dependencies=[Depends(get_current_user)])
async def delete_asset(
asset_id: str,
dao: DAO = Depends(get_dao)
@@ -65,7 +67,7 @@ async def delete_asset(
return None
@router.get("")
@router.get("", dependencies=[Depends(get_current_user)])
async def get_assets(request: Request, dao: DAO = Depends(get_dao), type: Optional[str] = None, limit: int = 10, offset: int = 0) -> AssetsResponse:
logger.info(f"get_assets called. Limit: {limit}, Offset: {offset}")
assets = await dao.assets.get_assets(type, limit, offset)
@@ -82,7 +84,7 @@ async def get_assets(request: Request, dao: DAO = Depends(get_dao), type: Option
@router.post("/upload", response_model=AssetResponse, status_code=status.HTTP_201_CREATED)
@router.post("/upload", response_model=AssetResponse, status_code=status.HTTP_201_CREATED, dependencies=[Depends(get_current_user)])
async def upload_asset(
file: UploadFile = File(...),
linked_char_id: Optional[str] = Form(None),
@@ -127,7 +129,7 @@ async def upload_asset(
)
@router.post("/regenerate_thumbnails")
@router.post("/regenerate_thumbnails", dependencies=[Depends(get_current_user)])
async def regenerate_thumbnails(dao: DAO = Depends(get_dao)):
"""
Regenerates thumbnails for all existing image assets that don't have one.
@@ -161,7 +163,7 @@ async def regenerate_thumbnails(dao: DAO = Depends(get_dao)):
return {"status": "completed", "processed": count, "updated": updated}
@router.post("/migrate_to_minio")
@router.post("/migrate_to_minio", dependencies=[Depends(get_current_user)])
async def migrate_to_minio(dao: DAO = Depends(get_dao)):
"""
Migrates assets from MongoDB to MinIO.

122
api/endpoints/auth.py Normal file
View File

@@ -0,0 +1,122 @@
from datetime import timedelta
from typing import Annotated
from fastapi import APIRouter, Depends, HTTPException, status
from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
from pydantic import BaseModel
from jose import JWTError, jwt
from repos.user_repo import UsersRepo, UserStatus
from utils.security import verify_password, create_access_token, ACCESS_TOKEN_EXPIRE_MINUTES, ALGORITHM, SECRET_KEY
from starlette.requests import Request
router = APIRouter(prefix="/api/auth", tags=["auth"])
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="/api/auth/token")
async def get_users_repo(request: Request) -> UsersRepo:
if not hasattr(request.app.state, "users_repo"):
raise HTTPException(status_code=500, detail="Users repo not initialized")
return request.app.state.users_repo
async def get_current_user(token: Annotated[str, Depends(oauth2_scheme)], repo: Annotated[UsersRepo, Depends(get_users_repo)]):
credentials_exception = HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="Could not validate credentials",
headers={"WWW-Authenticate": "Bearer"},
)
try:
payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
username: str = payload.get("sub")
if username is None:
raise credentials_exception
except JWTError:
raise credentials_exception
user = await repo.get_user_by_username(username)
if user is None:
raise credentials_exception
return user
async def get_current_admin(user: Annotated[dict, Depends(get_current_user)]):
if not user.get("is_admin"):
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail="Not enough permissions",
)
return user
class UserRegister(BaseModel):
username: str
password: str
full_name: str | None = None
class Token(BaseModel):
access_token: str
token_type: str
class UserResponse(BaseModel):
username: str
full_name: str | None = None
status: str
is_admin: bool = False
@router.get("/me", response_model=UserResponse)
async def read_users_me(current_user: Annotated[dict, Depends(get_current_user)]):
return current_user
@router.post("/register")
async def register(user_data: UserRegister, repo: Annotated[UsersRepo, Depends(get_users_repo)]):
try:
await repo.create_user(
username=user_data.username,
password=user_data.password,
full_name=user_data.full_name
)
except ValueError as e:
raise HTTPException(status_code=400, detail=str(e))
return {"message": "Registration successful. Please wait for administrator approval."}
@router.post("/token", response_model=Token)
async def login_for_access_token(
form_data: Annotated[OAuth2PasswordRequestForm, Depends()],
repo: Annotated[UsersRepo, Depends(get_users_repo)]
):
user = await repo.get_user_by_username(form_data.username)
if not user:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="Incorrect username or password",
headers={"WWW-Authenticate": "Bearer"},
)
# Проверяем пароль
if not verify_password(form_data.password, user["hashed_password"]):
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="Incorrect username or password",
headers={"WWW-Authenticate": "Bearer"},
)
# Проверка статуса
if user.get("status") != UserStatus.ALLOWED:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail="Account is not approved yet. Please contact administrator.",
)
access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
access_token = create_access_token(
data={"sub": user["username"]}, expires_delta=access_token_expires
)
return {"access_token": access_token, "token_type": "bearer"}

View File

@@ -16,7 +16,9 @@ import logging
logger = logging.getLogger(__name__)
router = APIRouter(prefix="/api/characters", tags=["Characters"])
from api.endpoints.auth import get_current_user
router = APIRouter(prefix="/api/characters", tags=["Characters"], dependencies=[Depends(get_current_user)])
@router.get("/", response_model=List[Character])

View File

@@ -11,11 +11,15 @@ from api.models.GenerationRequest import GenerationResponse, GenerationRequest,
from api.service.generation_service import GenerationService
from models.Generation import Generation
from starlette import status
import logging
logger = logging.getLogger(__name__)
router = APIRouter(prefix='/api/generations', tags=["Generation"])
from api.endpoints.auth import get_current_user
router = APIRouter(prefix='/api/generations', tags=["Generation"], dependencies=[Depends(get_current_user)])
@router.post("/prompt-assistant", response_model=PromptResponse)
@@ -69,3 +73,12 @@ async def get_generation(generation_id: str,
async def get_running_generations(request: Request,
generation_service: GenerationService = Depends(get_generation_service)):
return await generation_service.get_running_generations()
@router.delete("/{generation_id}", status_code=status.HTTP_204_NO_CONTENT, dependencies=[Depends(get_current_user)])
async def delete_generation(generation_id: str, generation_service: GenerationService = Depends(get_generation_service)):
logger.info(f"delete_generation called for ID: {generation_id}")
deleted = await generation_service.delete_generation(generation_id)
if not deleted:
raise HTTPException(status_code=404, detail="Generation not found")
return None

BIN
api/service/.DS_Store vendored Normal file

Binary file not shown.

View File

@@ -323,3 +323,21 @@ class GenerationService:
pass
except Exception as e:
logger.error(f"Error in progress simulation: {e}")
async def delete_generation(self, generation_id: str) -> bool:
"""
Soft delete generation by marking it as deleted.
"""
try:
generation = await self.dao.generations.get_generation(generation_id)
if not generation:
return False
generation.is_deleted = True
generation.updated_at = datetime.now(UTC)
await self.dao.generations.update_generation(generation)
return True
except Exception as e:
logger.error(f"Error deleting generation {generation_id}: {e}")
return False

10
main.py
View File

@@ -36,6 +36,8 @@ from routers.assets_router import router as assets_router # Роутер бот
from api.endpoints.assets_router import router as api_assets_router # Роутер FastAPI
from api.endpoints.character_router import router as api_char_router # Роутер FastAPI
from api.endpoints.generation_router import router as api_gen_router
from api.endpoints.auth import router as api_auth_router
from api.endpoints.admin import router as api_admin_router
load_dotenv()
logger = logging.getLogger(__name__)
@@ -126,11 +128,11 @@ async def lifespan(app: FastAPI):
# Инициализируем DAO для ассетов и кладем в state приложения
# Теперь в эндпоинтах можно делать request.app.state.assets_dao
app.state.mongo_client = mongo_client
app.state.mongo_client = mongo_client
app.state.gemini_client = gemini
app.state.bot = bot
app.state.s3_adapter = s3_adapter
app.state.users_repo = users_repo # Добавляем репозиторий в state
print("✅ DB & DAO initialized")
@@ -172,9 +174,15 @@ app.add_middleware(
)
# Подключаем роутер API
from api.endpoints.auth import router as auth_api_router
from api.endpoints.admin import router as admin_api_router
app.include_router(auth_api_router)
app.include_router(admin_api_router)
app.include_router(api_assets_router)
app.include_router(api_char_router)
app.include_router(api_gen_router)
app.include_router(api_admin_router)
app.include_router(api_auth_router)
# --- ХЕНДЛЕРЫ БОТА (Main Router) ---

BIN
models/.DS_Store vendored Normal file

Binary file not shown.

View File

@@ -33,6 +33,6 @@ class Generation(BaseModel):
token_usage: Optional[int] = None
input_token_usage: Optional[int] = None
output_token_usage: Optional[int] = None
is_deleted: bool = False
created_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
updated_at: datetime = Field(default_factory=lambda: datetime.now(UTC))

BIN
repos/.DS_Store vendored Normal file

Binary file not shown.

View File

@@ -26,12 +26,13 @@ class GenerationRepo:
async def get_generations(self, character_id: Optional[str] = None, status: Optional[GenerationStatus] = None,
limit: int = 10, offset: int = 10) -> List[Generation]:
args = {}
filter = {"is_deleted": False}
if character_id is not None:
args["linked_character_id"] = character_id
filter["linked_character_id"] = character_id
if status is not None:
args["status"] = status
res = await self.collection.find(args).sort("created_at", -1).skip(
filter["status"] = status
res = await self.collection.find(filter).sort("created_at", -1).skip(
offset).limit(limit).to_list(None)
generations: List[Generation] = []
for generation in res:

View File

@@ -1,8 +1,10 @@
from datetime import datetime, timedelta
from enum import Enum
from typing import Optional
from aiogram.types import User
from motor.motor_asyncio import AsyncIOMotorClient
from utils.security import get_password_hash
class UserStatus:
@@ -19,10 +21,49 @@ class UsersRepo:
async def get_user(self, user_id: int):
return await self.collection.find_one({"user_id": user_id})
async def get_user_by_username(self, username: str):
return await self.collection.find_one({"username": username})
async def create_user(self, username: str, password: str, full_name: Optional[str] = None):
"""Создает нового пользователя с username/паролем"""
existing = await self.get_user_by_username(username)
if existing:
raise ValueError("User with this username already exists")
user_doc = {
"username": username,
"hashed_password": get_password_hash(password),
"full_name": full_name,
"status": UserStatus.PENDING, # По умолчанию PENDING
"created_at": datetime.now(),
"is_email_user": False, # Теперь это просто "обычный" юзер, не телеграм (хотя поле можно переименовать)
"is_web_user": True,
"is_admin": False
}
result = await self.collection.insert_one(user_doc)
return await self.collection.find_one({"_id": result.inserted_id})
async def get_pending_users(self):
"""Возвращает список пользователей со статусом PENDING"""
cursor = self.collection.find({"status": UserStatus.PENDING})
return await cursor.to_list(length=100)
async def approve_user(self, username: str):
await self.collection.update_one(
{"username": username},
{"$set": {"status": UserStatus.ALLOWED}}
)
async def deny_user(self, username: str):
await self.collection.update_one(
{"username": username},
{"$set": {"status": UserStatus.DENIED}}
)
async def create_or_update_request(self, user: User):
"""
Обновляет дату последнего запроса и ставит статус PENDING.
Сохраняет всю инфу о юзере.
Сохраняет всю инфу о юзере (для Telegram пользователей).
"""
now = datetime.now()
data = {
@@ -30,7 +71,8 @@ class UsersRepo:
"username": user.username,
"full_name": user.full_name,
"status": UserStatus.PENDING,
"last_request_date": now
"last_request_date": now,
"is_email_user": False
}
await self.collection.update_one(
{"user_id": user.id},

View File

@@ -46,3 +46,7 @@ uvicorn==0.40.0
websockets==15.0.1
yarl==1.22.0
aioboto3==13.3.0
passlib[argon2]==1.7.4
python-jose[cryptography]==3.3.0
python-multipart==0.0.22
email-validator

View File

@@ -0,0 +1,22 @@
import pytest
from fastapi.testclient import TestClient
from main import app
client = TestClient(app)
def test_api_protection():
# 1. Assets
response = client.get("/api/assets")
assert response.status_code == 401
# 2. Characters
response = client.get("/api/characters")
assert response.status_code == 401
# 3. Generations
response = client.get("/api/generations")
assert response.status_code == 401
# 4. Upload Asset (POST)
response = client.post("/api/assets/upload")
assert response.status_code == 401

107
tests/test_auth_flow.py Normal file
View File

@@ -0,0 +1,107 @@
import pytest
from fastapi.testclient import TestClient
from unittest.mock import AsyncMock, MagicMock
from datetime import datetime
from main import app
from api.endpoints.auth import get_users_repo
from repos.user_repo import UsersRepo, UserStatus
from utils.security import get_password_hash
# Mock Repository
class MockUsersRepo:
def __init__(self):
self.users = {}
async def get_user_by_username(self, username: str):
return self.users.get(username)
async def create_user(self, username: str, password: str, full_name: str = None):
if username in self.users:
raise ValueError("User with this username already exists")
user = {
"username": username,
"hashed_password": get_password_hash(password),
"full_name": full_name,
"status": UserStatus.PENDING,
"is_email_user": False,
"is_admin": False,
"created_at": datetime.now()
}
self.users[username] = user
return user
async def get_pending_users(self):
return [u for u in self.users.values() if u["status"] == UserStatus.PENDING]
async def approve_user(self, username: str):
if username in self.users:
self.users[username]["status"] = UserStatus.ALLOWED
async def deny_user(self, username: str):
if username in self.users:
self.users[username]["status"] = UserStatus.DENIED
mock_repo = MockUsersRepo()
# Override Dependency
app.dependency_overrides[get_users_repo] = lambda: mock_repo
client = TestClient(app)
def test_auth_flow_with_approval():
# 1. Register (User)
user_data = {
"username": "newuser",
"password": "password123",
"full_name": "New User"
}
response = client.post("/auth/register", json=user_data)
assert response.status_code == 200
assert response.json()["message"] == "Registration successful. Please wait for administrator approval."
# 2. Try Login (User) -> Should Fail (Pending)
login_data = {
"username": "newuser",
"password": "password123"
}
response = client.post("/auth/token", data=login_data)
assert response.status_code == 403
assert "not approved" in response.json()["detail"]
# 3. Setup Admin (Backdoor for test)
mock_repo.users["admin"] = {
"username": "admin",
"hashed_password": get_password_hash("adminpass"),
"status": UserStatus.ALLOWED,
"is_admin": True,
"created_at": datetime.now()
}
# 4. Admin Login
admin_login = {
"username": "admin",
"password": "adminpass"
}
response = client.post("/auth/token", data=admin_login)
assert response.status_code == 200
admin_token = response.json()["access_token"]
admin_auth = {"Authorization": f"Bearer {admin_token}"}
# 5. List Pending (Admin)
response = client.get("/admin/approvals", headers=admin_auth)
assert response.status_code == 200
users = response.json()
assert len(users) >= 1
assert users[0]["username"] == "newuser"
assert users[0]["status"] == "pending"
# 6. Approve User (Admin)
response = client.post("/admin/approve/newuser", headers=admin_auth)
assert response.status_code == 200
assert response.json()["message"] == "User newuser approved"
# 7. Login User (Again) -> Should Success
response = client.post("/auth/token", data=login_data)
assert response.status_code == 200
assert "access_token" in response.json()

Binary file not shown.

35
utils/security.py Normal file
View File

@@ -0,0 +1,35 @@
from datetime import datetime, timedelta
from typing import Optional, Union, Any
from jose import jwt
from passlib.context import CryptContext
# Настройки безопасности (лучше вынести в config/env, но для старта здесь)
# SECRET_KEY должен быть сложным и секретным в продакшене!
SECRET_KEY = "CHANGE_ME_TO_A_SUPER_SECRET_KEY"
ALGORITHM = "HS256"
ACCESS_TOKEN_EXPIRE_MINUTES = 30 * 24 * 60 # 30 дней, например
pwd_context = CryptContext(schemes=["argon2"], deprecated="auto")
def verify_password(plain_password: str, hashed_password: str) -> bool:
return pwd_context.verify(plain_password, hashed_password)
def get_password_hash(password: str) -> str:
return pwd_context.hash(password)
def create_access_token(data: dict, expires_delta: Optional[timedelta] = None) -> str:
to_encode = data.copy()
if expires_delta:
expire = datetime.utcnow() + expires_delta
else:
expire = datetime.utcnow() + timedelta(minutes=15)
# Стандартное поле 'exp' для JWT
to_encode.update({"exp": expire})
encoded_jwt = jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
return encoded_jwt