main.py

import os
import asyncio
import json
import logging
from telegram import Bot, Update
from telegram.ext import Application, CommandHandler, MessageHandler
import httpx

from pikpakapi import PikPakApi

from typing import Union, Any, Dict, List, Optional
from fastapi import (
    FastAPI,
    APIRouter,
    Depends,
    Request,
    Query,
    Body,
    Path,
    Response,
    HTTPException,
    status,
    Request,
)
from fastapi.responses import StreamingResponse, HTMLResponse, JSONResponse
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from fastapi.templating import Jinja2Templates
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Extra


class PostRequest(BaseModel):
    class Config:
        extra = Extra.allow


class FileRequest(BaseModel):
    size: int = 100
    parent_id: str | None = ""
    next_page_token: str | None = ""
    additional_filters: Dict | None = {}

    class Config:
        extra = Extra.allow


class OfflineRequest(BaseModel):
    file_url: str = ""
    parent_id: str | None = ""
    name: str | None = ""

    class Config:
        extra = Extra.allow


security = HTTPBearer()
# SECRET_TOKEN = "SECRET_TOKEN"
SECRET_TOKEN = os.getenv("SECRET_TOKEN")
if SECRET_TOKEN is None:
    raise ValueError("请在环境变量中设置SECRET_TOKEN，确保安全!")

THUNDERX_USERNAME = os.getenv("THUNDERX_USERNAME")
if THUNDERX_USERNAME is None:
    raise ValueError("请在环境变量中设置THUNDERX_USERNAME，用户名【邮箱】用来登陆!")


THUNDERX_PASSWORD = os.getenv("THUNDERX_PASSWORD")
if THUNDERX_PASSWORD is None:
    raise ValueError("请在环境变量中设置THUNDERX_PASSWORD，密码用来登陆!")

PROXY_URL = os.getenv("PROXY_URL")
TG_BOT_TOKEN = os.getenv("TG_BOT_TOKEN")
TG_WEBHOOK_URL = os.getenv("TG_WEBHOOK_URL")


async def verify_token(
    request: Request, credentials: HTTPAuthorizationCredentials = Depends(security)
):
    # excluded_paths = ["/"]  # 需要排除的路径列表
    # if request.url.path in excluded_paths:
    #     return  # 直接跳过验证

    # 验证Bearer格式
    if credentials.scheme != "Bearer":
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid authentication scheme",
        )

    # 验证令牌内容
    if credentials.credentials != SECRET_TOKEN:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid or expired token"
        )


app = FastAPI()


app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

api_router = APIRouter(dependencies=[Depends(verify_token)])
front_router = APIRouter()

templates = Jinja2Templates(
    directory="templates", variable_start_string="{[", variable_end_string="]}"
)


async def log_token(THUNDERX_CLIENT, extra_data):
    logging.info(f"Token: {THUNDERX_CLIENT.encoded_token}, Extra Data: {extra_data}")


THUNDERX_CLIENT = None
TG_BOT_APPLICATION = None
TG_BASE_URL = "https://tg.alist.dpdns.org/bot"


###################TG机器人功能区###################
# 定义命令处理函数
async def start(update: Update, context):
    commands = (
        "欢迎使用我的机器人！\n\n" "以下是您可以使用的命令：\n" "/help - 获取帮助信息\n"
    )
    await update.message.reply_text(commands)


async def help(update: Update, context):
    await update.message.reply_text("这是一个帮助信息！")


@app.on_event("startup")
async def init_client():
    global THUNDERX_CLIENT
    global TG_BOT_APPLICATION
    if not os.path.exists("thunderx.txt"):
        THUNDERX_CLIENT = PikPakApi(
            username=THUNDERX_USERNAME,
            password=THUNDERX_PASSWORD,
            httpx_client_args=None,
            token_refresh_callback=log_token,
            token_refresh_callback_kwargs={"extra_data": "test"},
        )
        await THUNDERX_CLIENT.login()
        await THUNDERX_CLIENT.refresh_access_token()
        with open("thunderx.json", "w") as f:
            f.write(json.dumps(THUNDERX_CLIENT.to_dict(), indent=4))
    else:
        with open("thunderx.txt", "r") as f:
            data = json.load(f)
            THUNDERX_CLIENT = PikPakApi.from_dict(data)
            # await client.refresh_access_token()
            print(json.dumps(THUNDERX_CLIENT.get_user_info(), indent=4))

            print(
                json.dumps(
                    await THUNDERX_CLIENT.events(),
                    indent=4,
                )
            )

    if TG_BOT_TOKEN is None:
        print("未设置TG_BOT_TOKEN无法实现TG机器人功能！")
    else:
        TG_BOT_APPLICATION = (
            Application.builder().base_url(TG_BASE_URL).token(TG_BOT_TOKEN).build()
        )
        # await TG_BOT_APPLICATION.bot.delete_webhook()
        await TG_BOT_APPLICATION.bot.set_webhook(TG_WEBHOOK_URL)
        await TG_BOT_APPLICATION.initialize()
        # 将命令处理函数添加到 dispatcher
        TG_BOT_APPLICATION.add_handler(CommandHandler("start", start))
        TG_BOT_APPLICATION.add_handler(CommandHandler("help", help))


