Добавлены модели и функции для работы с аккаунтами в базе данных. Реализовано заполнение базы данных тестовыми аккаунтами, добавлено хеширование для идентификации TgAgent. Обновлены функции авторизации и получения аккаунта.

This commit is contained in:
Redsandyg 2025-06-03 10:57:23 +03:00
parent 37c855c601
commit 4c4a84eefe
2 changed files with 59 additions and 12 deletions

View File

@ -1,9 +1,10 @@
import random
from uuid import uuid4
from sqlmodel import Session
from main import AUTH_DB_ENGINE, TgAgent, Ref, Sale, Transaction
from main import AUTH_DB_ENGINE, TgAgent, Ref, Sale, Transaction, Account
from sqlalchemy import text
from datetime import datetime, timedelta
from hashlib import sha256
# Константа: список user_ids
@ -72,11 +73,28 @@ def fill_db():
session.execute(text("DELETE FROM sale"))
session.execute(text("DELETE FROM ref"))
session.execute(text("DELETE FROM tgagent"))
session.execute(text("DELETE FROM account"))
session.commit()
# 0. Accounts
accounts = []
for i in range(4):
acc = Account(
login=f"user{i+1}",
password="password123", # В реальном проекте пароли должны быть захешированы!
name=NAMES[i % len(NAMES)],
email=f"user{i+1}@example.com",
balance=round(random.uniform(1000, 10000), 2)
)
session.add(acc)
accounts.append(acc)
session.commit()
for acc in accounts:
session.refresh(acc)
# 1. TgAgents
tg_agents = []
for i, tg_agent_id in enumerate(USER_IDS):
dt = random.choice(date_list)
hash_value = sha256(f"{tg_agent_id}sold".encode()).hexdigest()
tg_agent = TgAgent(
tg_id=tg_agent_id,
chat_id=tg_agent_id, # chat_id совпадает с tg_id
@ -84,7 +102,8 @@ def fill_db():
name=NAMES[i % len(NAMES)],
login=LOGINS[i % len(LOGINS)],
create_dttm=dt,
update_dttm=dt
update_dttm=dt,
hash=hash_value
)
session.add(tg_agent)
tg_agents.append(tg_agent)

48
main.py
View File

@ -1,4 +1,4 @@
from fastapi import FastAPI, Depends, HTTPException, status, Query
from fastapi import FastAPI, Depends, HTTPException, status, Query, Body, Request
from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
from sqlmodel import SQLModel, Field, create_engine, Session, select
from passlib.context import CryptContext
@ -8,6 +8,7 @@ from models import RefAddRequest, RefResponse, RegisterRequest, Token, TokenReq
from uuid import uuid4
from fastapi.responses import JSONResponse
from sqlalchemy import func
from hashlib import sha256
# Конфигурация
AUTH_DATABASE_ADDRESS = "sqlite:///partner.db"
@ -20,6 +21,7 @@ class TgAgent(SQLModel, table=True):
phone: Optional[str] = None
name: Optional[str] = None
login: Optional[str] = None
hash: Optional[str] = None
create_dttm: datetime = Field(default_factory=datetime.utcnow)
update_dttm: datetime = Field(default_factory=datetime.utcnow)
@ -49,6 +51,14 @@ class Transaction(SQLModel, table=True):
create_dttm: datetime = Field(default_factory=datetime.utcnow)
update_dttm: datetime = Field(default_factory=datetime.utcnow)
class Account(SQLModel, table=True):
id: Optional[int] = Field(default=None, primary_key=True)
login: str = Field(index=True, unique=True)
password: str
name: Optional[str] = None
email: Optional[str] = None
balance: float = 0.0
# Создание движка базы данных
AUTH_DB_ENGINE = create_engine(AUTH_DATABASE_ADDRESS, echo=True)
SQLModel.metadata.create_all(AUTH_DB_ENGINE)
@ -74,20 +84,17 @@ def get_db():
yield session
# Авторизация
async def get_current_tg_agent(token: str = Depends(oauth2_scheme), db: Session = Depends(get_db)):
async def get_current_tg_agent(request: Request, db: Session = Depends(get_db)):
credentials_exception = HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="Could not validate credentials",
headers={"WWW-Authenticate": "Bearer"},
)
# Ожидаем токен вида 'session_for_{tg_id}'
if not token.startswith("session_for_"):
auth_header = request.headers.get("Authorization")
if not auth_header or not auth_header.startswith("Bearer "):
raise credentials_exception
try:
tg_id = int(token.replace("session_for_", ""))
except Exception:
raise credentials_exception
tg_agent = get_tg_agent_by_tg_id(db, tg_id)
hash_value = auth_header.replace("Bearer ", "").strip()
tg_agent = db.exec(select(TgAgent).where(TgAgent.hash == hash_value)).first()
if tg_agent is None:
raise credentials_exception
return tg_agent
@ -106,7 +113,8 @@ def register(req: RegisterRequest, db: Session = Depends(get_db)):
tg_agent = get_tg_agent_by_tg_id(db, tg_id)
if tg_agent:
raise HTTPException(status_code=400, detail="tg_id already registered")
new_tg_agent = TgAgent(tg_id=tg_id, chat_id=chat_id, phone=phone, name=name, login=login)
hash_value = sha256(f"{tg_id}sold".encode()).hexdigest()
new_tg_agent = TgAgent(tg_id=tg_id, chat_id=chat_id, phone=phone, name=name, login=login, hash=hash_value)
db.add(new_tg_agent)
db.commit()
db.refresh(new_tg_agent)
@ -446,3 +454,23 @@ def get_billing_chart_pie(db: Session = Depends(get_db)):
for row in result
]
return JSONResponse(content=data)
@app.get("/account", tags=["bff"])
def get_account(db: Session = Depends(get_db)):
account = db.exec(select(Account)).first()
if not account:
raise HTTPException(status_code=404, detail="Account not found")
return {
"id": account.id,
"login": account.login,
"name": account.name,
"email": account.email,
"balance": account.balance
}
@app.post("/tg_auth", tags=["partner-tg"])
def tg_auth(hash: str = Body(..., embed=True), db: Session = Depends(get_db)):
tg_agent = db.exec(select(TgAgent).where(TgAgent.hash == hash)).first()
if not tg_agent:
raise HTTPException(status_code=401, detail="Hash not found")
return {"msg": "Auth success", "tg_id": tg_agent.tg_id}