open-douyin/api/updload.py

import hashlib
import os
import pathlib
import re
import sys
sys.path.append(os.path.dirname(os.path.dirname(__file__)))
import jwt
from fastapi import FastAPI,APIRouter,  File, HTTPException, Depends, Request,Header, UploadFile  
from fastapi.responses import JSONResponse 
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials  
import aiofiles
from starlette.status import HTTP_401_UNAUTHORIZED 
from api.jwt import verify_jwt_token,get_current_user
from config import *
from db.docs import DocumentsRepository

upload_router = APIRouter()  
security = HTTPBearer()
@upload_router.post('/upload')
async def upload(open_id=Depends(verify_jwt_token),file: UploadFile = File(...)  ):  
    
    res = await save_to_user_dir(open_id, file)
    if res:
        # 保存文件到本地
        DocumentsRepository()
        return {"message": "upload success"}
    else:  
        # 保存文件失败，返回 400 Bad Request 或其他适当的错误状态码  
        raise HTTPException(status_code=400, detail="upload fail")

def is_valid_filename(s):  
    """检查字符串是否只包含合法的文件名字符"""  
    # 这里的正则表达式只允许字母、数字、下划线、点、破折号和空格  
    # 你可以根据需要调整这个正则表达式  
    return re.match(r"^[a-zA-Z0-9_.\- ]+$", s) is not None  
  
def get_user_dir(open_id):  
    if is_valid_filename(open_id):  
        # 如果open_id合法，直接使用它作为目录名  
        user_dir = os.path.join(DATA_DIR, open_id)  
    else:  
        # 否则，计算它的哈希值并用作目录名  
        hash_object = hashlib.md5(open_id.encode())  
        hex_dig = hash_object.hexdigest()[:8] # 只取前 8 个字符   
        user_dir = os.path.join(DATA_DIR, "hash8_" + hex_dig)  
      
    # 如果目录不存在，创建它  
    if not os.path.exists(user_dir):  
        os.makedirs(user_dir)  
      
    return user_dir  

    
async def save_to_user_dir(open_id, file:UploadFile):
    user_dir = get_user_dir(open_id)
    file_path = os.path.join(user_dir,"docs", file.filename)
      
    async with aiofiles.open(file_path, "wb") as buffer:  
        chunk = await file.read(8192)  
        while chunk:  
            await buffer.write(chunk)  
            chunk = await file.read(8192)  
        logger.info(f"{open_id} save to {file_path}")
        return True