#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ 将捕获的Cookie同步到MySQL数据库 从captured_account_cookies.json读取Cookie数据，更新到ai_authors表 """ import json import sys import os from datetime import datetime from typing import Dict, List, Optional # 导入统一的数据库管理器和日志配置 from database_config import DatabaseManager, DB_CONFIG from log_config import setup_cookie_sync_logger # 初始化日志记录器 logger = setup_cookie_sync_logger() # 设置UTF-8编码 if sys.platform == 'win32': import io sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8') sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8') class CookieSyncToDB: """Cookie同步到数据库""" def __init__(self, db_config: Optional[Dict] = None): """ 初始化数据库连接 Args: db_config: 数据库配置字典，默认使用database_config.DB_CONFIG """ self.script_dir = os.path.dirname(os.path.abspath(__file__)) self.cookies_file = os.path.join(self.script_dir, "captured_account_cookies.json") # 成功和失败记录文件 timestamp = datetime.now().strftime('%Y%m%d_%H%M%S') self.success_file = os.path.join(self.script_dir, f"sync_success_{timestamp}.json") self.failed_file = os.path.join(self.script_dir, f"sync_failed_{timestamp}.json") # 使用统一的数据库管理器 self.db_manager = DatabaseManager(db_config) self.db_config = self.db_manager.config def connect_db(self) -> bool: """连接数据库""" return self.db_manager.test_connection() def close_db(self): """关闭数据库连接""" print("[OK] 数据库操作完成") def load_cookies_from_file(self) -> Dict: """从JSON文件加载Cookie数据""" try: if not os.path.exists(self.cookies_file): print(f"[X] Cookie文件不存在: {self.cookies_file}") return {} with open(self.cookies_file, 'r', encoding='utf-8') as f: data = json.load(f) print(f"[OK] 从文件加载了 {len(data)} 个账号的Cookie") return data except Exception as e: print(f"[X] 加载Cookie文件失败: {e}") return {} def cookie_dict_to_string(self, cookies: Dict) -> str: """ 将Cookie字典转换为字符串格式 Args: cookies: Cookie字典 Returns: Cookie字符串，格式: "key1=value1; key2=value2" """ return '; '.join([f"{k}={v}" for k, v in cookies.items()]) def save_sync_records(self, success_records: List[Dict], failed_records: List[Dict]): """ 保存同步结果记录 Args: success_records: 成功记录列表 failed_records: 失败记录列表 """ try: # 保存成功记录 if success_records: with open(self.success_file, 'w', encoding='utf-8') as f: json.dump(success_records, f, ensure_ascii=False, indent=2) logger.info(f"成功记录已保存到: {self.success_file}") print(f"\n[OK] 成功记录已保存: {self.success_file}") # 保存失败记录 if failed_records: with open(self.failed_file, 'w', encoding='utf-8') as f: json.dump(failed_records, f, ensure_ascii=False, indent=2) logger.info(f"失败记录已保存到: {self.failed_file}") print(f"[OK] 失败记录已保存: {self.failed_file}") except Exception as e: logger.error(f"保存记录文件失败: {e}", exc_info=True) print(f"[X] 保存记录文件失败: {e}") def find_author_by_name(self, author_name: str, channel: int = 1) -> Optional[Dict]: """ 根据作者名称和渠道查找数据库记录 Args: author_name: 作者名称 channel: 渠道（1=百家号，默认1） Returns: 作者记录字典，未找到返回None """ try: sql = "SELECT * FROM ai_authors WHERE author_name = %s AND channel = %s LIMIT 1" result = self.db_manager.execute_query(sql, (author_name, channel), fetch_one=True) return result except Exception as e: print(f"[X] 查询作者失败: {e}") return None def find_author_by_app_id(self, app_id: str) -> Optional[Dict]: """ 根据app_id查找数据库记录 Args: app_id: 百家号app_id Returns: 作者记录字典，未找到返回None """ try: sql = "SELECT * FROM ai_authors WHERE app_id = %s LIMIT 1" result = self.db_manager.execute_query(sql, (app_id,), fetch_one=True) return result except Exception as e: print(f"[X] 查询作者失败: {e}") return None def update_author_cookie(self, author_id: int, cookie_string: str, app_id: Optional[str] = None, app_token: Optional[str] = None) -> bool: """ 更新作者的Cookie信息 Args: author_id: 作者ID cookie_string: Cookie字符串 app_id: 百家号app_id（可选） app_token: 百家号app_token（可选） Returns: 是否更新成功 """ try: # 构建更新SQL update_fields = ["toutiao_cookie = %s", "updated_at = NOW()"] params = [cookie_string] # 如果提供了app_id和app_token，也一并更新 if app_id: update_fields.append("app_id = %s") params.append(app_id) if app_token: update_fields.append("app_token = %s") params.append(app_token) params.append(author_id) sql = f"UPDATE ai_authors SET {', '.join(update_fields)} WHERE id = %s" self.db_manager.execute_update(sql, tuple(params)) return True except Exception as e: print(f"[X] 更新Cookie失败: {e}") return False def insert_new_author(self, author_name: str, cookie_string: str, app_id: Optional[str] = None, nick: Optional[str] = None, domain: Optional[str] = None, level: Optional[str] = None) -> bool: """ 插入新作者记录 Args: author_name: 作者名称 cookie_string: Cookie字符串 app_id: 百家号app_id nick: 昵称 domain: 领域 level: 等级 Returns: 是否插入成功 """ try: # 构建插入SQL sql = """ INSERT INTO ai_authors (author_name, app_id, app_token, department_id, department_name, department, toutiao_cookie, channel, status, created_at, updated_at) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, NOW(), NOW()) """ # 参数 params = ( author_name, app_id or '', '', # app_token 暂时为空 0, # department_id 默认0 domain or '其它', # department_name 使用领域 '', # department 暂时为空 cookie_string, 1, # channel: 1=baidu 'active' # status ) self.db_manager.execute_update(sql, params) return True except Exception as e: print(f"[X] 插入新作者失败: {e}") return False def sync_cookies(self, auto_create: bool = True) -> Dict: """ 同步Cookie到数据库 Args: auto_create: 当作者不存在时是否自动创建，默认True Returns: 同步统计信息 """ stats = { 'total': 0, 'updated': 0, 'created': 0, 'failed': 0, 'skipped': 0 } # 成功和失败记录 success_records = [] failed_records = [] # 加载Cookie数据 cookies_data = self.load_cookies_from_file() if not cookies_data: print("[X] 没有可同步的Cookie数据") return stats stats['total'] = len(cookies_data) print("\n" + "="*70) print("开始同步Cookie到数据库") print("="*70) # 遍历每个账号 for idx, (account_name, account_info) in enumerate(cookies_data.items(), 1): print(f"\n[{idx}/{stats['total']}] 处理账号: {account_name}") # 提取Cookie信息 cookies = account_info.get('cookies', {}) if not cookies: print(" [!] 该账号没有Cookie，跳过") stats['skipped'] += 1 continue # 转换Cookie为字符串 cookie_string = self.cookie_dict_to_string(cookies) # 提取其他信息 app_id = account_info.get('app_id', '') nick = account_info.get('nick', '') domain = account_info.get('domain', '') level = account_info.get('level', '') # 查找作者（使用 author_name + channel 作为唯一键） channel = 1 # 百家号固定为channel=1 author = self.find_author_by_name(account_name, channel) if author: print(f" [√] 找到作者: {author['author_name']} (ID: {author['id']}, Channel: {author['channel']})") # 更新或创建 if author: # 更新现有记录 success = self.update_author_cookie( author['id'], cookie_string, app_id if app_id else None, None # app_token暂不更新 ) if success: print(f" [OK] Cookie已更新") stats['updated'] += 1 # 记录成功 success_records.append({ 'account_name': account_name, 'app_id': app_id, 'nick': nick, 'domain': domain, 'action': 'updated', 'db_author_id': author['id'], 'db_author_name': author['author_name'] }) else: print(f" [X] Cookie更新失败") stats['failed'] += 1 # 记录失败 failed_records.append({ 'account_name': account_name, 'app_id': app_id, 'nick': nick, 'reason': '数据库更新失败', 'db_author_id': author['id'] }) else: # 作者不存在 if auto_create: print(f" [*] 作者不存在，创建新记录...") success = self.insert_new_author( account_name, cookie_string, app_id, nick, domain, level ) if success: print(f" [OK] 新作者已创建") stats['created'] += 1 # 记录成功 success_records.append({ 'account_name': account_name, 'app_id': app_id, 'nick': nick, 'domain': domain, 'action': 'created' }) else: print(f" [X] 创建作者失败") stats['failed'] += 1 # 记录失败 failed_records.append({ 'account_name': account_name, 'app_id': app_id, 'nick': nick, 'reason': '数据库插入失败' }) else: print(f" [!] 作者不存在，跳过（auto_create=False）") stats['skipped'] += 1 # 记录失败（数据库中不存在） failed_records.append({ 'account_name': account_name, 'app_id': app_id, 'nick': nick, 'reason': '数据库中不存在该账号，且未开启自动创建' }) # 保存记录文件 self.save_sync_records(success_records, failed_records) return stats def run(self, auto_create: bool = True): """ 执行同步任务 Args: auto_create: 是否自动创建不存在的作者 """ logger.info("\n" + "="*70) logger.info("Cookie同步到数据库工具") logger.info("="*70) print("\n" + "="*70) print("Cookie同步到数据库工具") print("="*70) # 连接数据库 if not self.connect_db(): logger.error("数据库连接失败，退出") return try: # 同步Cookie logger.info(f"开始同步Cookie，自动创建={auto_create}") stats = self.sync_cookies(auto_create=auto_create) # 显示统计 print("\n" + "="*70) print("同步完成") print("="*70) print(f" 总账号数: {stats['total']}") print(f" 已更新: {stats['updated']}") print(f" 已创建: {stats['created']}") print(f" 失败: {stats['failed']}") print(f" 跳过: {stats['skipped']}") print("="*70 + "\n") logger.info("同步结果:") logger.info(f" 总账号数: {stats['total']}") logger.info(f" 已更新: {stats['updated']}") logger.info(f" 已创建: {stats['created']}") logger.info(f" 失败: {stats['failed']}") logger.info(f" 跳过: {stats['skipped']}") finally: # 关闭数据库连接 self.close_db() def main(): """主函数""" print("\n" + "="*70) print("Cookie同步到数据库配置") print("="*70) # 使用默认配置还是自定义配置 print("\n请选择数据库配置方式:") print(" 1. 使用默认配置 (8.149.233.36/ai_statistics_read)") print(" 2. 自定义配置") choice = input("\n请选择 (1/2, 默认1): ").strip() or '1' if choice == '2': # 自定义数据库配置 print("\n请输入数据库连接信息:\n") host = input("数据库地址: ").strip() port = input("端口 (默认: 3306): ").strip() or '3306' user = input("用户名: ").strip() password = input("密码: ").strip() database = input("数据库名: ").strip() db_config = { 'host': host, 'port': int(port), 'user': user, 'password': password, 'database': database, 'charset': 'utf8mb4' } else: # 使用默认配置 db_config = None print("\n使用默认数据库配置...") # 询问是否自动创建不存在的作者 print("\n" + "="*70) auto_create_input = input("当作者不存在时是否自动创建？(y/n, 默认y): ").strip().lower() auto_create = auto_create_input != 'n' # 创建同步器并执行 syncer = CookieSyncToDB(db_config) print("\n配置确认:") print(f" 数据库: {syncer.db_config['host']}:{syncer.db_config.get('port', 3306)}/{syncer.db_config['database']}") print(f" 用户: {syncer.db_config['user']}") print(f" 自动创建: {'是' if auto_create else '否'}") print("="*70) confirm = input("\n确认开始同步？(y/n): ").strip().lower() if confirm != 'y': print("\n已取消") return syncer.run(auto_create=auto_create) if __name__ == '__main__': main()