# FastAPI 路由：接收来自 Telegram 的 Webhook 回调
@app.post("/webhook")
async def webhook(request: Request):
    # 从请求获取 JSON 数据
    data = await request.json()

    # 将 Telegram Update 转换为 Update 对象
    update = Update.de_json(data, TG_BOT_APPLICATION.bot)

    # 将 Update 对象传递给 Application 进行处理
    await TG_BOT_APPLICATION.process_update(update)

    return JSONResponse({"status": "ok"})


@front_router.get(
    "/",
    response_class=HTMLResponse,
    summary="前台页面",
    description="前台管理页面，需要在设置里设置SECRET_TOKEN才能正常请求",
    tags=["前端"],
)
async def home(request: Request):
    return templates.TemplateResponse("index.html", {"request": request})


@api_router.post(
    "/files", summary="文件列表", description="获取文件列表", tags=["文件"]
)
async def get_files(item: FileRequest):
    return await THUNDERX_CLIENT.file_list(
        item.size, item.parent_id, item.next_page_token, item.additional_filters
    )


@api_router.post(
    "/file_star_list", summary="加星文件列表", description="加星文件列表", tags=["文件"]
)
async def file_star_list(
    size: int = Query(default=100, title="显示数量", description="显示数量"),
    next_page_token: str | None = Query(
        default=None, title="分页Token", description="分页Token"
    ),
):
    return await THUNDERX_CLIENT.file_star_list(size, next_page_token)


@api_router.get(
    "/files/{file_id}", summary="文件信息", description="获取文件信息", tags=["文件"]
)
async def get_file_info(file_id: str = Path(..., title="文件ID", description="文件ID")):
    return await THUNDERX_CLIENT.get_download_url(file_id)


@api_router.delete(
    "/files/{file_id}", summary="删除文件", description="删除文件", tags=["文件"]
)
async def delete_file_info(
    file_id: str = Path(..., title="文件ID", description="文件ID")
):
    return await THUNDERX_CLIENT.delete_forever([file_id])


@api_router.post(
    "/file_rename", summary="重命名文件", description="重命名文件", tags=["文件"]
)
async def file_rename(
    file_id: str = Query(title="文件ID", description="文件ID"),
    new_file_name: str = Query(title="新文件名", description="新文件名"),
):
    return await THUNDERX_CLIENT.file_rename(file_id, new_file_name)


@api_router.post(
    "/file_batch_copy",
    summary="批量复制文件",
    description="批量复制文件",
    tags=["文件"],
)
async def file_batch_copy(
    ids: List[str] = Body(title="文件ID列表", description="文件ID列表"),
    to_parent_id: str = Query(
        title="复制到的文件夹id, 默认为根目录",
        description="复制到的文件夹id, 默认为根目录",
    ),
):
    return await THUNDERX_CLIENT.file_batch_copy(ids, to_parent_id)


@api_router.post(
    "/file_batch_move",
    summary="批量移动文件",
    description="批量移动文件",
    tags=["文件"],
)
async def file_batch_move(
    ids: List[str] = Body(title="文件ID列表", description="文件ID列表"),
    to_parent_id: str = Query(
        title="移动到的文件夹id, 默认为根目录",
        description="移动到的文件夹id, 默认为根目录",
    ),
):
    return await THUNDERX_CLIENT.file_batch_move(ids, to_parent_id)


@api_router.post(
    "/create_folder", summary="新建文件夹", description="新建文件夹", tags=["文件"]
)
async def create_folder(
    name: str = Query(title="文件夹名称", description="文件夹名称"),
    parent_id: str = Query(
        title="父文件夹id, 默认创建到根目录", description="父文件夹id, 默认创建到根目录"
    ),
):
    return await THUNDERX_CLIENT.create_folder(name, parent_id)


@api_router.post(
    "/delete_to_trash",
    summary="将文件夹、文件移动到回收站",
    description="将文件夹、文件移动到回收站",
    tags=["文件"],
)
async def delete_to_trash(
    ids: List[str] = Body(title="文件ID列表", description="文件ID列表")
):
    return await THUNDERX_CLIENT.delete_to_trash(ids)


@api_router.post(
    "/delete_forever",
    summary="将文件夹、文件彻底删除",
    description="将文件夹、文件彻底删除",
    tags=["文件"],
)
async def delete_forever(
    ids: List[str] = Body(title="文件ID列表", description="文件ID列表")
):
    return await THUNDERX_CLIENT.delete_forever(ids)


