core/cron.py

"""
定时任务模块 - 定期更新所有API密钥
"""
import json
import os
import time
import sched
import threading
import sqlite3
from core.update import update
from utils.db import get_db_connection
from config import API_KEYS_FILE

# 定义更新间隔（12小时，单位：秒）
UPDATE_INTERVAL = 12 * 60 * 60

def update_all_keys():
    """
    更新所有API密钥
    """
    print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] 开始更新所有API密钥...")
    
    # 从SQLite数据库获取所有密钥ID
    conn = get_db_connection()
    key_ids = []
    
    try:
        cursor = conn.cursor()
        cursor.execute('SELECT id FROM api_keys')
        rows = cursor.fetchall()
        key_ids = [row['id'] for row in rows]
    except sqlite3.Error as e:
        print(f"从数据库获取密钥时出错: {str(e)}")
        
        # 备用方案：尝试从JSON文件读取
        if os.path.exists(API_KEYS_FILE):
            try:
                with open(API_KEYS_FILE, "r", encoding="utf-8") as f:
                    data = json.load(f)
                key_ids = [key["id"] for key in data.get("api_keys", [])]
            except Exception as e:
                print(f"读取API密钥文件失败: {str(e)}")
                return
    finally:
        if conn:
            conn.close()
    
    # 如果没有找到密钥，则退出
    if not key_ids:
        print("没有找到API密钥，跳过更新")
        return
    
    # 逐个更新API密钥
    for key_id in key_ids:
        result = update(key_id)
        if result["success"]:
            print(f"  - 密钥 {key_id} 更新成功")
        else:
            print(f"  - 密钥 {key_id} 更新失败: {result['message']}")
    
    print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] 完成所有API密钥更新")

def run_scheduler():
    """
    运行定时任务
    """
    scheduler = sched.scheduler(time.time, time.sleep)
    
    def scheduled_update():
        update_all_keys()
        # 再次安排下一次更新
        scheduler.enter(UPDATE_INTERVAL, 1, scheduled_update, ())
    
    # 首次安排更新
    scheduler.enter(UPDATE_INTERVAL, 1, scheduled_update, ())
    
    print(f"定时任务已启动，每 {UPDATE_INTERVAL // 3600} 小时更新一次API密钥")
    scheduler.run()

if __name__ == "__main__":
    # 在单独的线程中运行定时任务，避免阻塞主线程
    scheduler_thread = threading.Thread(target=run_scheduler)
    scheduler_thread.daemon = True  # 设置为守护线程，当主线程退出时自动结束
    scheduler_thread.start()
    
    # 保持主线程运行（可选，根据实际需求）
    while True:
        time.sleep(60)