329 lines
9.8 KiB
Python
329 lines
9.8 KiB
Python
#!/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',
|
||
extracted BOOLEAN DEFAULT 0,
|
||
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),
|
||
INDEX idx_extracted (extracted, status, 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() |