@api_router.post(
    "/untrash",
    summary="将文件夹、文件移出回收站",
    description="将文件夹、文件移出回收站",
    tags=["文件"],
)
async def untrash(ids: List[str] = Body(title="文件ID列表", description="文件ID列表")):
    return await THUNDERX_CLIENT.untrash(ids)


@api_router.post(
    "/file_batch_star",
    summary="批量给文件加星标",
    description="批量给文件加星标",
    tags=["文件"],
)
async def file_batch_star(
    ids: List[str] = Body(title="文件ID列表", description="文件ID列表")
):
    return await THUNDERX_CLIENT.file_batch_star(ids)


@api_router.post(
    "/file_batch_unstar",
    summary="批量给文件加星标",
    description="批量给文件加星标",
    tags=["文件"],
)
async def file_batch_unstar(
    ids: List[str] = Body(title="文件ID列表", description="文件ID列表")
):
    return await THUNDERX_CLIENT.file_batch_unstar(ids)


@api_router.post(
    "/emptytrash", summary="清空回收站", description="清空回收站【慎用】", tags=["文件"]
)
async def emptytrash():
    return await THUNDERX_CLIENT.emptytrash()


##############  分享 ################
@api_router.post(
    "/get_share_list",
    summary="获取账号分享列表",
    description="获取账号分享列表",
    tags=["分享"],
)
async def get_share_list(
    page_token: str | None = Query(
        default=None, title="分页Token", description="分页Token"
    )
):
    return await THUNDERX_CLIENT.get_share_list(page_token)


@api_router.post(
    "/file_batch_share", summary="创建分享", description="创建分享", tags=["分享"]
)
async def file_batch_share(
    ids: List[str] = Body(default=None, title="文件ID列表", description="文件ID列表"),
    need_password: bool | None = Query(
        default=False, title="是否需要密码", description="是否需要密码"
    ),
    expiration_days: int | None = Query(
        default=-1, title="过期时间", description="过期时间【天数，默认永远】"
    ),
):
    return await THUNDERX_CLIENT.file_batch_share(ids, need_password, expiration_days)


@api_router.post(
    "/share_batch_delete", summary="取消分享", description="取消分享", tags=["分享"]
)
async def share_batch_delete(
    ids: List[str] = Body(title="文件ID列表", description="文件ID列表")
):
    return await THUNDERX_CLIENT.share_batch_delete(ids)


@api_router.post(
    "/get_share_folder",
    summary="获取分享信息",
    description="获取分享信息",
    tags=["分享"],
)
async def get_share_folder(
    share_id: str = Query(title="分享ID", description="分享ID"),
    pass_code_token: str | None = Query(default=None, title="密码", description="密码"),
    parent_id: str | None = Query(default=None, title="父ID", description="父ID"),
):
    return await THUNDERX_CLIENT.get_share_folder(share_id, pass_code_token, parent_id)


@api_router.post(
    "/restore", summary="转存分享文件", description="转存分享文件", tags=["分享"]
)
async def restore(
    share_id: str, pass_code_token: str | None = None, file_ids: List[str] | None = None
):
    return await THUNDERX_CLIENT.restore(share_id, pass_code_token, file_ids)


##############  离线任务 ################


@api_router.get(
    "/offline", summary="离线任务列表", description="离线任务列表", tags=["离线任务"]
)
async def offline_list(size: int = 10000, next_page_token: str | None = None):
    return await THUNDERX_CLIENT.offline_list(
        size=size,
        next_page_token=next_page_token,
        phase=None,
    )


@api_router.post(
    "/offline", summary="添加离线任务", description="添加离线任务", tags=["离线任务"]
)
async def offline(item: OfflineRequest):
    return await THUNDERX_CLIENT.offline_download(
        item.file_url, item.parent_id, item.name
    )


@api_router.post(
    "/delete_tasks",
    summary="删除离线任务",
    description="删除离线任务",
    tags=["离线任务"],
)
async def delete_tasks(task_ids: List[str], delete_files: bool = False):
    return await THUNDERX_CLIENT.delete_tasks(task_ids, delete_files)


##############  账号 ################
@api_router.get(
    "/userinfo", summary="用户信息", description="获取用户登陆信息", tags=["账号"]
)
async def userinfo():
    return THUNDERX_CLIENT.get_user_info()


@api_router.get(
    "/quota", summary="空间使用信息", description="获取空间使用信息", tags=["账号"]
)
async def quota_info():
    return await THUNDERX_CLIENT.get_quota_info()


@api_router.get(
    "/invite_code", summary="查看邀请码", description="查看邀请码", tags=["账号"]
)
async def get_invite_code():
    return await THUNDERX_CLIENT.get_invite_code()


app.include_router(front_router)
app.include_router(api_router)