backend_v0.1

2026-01-24 11:02:08 -08:00
commit d22fa8a286
26 changed files with 1088 additions and 0 deletions
--- a/app/pycache/auth.cpython-311.pyc
+++ b/app/pycache/auth.cpython-311.pyc
--- a/app/pycache/database.cpython-311.pyc
+++ b/app/pycache/database.cpython-311.pyc
--- a/app/pycache/main.cpython-311.pyc
+++ b/app/pycache/main.cpython-311.pyc
--- a/app/pycache/models.cpython-311.pyc
+++ b/app/pycache/models.cpython-311.pyc
--- a/app/pycache/schemas.cpython-311.pyc
+++ b/app/pycache/schemas.cpython-311.pyc
--- a/app/auth.py
+++ b/app/auth.py
@@ -0,0 +1,108 @@
+from datetime import datetime, timedelta
+from jose import jwt
+from passlib.context import CryptContext
+
+SECRET_KEY = "your-secret-key"
+ALGORITHM = "HS256"
+
+pwd_context = CryptContext(
+    schemes=["argon2", "bcrypt"],  # 添加 argon2 作为备选方案
+    default="argon2",  # 使用 argon2 作为默认方案
+    deprecated="auto",
+)
+
+def hash_password(password: str) -> str:
+    return pwd_context.hash(password)
+
+
+def verify_password(plain_password, hashed_password):
+    return pwd_context.verify(plain_password, hashed_password)
+
+def create_access_token(data: dict):
+    to_encode = data.copy()
+    expire = datetime.utcnow() + timedelta(minutes=30)
+    to_encode.update({"exp": expire})
+    return jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
+
+from Crypto.PublicKey import RSA
+
+def generate_key_pair():
+    key = RSA.generate(2048)
+    private_key = key.export_key().decode('utf-8')
+    public_key = key.publickey().export_key().decode('utf-8')
+    return private_key, public_key
+
+from Crypto.Cipher import PKCS1_OAEP, AES
+from Crypto.Hash import SHA256
+from Crypto.Random import get_random_bytes
+import base64
+
+def encrypt_data(data: str, public_key_pem: str) -> str:
+    # 1. 准备 RSA 公钥
+    recipient_key = RSA.import_key(public_key_pem)
+    cipher_rsa = PKCS1_OAEP.new(recipient_key)
+
+    # 2. 生成随机 AES 会话密钥 (32 bytes)
+    session_key = get_random_bytes(32)
+
+    # 3. 用 RSA 加密 AES 密钥
+    enc_session_key = cipher_rsa.encrypt(session_key)
+
+    # 4. 用 AES-GCM 加密实际数据
+    cipher_aes = AES.new(session_key, AES.MODE_GCM)
+    ciphertext, tag = cipher_aes.encrypt_and_digest(data.encode('utf-8'))
+
+    # 5. 组合结果: [RSA加密的密钥(256B)] + [Nonce(16B)] + [Tag(16B)] + [密文]
+    combined = enc_session_key + cipher_aes.nonce + tag + ciphertext
+    return base64.b64encode(combined).decode('utf-8')
+
+def decrypt_data(encrypted_data_b64: str, private_key_pem: str) -> str:
+    # 1. 解码合并的二进制数据
+    encrypted_data = base64.b64decode(encrypted_data_b64)
+    
+    # 2. 分解组件 (RSA 2048位产生的密文固定为 256 字节)
+    rsa_key_len = 256 
+    enc_session_key = encrypted_data[:rsa_key_len]
+    nonce = encrypted_data[rsa_key_len : rsa_key_len + 16]
+    tag = encrypted_data[rsa_key_len + 16 : rsa_key_len + 32]
+    ciphertext = encrypted_data[rsa_key_len + 32 :]
+
+    # 3. 用 RSA 私钥解密 AES 会话密钥
+    private_key = RSA.import_key(private_key_pem)
+    cipher_rsa = PKCS1_OAEP.new(private_key)
+    session_key = cipher_rsa.decrypt(enc_session_key)
+
+    # 4. 用 AES 会话密钥解密数据
+    cipher_aes = AES.new(session_key, AES.MODE_GCM, nonce=nonce)
+    dec_data = cipher_aes.decrypt_and_verify(ciphertext, tag)
+    
+    return dec_data.decode('utf-8')
+
+from fastapi import Depends, HTTPException, status
+from fastapi.security import OAuth2PasswordBearer
+from sqlalchemy.future import select
+from sqlalchemy.ext.asyncio import AsyncSession
+from jose import JWTError
+from . import database, models
+
+oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
+
+async def get_current_user(token: str = Depends(oauth2_scheme), db: AsyncSession = Depends(database.get_db)):
+    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
+    
+    result = await db.execute(select(models.User).where(models.User.username == username))
+    user = result.scalars().first()
+    if user is None:
+        raise credentials_exception
+    return user
--- a/app/database.py
+++ b/app/database.py
@@ -0,0 +1,50 @@
+# database.py
+import os
+from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
+from sqlalchemy.orm import sessionmaker, declarative_base
+
+# 优先从环境变量读取（Docker 部署推荐）
+SQLALCHEMY_DATABASE_URL = os.getenv("DATABASE_URL", "postgresql+asyncpg://user:password@db:5432/fastapi_db")
+
+engine = create_async_engine(SQLALCHEMY_DATABASE_URL, echo=True)
+AsyncSessionLocal = sessionmaker(bind=engine, class_=AsyncSession, expire_on_commit=False)
+
+Base = declarative_base()
+
+async def get_db():
+    async with AsyncSessionLocal() as session:
+        yield session
+        # 注意：通常不在 get_db 里统一 commit，建议在 endpoint 里手动 commit
+
+async def init_db():
+    async with engine.begin() as conn:
+        # 导入模型以确保 metadata 注册了表
+        from . import models 
+        # 自动创建表
+        await conn.run_sync(Base.metadata.create_all)
+    
+    # 创建默认管理员用户
+    async with AsyncSessionLocal() as session:
+        from . import auth
+        from sqlalchemy.future import select
+        
+        # 检查是否已存在 admin 用户
+        result = await session.execute(
+            select(models.User).where(models.User.username == "admin")
+        )
+        existing_admin = result.scalars().first()
+        
+        if not existing_admin:
+            # 创建管理员用户
+            private_key, public_key = auth.generate_key_pair()
+            admin_user = models.User(
+                username="admin",
+                hashed_password=auth.hash_password("admin123"),
+                private_key=private_key,
+                public_key=public_key,
+                is_admin=True,
+                guale=False
+            )
+            session.add(admin_user)
+            await session.commit()
+            print("✅ Default admin user created (username: admin, password: admin123)")
--- a/app/main.py
+++ b/app/main.py
@@ -0,0 +1,214 @@
+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
+
+@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}"}
+
--- a/app/models.py
+++ b/app/models.py
@@ -0,0 +1,37 @@
+from sqlalchemy import Column, Integer, String, ForeignKey, Text, Table, Boolean
+from sqlalchemy.orm import relationship
+
+from .database import Base
+
+
+class User(Base):
+    __tablename__ = "users"
+
+    id = Column(Integer, primary_key=True, index=True)
+    username = Column(String, unique=True, index=True)
+    hashed_password = Column(String)
+    # System keys
+    public_key = Column(String)
+    private_key = Column(String)  # Encrypted or raw? Storing raw for now as per req
+    is_admin = Column(Boolean, default=False)
+    guale = Column(Boolean, default=False)
+
+    assets = relationship("Asset", foreign_keys="Asset.author_id", back_populates="author")
+    inherited_assets = relationship("Asset", foreign_keys="Asset.heir_id", back_populates="heir")
+
+
+
+class Asset(Base):
+    __tablename__ = "assets"
+
+    id = Column(Integer, primary_key=True, index=True)
+    title = Column(String, index=True)
+    content_outer_encrypted = Column(Text)
+    author_id = Column(Integer, ForeignKey("users.id"))
+    heir_id = Column(Integer, ForeignKey("users.id"))
+
+    # Key shard for this asset
+    private_key_shard = Column(String)
+
+    author = relationship("User", foreign_keys=[author_id], back_populates="assets")
+    heir = relationship("User", foreign_keys=[heir_id], back_populates="inherited_assets")
--- a/app/note
+++ b/app/note
@@ -0,0 +1 @@
+sudo docker compose exec web python reset_db.py
--- a/app/schemas.py
+++ b/app/schemas.py
@@ -0,0 +1,62 @@
+from pydantic import BaseModel, ConfigDict
+from typing import List, Optional
+
+# Heir Schemas
+class HeirBase(BaseModel):
+    name: str
+
+class HeirCreate(HeirBase):
+    pass
+
+class HeirOut(HeirBase):
+    id: int
+    user_id: int
+    model_config = ConfigDict(from_attributes=True)
+
+# User Schemas
+class UserCreate(BaseModel):
+    username: str
+    password: str
+
+class UserLogin(BaseModel):
+    username: str
+    password: str
+
+class UserOut(BaseModel):
+    id: int
+    username: str
+    public_key: Optional[str] = None
+    is_admin: bool = False
+    guale: bool = False
+    #heirs: List[HeirOut] = []
+    model_config = ConfigDict(from_attributes=True)
+
+# Asset Schemas (renamed from Article)
+class AssetBase(BaseModel):
+    title: str
+
+class AssetCreate(AssetBase):
+    private_key_shard: str
+    content_inner_encrypted: str
+
+class AssetOut(AssetBase):
+    id: int
+    author_id: int
+    private_key_shard: str
+    content_outer_encrypted: str
+    model_config = ConfigDict(from_attributes=True)
+
+class AssetClaim(BaseModel):
+    asset_id: int
+    private_key_shard: str
+
+class AssetClaimOut(AssetClaim):
+    id: int
+    result: str
+
+class AssetAssign(BaseModel):
+    asset_id: int
+    heir_name: str
+
+class DeclareGuale(BaseModel):
+    username: str
				`@@ -0,0 +1 @@`
				`sudo docker compose exec web python reset_db.py`