266 lines
8.1 KiB
Python
266 lines
8.1 KiB
Python
from fastapi import FastAPI, Depends, HTTPException, status
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
from sqlalchemy.future import select
|
|
from sqlalchemy.orm import selectinload
|
|
from . import models, schemas, auth, database
|
|
from passlib.context import CryptContext
|
|
from sqlalchemy.exc import IntegrityError
|
|
from contextlib import asynccontextmanager
|
|
import httpx
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(app: FastAPI):
|
|
# 启动时执行：创建表
|
|
await database.init_db()
|
|
yield
|
|
# 关闭时执行（如果需要）
|
|
|
|
app = FastAPI(lifespan=lifespan)
|
|
|
|
|
|
@app.post("/register", response_model=schemas.UserOut)
|
|
async def register(user: schemas.UserCreate, db: AsyncSession = Depends(database.get_db)):
|
|
hashed_pwd = auth.hash_password(user.password)
|
|
private_key, public_key = auth.generate_key_pair()
|
|
|
|
new_user = models.User(
|
|
username=user.username,
|
|
hashed_password=hashed_pwd,
|
|
private_key=private_key,
|
|
public_key=public_key
|
|
)
|
|
db.add(new_user)
|
|
try:
|
|
await db.commit()
|
|
await db.refresh(new_user)
|
|
return new_user
|
|
|
|
except IntegrityError:
|
|
# 发生唯一约束冲突时回滚并报错
|
|
await db.rollback()
|
|
raise HTTPException(
|
|
status_code=status.HTTP_400_BAD_REQUEST,
|
|
detail="用户名已存在"
|
|
)
|
|
|
|
|
|
|
|
@app.post("/token")
|
|
async def login(form_data: schemas.UserLogin, db: AsyncSession = Depends(database.get_db)):
|
|
result = await db.execute(select(models.User).where(models.User.username == form_data.username))
|
|
user = result.scalars().first()
|
|
if not user or not auth.verify_password(form_data.password, user.hashed_password):
|
|
raise HTTPException(status_code=400, detail="Incorrect username or password")
|
|
|
|
access_token = auth.create_access_token(data={"sub": user.username})
|
|
return {"access_token": access_token, "token_type": "bearer"}
|
|
|
|
|
|
@app.post("/assets/", response_model=schemas.AssetOut)
|
|
async def create_asset(
|
|
asset: schemas.AssetCreate,
|
|
current_user: models.User = Depends(auth.get_current_user),
|
|
db: AsyncSession = Depends(database.get_db)
|
|
):
|
|
# Encrypt the inner content using user's public key
|
|
encrypted_content = auth.encrypt_data(asset.content_inner_encrypted, current_user.public_key)
|
|
|
|
new_asset = models.Asset(
|
|
title=asset.title,
|
|
content_outer_encrypted=encrypted_content,
|
|
private_key_shard=asset.private_key_shard,
|
|
author_id=current_user.id
|
|
)
|
|
db.add(new_asset)
|
|
await db.commit()
|
|
await db.refresh(new_asset)
|
|
return new_asset
|
|
|
|
|
|
@app.post("/assets/claim")
|
|
async def claim_asset(
|
|
asset_claim: schemas.AssetClaim,
|
|
current_user: models.User = Depends(auth.get_current_user),
|
|
db: AsyncSession = Depends(database.get_db)
|
|
):
|
|
# Fetch asset with author loaded
|
|
result = await db.execute(
|
|
select(models.Asset)
|
|
.options(selectinload(models.Asset.author))
|
|
.where(models.Asset.id == asset_claim.asset_id)
|
|
)
|
|
asset = result.scalars().first()
|
|
|
|
if not asset:
|
|
raise HTTPException(status_code=404, detail="Asset not found")
|
|
|
|
# 1. 验证用户是否是继承人
|
|
if asset.heir_id != current_user.id:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_403_FORBIDDEN,
|
|
detail="You are not the designated heir for this asset"
|
|
)
|
|
|
|
# 2. 验证所有人是否已经挂了 (guale)
|
|
if not asset.author.guale:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_403_FORBIDDEN,
|
|
detail="The owner of this asset is still alive. You cannot claim it yet."
|
|
)
|
|
|
|
# 3. 验证通过后用asset所有人的private_key解密内容
|
|
try:
|
|
decrypted_content = auth.decrypt_data(
|
|
asset.content_outer_encrypted,
|
|
asset.author.private_key
|
|
)
|
|
except Exception as e:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
detail=f"Failed to decrypt asset: {str(e)}"
|
|
)
|
|
|
|
return {
|
|
"asset_id": asset.id,
|
|
"title": asset.title,
|
|
"decrypted_content": decrypted_content,
|
|
"server_shard_key": asset.private_key_shard
|
|
}
|
|
|
|
|
|
@app.post("/assets/assign")
|
|
async def assign_asset(
|
|
assignment: schemas.AssetAssign,
|
|
current_user: models.User = Depends(auth.get_current_user),
|
|
db: AsyncSession = Depends(database.get_db)
|
|
):
|
|
# Fetch Asset
|
|
result = await db.execute(
|
|
select(models.Asset)
|
|
.options(selectinload(models.Asset.heir))
|
|
.where(models.Asset.id == assignment.asset_id)
|
|
)
|
|
asset = result.scalars().first()
|
|
|
|
if not asset:
|
|
raise HTTPException(status_code=404, detail="Asset not found")
|
|
|
|
if asset.author_id != current_user.id:
|
|
raise HTTPException(status_code=403, detail="Not authorized to assign this asset")
|
|
|
|
|
|
heir_result = await db.execute(
|
|
select(models.User).where(
|
|
models.User.username == assignment.heir_name
|
|
)
|
|
)
|
|
heir_user = heir_result.scalars().first()
|
|
|
|
if not heir_user:
|
|
raise HTTPException(status_code=404, detail="Heir not found")
|
|
|
|
if heir_user.id == current_user.id:
|
|
asset.heir = None
|
|
await db.commit()
|
|
#raise HTTPException(status_code=403, detail="You cannot assign an asset to yourself")
|
|
return {"message": "Asset unassigned"}
|
|
|
|
asset.heir = heir_user
|
|
await db.commit()
|
|
|
|
return {"message": f"Asset assigned to {assignment.heir_name}"}
|
|
|
|
|
|
|
|
|
|
@app.post("/admin/declare-guale")
|
|
async def declare_user_guale(
|
|
declare: schemas.DeclareGuale,
|
|
current_user: models.User = Depends(auth.get_current_user),
|
|
db: AsyncSession = Depends(database.get_db)
|
|
):
|
|
# Check if current user is admin
|
|
if not current_user.is_admin:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_403_FORBIDDEN,
|
|
detail="Only administrators can declare users as deceased"
|
|
)
|
|
|
|
# Find the target user
|
|
result = await db.execute(
|
|
select(models.User).where(models.User.username == declare.username)
|
|
)
|
|
target_user = result.scalars().first()
|
|
|
|
if not target_user:
|
|
raise HTTPException(status_code=404, detail="User not found")
|
|
|
|
# Set guale to True
|
|
target_user.guale = True
|
|
await db.commit()
|
|
|
|
return {
|
|
"message": f"User {declare.username} has been declared as deceased",
|
|
"username": target_user.username,
|
|
"guale": target_user.guale
|
|
}
|
|
|
|
# 用于测试热加载
|
|
@app.post("/post1")
|
|
async def test1():
|
|
a=2
|
|
b=3
|
|
c = a+b
|
|
return {"msg": f"this is a msg {c}"}
|
|
|
|
|
|
@app.post("/ai/proxy", response_model=schemas.AIResponse)
|
|
async def ai_proxy(
|
|
ai_request: schemas.AIRequest,
|
|
current_user: models.User = Depends(auth.get_current_user),
|
|
db: AsyncSession = Depends(database.get_db)
|
|
):
|
|
"""
|
|
Proxy relay for AI requests.
|
|
Fetches AI configuration from the database.
|
|
"""
|
|
# Fetch active AI config
|
|
result = await db.execute(
|
|
select(models.AIConfig).where(models.AIConfig.is_active == True)
|
|
)
|
|
config = result.scalars().first()
|
|
|
|
if not config:
|
|
raise HTTPException(status_code=500, detail="AI configuration not found")
|
|
|
|
headers = {
|
|
"Authorization": f"Bearer {config.api_key}",
|
|
"Content-Type": "application/json"
|
|
}
|
|
|
|
# Prepare payload
|
|
payload = ai_request.model_dump()
|
|
payload["model"] = config.default_model
|
|
|
|
async with httpx.AsyncClient() as client:
|
|
try:
|
|
response = await client.post(
|
|
config.api_url,
|
|
json=payload,
|
|
headers=headers,
|
|
timeout=30.0
|
|
)
|
|
response.raise_for_status()
|
|
return response.json()
|
|
except httpx.HTTPStatusError as e:
|
|
raise HTTPException(
|
|
status_code=e.response.status_code,
|
|
detail=f"AI provider returned an error: {e.response.text}"
|
|
)
|
|
except httpx.RequestError as e:
|
|
raise HTTPException(
|
|
status_code=500,
|
|
detail=f"An error occurred while requesting AI provider: {str(e)}"
|
|
)
|
|
|