ehentai-fastapi/main.py

from fastapi import FastAPI, Request
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from fastapi.responses import JSONResponse, FileResponse
import uvicorn
import glob
import os
from pydantic import BaseModel
import step2
from utils import *

app = FastAPI(title="EH-Downloader", version="1.0.0")

# 在应用启动时检查并创建data文件夹和targets.txt，以及proxy.txt
@app.on_event("startup")
async def startup_event():
    # 检查并创建data文件夹
    data_dir = "data"
    if not os.path.exists(data_dir):
        os.makedirs(data_dir)
        print(f"创建目录: {data_dir}")
    
    # 检查并创建targets.txt文件
    targets_file = os.path.join(data_dir, "targets.txt")
    if not os.path.exists(targets_file):
        with open(targets_file, 'w', encoding='utf-8') as f:
            f.write("# 在这里添加目标URL，每行一个\n")
            f.write("# 示例:\n")
            f.write("https://e-hentai.org/g/3550066/47d6393550\n")
        print(f"创建文件: {targets_file}")
    else:
        print(f"文件已存在: {targets_file}")
    
    # 检查并创建proxy.txt文件
    proxy_file = "proxy.txt"
    if not os.path.exists(proxy_file):
        with open(proxy_file, 'w', encoding='utf-8') as f:
            f.write("127.0.0.1:7890\n")
        print(f"创建文件: {proxy_file}")
    else:
        print(f"文件已存在: {proxy_file}")

# 挂载静态文件和模板
app.mount("/static", StaticFiles(directory="static"), name="static")
templates = Jinja2Templates(directory="templates")

# favicon 路由
@app.get("/favicon.ico", include_in_schema=False)
async def favicon():
    return FileResponse("static/favicon.ico")

@app.get("/")
async def home(request: Request):
    """主页面"""
    # 读取proxy.txt中的代理列表
    proxies = []
    try:
        with open("proxy.txt", 'r', encoding='utf-8') as f:
            proxies = [line.strip() for line in f.readlines() if line.strip()]
    except Exception as e:
        print(f"读取proxy.txt失败: {e}")
        proxies = ["127.0.0.1:7890"]
    
    # 如果没有代理配置，使用默认值
    if not proxies:
        proxies = ["127.0.0.1:7890"]
    
    return templates.TemplateResponse("index.html", {
        "request": request,
        "proxies": proxies,
        "default_proxy": proxies[0] if proxies else "127.0.0.1:7890"
    })

@app.post("/load_urls")
async def load_urls():
    """读取 targets.txt 文件中的URL"""
    try:
        file_path = "data/targets.txt"
        
        # 读取文件内容
        with open(file_path, 'r', encoding='utf-8') as f:
            urls = [line.strip() for line in f.readlines() if line.strip()]
        
        # 过滤掉空行和注释行（以#开头的行）
        urls = [url for url in urls if url and not url.startswith('#')]
        
        if not urls:
            return JSONResponse({
                "success": True,
                "message": "targets.txt 文件为空，请在data/targets.txt中添加URL",
                "urls": []
            })
        
        return JSONResponse({
            "success": True,
            "message": f"成功读取 {len(urls)} 个URL",
            "urls": urls
        })
        
    except Exception as e:
        return JSONResponse({
            "success": False,
            "message": f"读取文件时出错: {str(e)}",
            "urls": []
        })

@app.post("/clear")
async def clear_output():
    """清除输出"""
    return JSONResponse({
        "success": True,
        "message": "输出已清除",
        "output": ""
    })

class ProxyRequest(BaseModel):
    proxy: str  # 修改为单个proxy字段

@app.post("/download_urls")
async def download_urls(req: ProxyRequest):
    # 解析proxy字符串为ip和port
    if ":" in req.proxy:
        ip, port = req.proxy.split(":", 1)
        proxy = f"http://{ip}:{port}"
    else:
        proxy = None
    msg = await run_step1(proxy)
    return JSONResponse({"success": True, "message": msg})

@app.post("/download_images")
async def download_images(req: ProxyRequest):
    # 解析proxy字符串为ip和port
    if ":" in req.proxy:
        ip, port = req.proxy.split(":", 1)
        proxy = f"http://{ip}:{port}"
    else:
        proxy = None
    msg = await run_step2(proxy)
    return JSONResponse({"success": True, "message": msg})

@app.post("/clean_files")
async def clean_files():
    """清理项目目录下的所有 .log 和 .json 文件"""
    try:
        deleted_files = []
        error_files = []
        
        # 查找当前目录及所有子目录中的 .log 和 .json 文件
        patterns = ["**/*.log", "**/*.json"]
        
        for pattern in patterns:
            for file_path in glob.glob(pattern, recursive=True):
                try:
                    # 跳过 data/targets.txt 文件，因为这是配置文件
                    if file_path == "data/targets.txt":
                        continue
                    
                    os.remove(file_path)
                    deleted_files.append(file_path)
                    print(f"已删除文件: {file_path}")
                except Exception as e:
                    error_files.append(f"{file_path}: {str(e)}")
                    print(f"删除文件失败 {file_path}: {str(e)}")
        
        if error_files:
            return JSONResponse({
                "success": False,
                "message": f"清理完成，但部分文件删除失败",
                "deleted_count": len(deleted_files),
                "error_count": len(error_files),
                "deleted_files": deleted_files,
                "error_files": error_files
            })
        else:
            return JSONResponse({
                "success": True,
                "message": f"成功清理 {len(deleted_files)} 个文件",
                "deleted_count": len(deleted_files),
                "error_count": 0,
                "deleted_files": deleted_files
            })
            
    except Exception as e:
        return JSONResponse({
            "success": False,
            "message": f"清理过程中出错: {str(e)}",
            "deleted_count": 0,
            "error_count": 0
        })

@app.post("/check_incomplete")
async def check_incomplete():
    result = await step2.scan_tasks()

    """检查未完成文件"""
    return JSONResponse({
        "success": True,
        "message": "检查未完成文件功能已就绪",
        "data": f"共 {len(result)} 个文件未下载"
    })

if __name__ == "__main__":
    uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True)