#!/usr/bin/env python3 """ 数据库初始化脚本 用于自动创建MySQL数据库、用户和表结构，并配置Redis """ import os import sys import getpass import asyncio import yaml import pymysql import redis from loguru import logger # 默认配置 DEFAULT_CONFIG = { "database": { "host": "localhost", "port": 3306, "username": "auto_cursor_reg", "password": "auto_cursor_pass", "database": "auto_cursor_reg", "pool_size": 10, "use_redis": True }, "redis": { "host": "127.0.0.1", "port": 6379, "password": "", "db": 0 } } # 数据库表结构 EMAIL_ACCOUNTS_TABLE = ''' CREATE TABLE IF NOT EXISTS email_accounts ( id INT AUTO_INCREMENT PRIMARY KEY, email VARCHAR(255) UNIQUE NOT NULL, password VARCHAR(255) NOT NULL, client_id VARCHAR(255) NOT NULL, refresh_token TEXT NOT NULL, in_use BOOLEAN DEFAULT 0, cursor_password VARCHAR(255), cursor_cookie TEXT, cursor_token TEXT, sold BOOLEAN DEFAULT 0, status VARCHAR(20) DEFAULT 'pending', created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, INDEX idx_status_inuse_sold (status, in_use, sold) ) ''' def setup_logger(): """设置日志""" logger.remove() logger.add(sys.stderr, level="INFO") logger.add("init_database.log", rotation="1 MB", level="DEBUG") def get_mysql_root_credentials(): """获取MySQL root用户凭据""" print("\n===== MySQL配置 =====") print("请输入MySQL root用户凭据以创建数据库和用户") mysql_host = input("MySQL 主机地址 [localhost]: ") or "localhost" mysql_port = input("MySQL 端口 [3306]: ") or "3306" try: mysql_port = int(mysql_port) except ValueError: mysql_port = 3306 print("端口无效，使用默认端口 3306") mysql_root = input("MySQL root用户名 [root]: ") or "root" mysql_root_password = getpass.getpass("MySQL root密码: ") return { "host": mysql_host, "port": mysql_port, "user": mysql_root, "password": mysql_root_password } def get_app_database_config(): """获取应用数据库配置""" print("\n请设置应用程序的数据库配置:") db_name = input("数据库名称 [auto_cursor_reg]: ") or "auto_cursor_reg" db_user = input("数据库用户名 [auto_cursor_reg]: ") or "auto_cursor_reg" db_pass = getpass.getpass(f"为用户 {db_user} 设置密码 [auto_cursor_pass]: ") if not db_pass: db_pass = "auto_cursor_pass" return { "database": db_name, "username": db_user, "password": db_pass } def get_redis_config(): """获取Redis配置""" use_redis = input("\n是否使用Redis缓存? (y/n) [y]: ").lower() != "n" if not use_redis: return {"use_redis": False} print("\n===== Redis配置 =====") redis_host = input("Redis 主机地址 [127.0.0.1]: ") or "127.0.0.1" redis_port = input("Redis 端口 [6379]: ") or "6379" try: redis_port = int(redis_port) except ValueError: redis_port = 6379 print("端口无效，使用默认端口 6379") redis_password = getpass.getpass("Redis 密码 (如无密码请直接回车): ") redis_db = input("Redis 数据库索引 [0]: ") or "0" try: redis_db = int(redis_db) except ValueError: redis_db = 0 print("数据库索引无效，使用默认值 0") return { "use_redis": True, "redis": { "host": redis_host, "port": redis_port, "password": redis_password, "db": redis_db } } def test_mysql_connection(config): """测试MySQL连接""" try: conn = pymysql.connect(**config) conn.close() return True except Exception as e: logger.error(f"MySQL连接失败: {str(e)}") return False def test_redis_connection(config): """测试Redis连接""" try: r = redis.Redis( host=config["host"], port=config["port"], password=config["password"] if config["password"] else None, db=config["db"] ) r.ping() r.close() return True except Exception as e: logger.error(f"Redis连接失败: {str(e)}") return False def setup_mysql(root_config, app_config): """设置MySQL数据库和用户""" logger.info("开始设置MySQL数据库...") try: # 连接到MySQL conn = pymysql.connect(**root_config) cursor = conn.cursor() db_name = app_config["database"] db_user = app_config["username"] db_pass = app_config["password"] # 创建数据库 logger.info(f"创建数据库: {db_name}") cursor.execute(f"CREATE DATABASE IF NOT EXISTS `{db_name}` CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci") # 创建用户并授权 logger.info(f"创建用户: {db_user}") # 检查用户是否已存在 cursor.execute(f"SELECT User FROM mysql.user WHERE User = '{db_user}'") user_exists = cursor.fetchone() if user_exists: logger.info(f"用户 {db_user} 已存在，更新密码") cursor.execute(f"ALTER USER '{db_user}'@'localhost' IDENTIFIED BY '{db_pass}'") cursor.execute(f"ALTER USER '{db_user}'@'%' IDENTIFIED BY '{db_pass}'") else: # 创建用户 (同时创建本地和远程连接权限) try: cursor.execute(f"CREATE USER '{db_user}'@'localhost' IDENTIFIED BY '{db_pass}'") cursor.execute(f"CREATE USER '{db_user}'@'%' IDENTIFIED BY '{db_pass}'") except pymysql.err.MySQLError as e: logger.warning(f"创建用户时出现警告 (可能用户已存在): {str(e)}") # 授权 cursor.execute(f"GRANT ALL PRIVILEGES ON `{db_name}`.* TO '{db_user}'@'localhost'") cursor.execute(f"GRANT ALL PRIVILEGES ON `{db_name}`.* TO '{db_user}'@'%'") cursor.execute("FLUSH PRIVILEGES") # 切换到新创建的数据库 cursor.execute(f"USE `{db_name}`") # 创建表 logger.info("创建数据表: email_accounts") cursor.execute(EMAIL_ACCOUNTS_TABLE) conn.commit() cursor.close() conn.close() logger.success("MySQL数据库设置成功") return True except Exception as e: logger.error(f"设置MySQL失败: {str(e)}") return False def update_config_file(db_config, redis_config=None): """更新配置文件""" logger.info("更新配置文件...") config_file = "config.yaml" try: # 读取现有配置 with open(config_file, "r", encoding="utf-8") as f: config = yaml.safe_load(f) # 更新数据库配置 config["database"].update({ "host": db_config.get("host", "localhost"), "port": db_config.get("port", 3306), "username": db_config["username"], "password": db_config["password"], "database": db_config["database"], "use_redis": db_config.get("use_redis", False) }) # 如果启用Redis，更新Redis配置 if redis_config and db_config.get("use_redis"): config["redis"] = redis_config # 备份原配置文件 if os.path.exists(config_file): os.rename(config_file, f"{config_file}.bak") logger.info(f"已备份原配置文件为 {config_file}.bak") # 写入新配置 with open(config_file, "w", encoding="utf-8") as f: yaml.dump(config, f, default_flow_style=False, allow_unicode=True) logger.success(f"配置文件已更新: {config_file}") return True except Exception as e: logger.error(f"更新配置文件失败: {str(e)}") return False def main(): """主函数""" setup_logger() logger.info("开始初始化数据库") # 获取MySQL root凭据 root_config = get_mysql_root_credentials() # 测试MySQL连接 if not test_mysql_connection(root_config): logger.error("无法连接到MySQL，请检查凭据和服务状态") return # 获取应用数据库配置 db_config = get_app_database_config() # 获取Redis配置 redis_info = get_redis_config() db_config["use_redis"] = redis_info["use_redis"] # 如果启用了Redis，测试连接 if redis_info["use_redis"]: redis_config = redis_info["redis"] if not test_redis_connection(redis_config): use_anyway = input("Redis连接测试失败，是否继续? (y/n) [n]: ").lower() == "y" if not use_anyway: logger.warning("用户取消了初始化") return else: redis_config = None # 设置MySQL if not setup_mysql(root_config, db_config): return # 合并配置 final_db_config = { "host": root_config["host"], "port": root_config["port"], "username": db_config["username"], "password": db_config["password"], "database": db_config["database"], "use_redis": db_config["use_redis"] } # 更新配置文件 update_config_file(final_db_config, redis_config) logger.success("数据库初始化完成!") print("\n===== 初始化完成 =====") print("您现在可以运行以下命令导入邮箱账号:") print(" python import_emails.py") print("\n然后运行主程序开始注册:") print(" python main.py") if __name__ == "__main__": main()