#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ 将单个账号的Cookie输入到MySQL数据库 支持手动输入Cookie信息或从剪贴板粘贴 """ import json import sys import os from datetime import datetime from typing import Dict, 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 SingleCookieToDB: """单个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.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 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 cookie_string_to_dict(self, cookie_string: str) -> Dict: """ 将Cookie字符串转换为字典格式 Args: cookie_string: Cookie字符串,格式: "key1=value1; key2=value2" Returns: Cookie字典 """ cookies = {} for item in cookie_string.split(';'): item = item.strip() if '=' in item: key, value = item.split('=', 1) cookies[key.strip()] = value.strip() return cookies 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 update_author_cookie(self, author_id: int, cookie_string: str, app_id: Optional[str] = None) -> bool: """ 更新作者的Cookie信息 Args: author_id: 作者ID cookie_string: Cookie字符串 app_id: 百家号app_id(可选) Returns: 是否更新成功 """ try: # 构建更新SQL update_fields = ["toutiao_cookie = %s", "updated_at = NOW()"] params = [cookie_string] # 如果提供了app_id,也一并更新 if app_id: update_fields.append("app_id = %s") params.append(app_id) params.append(author_id) sql = f"UPDATE ai_authors SET {', '.join(update_fields)} WHERE id = %s" self.db_manager.execute_update(sql, tuple(params)) logger.info(f"成功更新作者ID={author_id}的Cookie") return True except Exception as e: logger.error(f"更新Cookie失败: {e}", exc_info=True) 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) -> bool: """ 插入新作者记录 Args: author_name: 作者名称(用于数据库author_name字段) cookie_string: Cookie字符串 app_id: 百家号app_id nick: 昵称 domain: 领域 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) logger.info(f"成功创建新作者: {author_name}") return True except Exception as e: logger.error(f"插入新作者失败: {e}", exc_info=True) print(f"[X] 插入新作者失败: {e}") return False def add_cookie(self, account_info: Dict, auto_create: bool = True) -> bool: """ 添加单个账号的Cookie到数据库 Args: account_info: 账号信息字典,包含cookies、username、nick等字段 auto_create: 当作者不存在时是否自动创建,默认True Returns: 是否添加成功 """ # 提取Cookie信息 cookies = account_info.get('cookies', {}) if not cookies: print("[X] Cookie信息为空") return False # 转换Cookie为字符串(如果是字典格式) if isinstance(cookies, dict): cookie_string = self.cookie_dict_to_string(cookies) else: cookie_string = str(cookies) # 提取其他信息(使用username和nick作为author_name进行匹配) username = account_info.get('username', '').strip() nick = account_info.get('nick', '').strip() app_id = account_info.get('app_id', '').strip() domain = account_info.get('domain', '').strip() # 验证username或nick至少有一个存在 if not username and not nick: print("[X] username和nick至少需要提供一个") return False print(f"\n账号信息:") print(f" Username: {username}") print(f" 昵称: {nick}") print(f" App ID: {app_id}") print(f" 领域: {domain}") # 查找作者(使用双重匹配机制:先username,后nick) channel = 1 # 百家号固定为channel=1 author = None matched_field = None # 1. 首先尝试使用username匹配 if username: author = self.find_author_by_name(username, channel) if author: matched_field = 'username' print(f"\n[√] 通过username匹配到作者: {author['author_name']} (ID: {author['id']}, Channel: {author['channel']})") # 2. 如果username匹配失败,尝试使用nick匹配 if not author and nick: author = self.find_author_by_name(nick, channel) if author: matched_field = 'nick' print(f"\n[√] 通过nick匹配到作者: {author['author_name']} (ID: {author['id']}, Channel: {author['channel']})") # 3. 如果都没匹配到 if not author: print(f"\n[!] 未找到匹配的作者(已尝试username和nick)") # 更新或创建 if author: # 更新现有记录 print(f"\n正在更新作者Cookie...") success = self.update_author_cookie( author['id'], cookie_string, app_id if app_id else None ) if success: print(f"[OK] Cookie已更新(匹配字段: {matched_field})") return True else: print(f"[X] Cookie更新失败") return False else: # 作者不存在,考虑创建 if auto_create: # 优先使用username,如果没有则使用nick author_name_to_create = username if username else nick print(f"\n正在创建新作者(author_name: {author_name_to_create})...") success = self.insert_new_author( author_name_to_create, cookie_string, app_id, nick, domain ) if success: print(f"[OK] 新作者已创建 (author_name: {author_name_to_create})") return True else: print(f"[X] 创建作者失败") return False else: print(f"[X] 作者不存在,且未开启自动创建") return False def run_interactive(self): """交互式运行模式""" print("\n" + "="*70) print("添加单个账号Cookie到数据库") print("="*70) # 连接数据库 if not self.connect_db(): logger.error("数据库连接失败,退出") return try: # 询问是否自动创建不存在的作者 print("\n当作者不存在时是否自动创建?") auto_create_input = input("(y/n, 默认y): ").strip().lower() auto_create = auto_create_input != 'n' # 输入账号信息 print("\n" + "="*70) print("请输入账号信息:") print("="*70) username = input("\n1. Username (用于匹配数据库author_name): ").strip() nick = input("2. 昵称 (备用匹配字段): ").strip() app_id = input("3. App ID (可选): ").strip() domain = input("4. 领域 (可选): ").strip() # 输入Cookie print("\n" + "="*70) print("请输入Cookie信息:") print("提示: 可以输入以下任意格式:") print(" 1. Cookie字符串: key1=value1; key2=value2") print(" 2. JSON格式: {\"key1\": \"value1\", \"key2\": \"value2\"}") print(" 3. 多行输入,输入完成后输入 END 结束") print("="*70) cookie_lines = [] while True: line = input().strip() if line.upper() == 'END': break if line: cookie_lines.append(line) cookie_input = ' '.join(cookie_lines) # 解析Cookie cookies = {} if cookie_input.startswith('{'): # JSON格式 try: cookies = json.loads(cookie_input) except json.JSONDecodeError: print("[X] Cookie JSON格式解析失败") return else: # 字符串格式 cookies = self.cookie_string_to_dict(cookie_input) if not cookies: print("[X] Cookie为空,操作取消") return # 构建账号信息 account_info = { 'username': username, 'nick': nick, 'app_id': app_id, 'domain': domain, 'cookies': cookies } # 确认信息 print("\n" + "="*70) print("确认账号信息:") print("="*70) print(f" Username: {username}") print(f" 昵称: {nick}") print(f" App ID: {app_id}") print(f" 领域: {domain}") print(f" Cookie条目数: {len(cookies)}") print(f" 自动创建: {'是' if auto_create else '否'}") print("="*70) confirm = input("\n确认添加到数据库?(y/n): ").strip().lower() if confirm != 'y': print("\n已取消") return # 添加Cookie success = self.add_cookie(account_info, auto_create) if success: print("\n" + "="*70) print("添加成功!") print("="*70) else: print("\n" + "="*70) print("添加失败,请查看错误信息") print("="*70) 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使用默认数据库配置...") # 创建同步器并执行 syncer = SingleCookieToDB(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("="*70) # 运行交互式模式 syncer.run_interactive() if __name__ == '__main__': main()