auto_cursor/setup_environment.py

#!/usr/bin/env python3
"""
Cursor自动化服务 - 交互式环境配置脚本
用于配置和初始化新服务器环境、数据库和Redis
"""
import os
import sys
import subprocess
import getpass
import time
import yaml
import socket
import random
import string
from typing import Dict, Any, Optional, Tuple

# 设置颜色输出
class Colors:
    HEADER = '\033[95m'
    BLUE = '\033[94m'
    GREEN = '\033[92m'
    YELLOW = '\033[93m'
    RED = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'

def print_color(text, color):
    """输出彩色文本"""
    print(f"{color}{text}{Colors.ENDC}")

def print_title(title):
    """打印标题"""
    width = 60
    print("\n" + "=" * width)
    print(f"{Colors.BOLD}{Colors.HEADER}{title.center(width)}{Colors.ENDC}")
    print("=" * width + "\n")

def print_step(step, description):
    """打印步骤信息"""
    print(f"{Colors.BOLD}{Colors.BLUE}[步骤 {step}]{Colors.ENDC} {description}")

def print_success(message):
    """打印成功信息"""
    print(f"{Colors.GREEN}✓ {message}{Colors.ENDC}")

def print_warning(message):
    """打印警告信息"""
    print(f"{Colors.YELLOW}⚠ {message}{Colors.ENDC}")

def print_error(message):
    """打印错误信息"""
    print(f"{Colors.RED}✗ {message}{Colors.ENDC}")

# 检测虚拟环境
def is_in_virtualenv():
    """检查是否在虚拟环境中运行"""
    return hasattr(sys, 'real_prefix') or (hasattr(sys, 'base_prefix') and sys.base_prefix != sys.prefix)

def setup_virtualenv():
    """设置虚拟环境"""
    print_title("虚拟环境检测")
    
    if is_in_virtualenv():
        print_success("已在虚拟环境中运行")
        return True
    
    print_warning("当前不在虚拟环境中运行")
    create_venv = input("是否创建并使用虚拟环境? (推荐) (y/n, 默认: y): ").strip().lower()
    
    if create_venv == 'n':
        print_warning("跳过虚拟环境创建，将直接在系统Python环境中安装依赖")
        return True
    
    venv_path = input("请输入虚拟环境路径 (默认: ./venv): ").strip() or "./venv"
    
    try:
        # 检查venv模块
        try:
            import venv
            has_venv = True
        except ImportError:
            has_venv = False
            
        if not has_venv:
            print_warning("Python venv模块不可用，尝试安装...")
            if sys.platform.startswith('linux'):
                print("在Linux上安装venv模块...")
                try:
                    subprocess.check_call(["sudo", "apt", "install", "-y", "python3-venv", "python3-full"])
                except:
                    try:
                        subprocess.check_call(["sudo", "yum", "install", "-y", "python3-venv"])
                    except:
                        print_error("无法自动安装venv模块，请手动安装后重试")
                        print_warning("Ubuntu/Debian: sudo apt install python3-venv python3-full")
                        print_warning("CentOS/RHEL: sudo yum install python3-venv")
                        return False
            else:
                print_error("请安装Python venv模块后重试")
                return False
        
        # 创建虚拟环境
        print_step("创建", f"正在创建虚拟环境: {venv_path}")
        if sys.platform.startswith('win'):
            subprocess.check_call([sys.executable, "-m", "venv", venv_path])
        else:
            subprocess.check_call([sys.executable, "-m", "venv", venv_path])
        
        # 计算激活脚本路径
        if sys.platform.startswith('win'):
            activate_script = os.path.join(venv_path, "Scripts", "activate.bat")
            python_path = os.path.join(venv_path, "Scripts", "python.exe")
        else:
            activate_script = os.path.join(venv_path, "bin", "activate")
            python_path = os.path.join(venv_path, "bin", "python")
        
        print_success(f"虚拟环境创建成功: {venv_path}")
        print_warning("请退出当前程序，然后执行以下命令激活虚拟环境并重新运行:")
        
        if sys.platform.startswith('win'):
            print(f"\n{Colors.BOLD}Windows:{Colors.ENDC}")
            print(f"  {activate_script}")
            print(f"  python setup_environment.py")
        else:
            print(f"\n{Colors.BOLD}Linux/Mac:{Colors.ENDC}")
            print(f"  source {activate_script}")
            print(f"  python setup_environment.py")
            
        return False
        
    except Exception as e:
        print_error(f"创建虚拟环境失败: {str(e)}")
        return False

# 尝试导入必要的库，如果不存在则安装
required_packages = [
    "loguru",
    "pymysql",
    "aiomysql",
    "redis",
    "pyyaml",
    "cryptography"  # 添加这一行
]

def check_and_install_packages():
    """检查并安装必要的Python包"""
    print_step(1, "检查并安装必要的Python包")
    
    for package in required_packages:
        try:
            __import__(package.split(">=")[0])
            print_success(f"已安装: {package}")
        except ImportError:
            print_warning(f"未找到 {package}，正在安装...")
            try:
                subprocess.check_call([sys.executable, "-m", "pip", "install", package])
                print_success(f"成功安装 {package}")
            except subprocess.CalledProcessError as e:
                print_error(f"安装 {package} 失败: {str(e)}")
                sys.exit(1)
    
    # 安装后重新导入必要的库
    try:
        global yaml
        import yaml
        from loguru import logger
        print_success("所有必要的包已安装")
    except ImportError as e:
        print_error(f"无法导入必要的库: {str(e)}")
        sys.exit(1)

def generate_password(length=16):
    """生成随机密码"""
    chars = string.ascii_letters + string.digits + "!@#$%^&*()_+-="
    return ''.join(random.choice(chars) for _ in range(length))

def get_default_hostname():
    """获取默认主机名"""
    try:
        return socket.gethostname()
    except:
        return "cursor-server"

def read_config():
    """读取现有配置文件"""
    config_path = "config.yaml"
    if os.path.exists(config_path):
        try:
            with open(config_path, 'r', encoding='utf-8') as f:
                return yaml.safe_load(f)
        except Exception as e:
            print_warning(f"读取配置文件失败: {str(e)}")
    return {}

def configure_server_settings(existing_config):
    """配置服务器设置"""
    print_step(2, "配置服务器标识")
    
    existing_hostname = None
    if 'server_config' in existing_config and 'hostname' in existing_config['server_config']:
        existing_hostname = existing_config['server_config']['hostname']
    
    default_hostname = existing_hostname or get_default_hostname()
    hostname = input(f"请输入服务器唯一标识 (默认: {default_hostname}): ").strip()
    if not hostname:
        hostname = default_hostname
    
    if not 'server_config' in existing_config:
        existing_config['server_config'] = {}
    
    existing_config['server_config']['hostname'] = hostname
    print_success(f"服务器标识已设置为: {hostname}")
    return existing_config

def test_mysql_connection(host, port, user, password, database=None):
    """测试MySQL连接"""
    import pymysql
    try:
        conn_args = {
            'host': host,
            'port': port,
            'user': user,
            'password': password,
            'connect_timeout': 5
        }
        if database:
            conn_args['database'] = database
        
        conn = pymysql.connect(**conn_args)
        conn.close()
        return True, None
    except Exception as e:
        return False, str(e)

def configure_mysql_settings(existing_config):
    """配置MySQL设置"""
    print_step(3, "配置MySQL数据库")
    
    db_config = existing_config.get('database', {})
    
    print("请输入MySQL配置信息:")
    host = input(f"主机地址 (默认: {db_config.get('host', 'localhost')}): ").strip() or db_config.get('host', 'localhost')
    port = input(f"端口 (默认: {db_config.get('port', 3306)}): ").strip() or db_config.get('port', 3306)
    try:
        port = int(port)
    except ValueError:
        print_warning("端口必须是数字，使用默认值3306")
        port = 3306
    
    # 获取root信息以创建数据库和用户
    print("\n需要MySQL root权限来创建数据库和用户:")
    root_user = input("MySQL root用户名 (默认: root): ").strip() or "root"
    root_password = getpass.getpass("MySQL root密码: ")
    
    # 测试root连接
    success, error_msg = test_mysql_connection(host, port, root_user, root_password)
    if not success:
        print_error(f"无法连接到MySQL: {error_msg}")
        print_warning("请确认MySQL服务已启动且用户名密码正确")
        retry = input("是否重试MySQL配置? (y/n): ").lower()
        if retry == 'y':
            return configure_mysql_settings(existing_config)
        else:
            sys.exit(1)
    
    print_success("MySQL连接测试成功")
    
    # 应用数据库和用户设置
    db_name = input(f"数据库名称 (默认: {db_config.get('database', 'auto_cursor_reg')}): ").strip() or db_config.get('database', 'auto_cursor_reg')
    db_user = input(f"应用用户名 (默认: {db_config.get('username', 'auto_cursor_reg')}): ").strip() or db_config.get('username', 'auto_cursor_reg')
    
    # 为应用用户生成密码
    default_password = db_config.get('password')
    if default_password:
        use_existing = input(f"使用现有密码? (y/n, 默认: y): ").lower() != 'n'
        if use_existing:
            db_password = default_password
        else:
            suggested_password = generate_password(12)
            db_password = input(f"应用用户密码 (建议: {suggested_password}): ").strip() or suggested_password
    else:
        suggested_password = generate_password(12)
        db_password = input(f"应用用户密码 (建议: {suggested_password}): ").strip() or suggested_password
    
    # 更新配置
    existing_config['database'] = {
        'host': host,
        'port': port,
        'username': db_user,
        'password': db_password,
        'database': db_name,
        'pool_size': db_config.get('pool_size', 10),
        'use_redis': db_config.get('use_redis', True)
    }
    
    # 保存root信息用于创建数据库和用户
    root_info = {
        'host': host,
        'port': port,
        'user': root_user,
        'password': root_password
    }
    
    return existing_config, root_info

def test_redis_connection(host, port, password=None, db=0):
    """测试Redis连接"""
    try:
        import redis
        client = redis.Redis(
            host=host,
            port=port,
            password=password if password else None,
            db=db,
            socket_timeout=5
        )
        client.ping()
        client.close()
        return True, None
    except Exception as e:
        return False, str(e)

def configure_redis_settings(existing_config):
    """配置Redis设置"""
    print_step(4, "配置Redis缓存")
    
    db_config = existing_config.get('database', {})
    redis_config = existing_config.get('redis', {})
    
    use_redis = input(f"是否使用Redis缓存? (y/n, 默认: {'y' if db_config.get('use_redis', True) else 'n'}): ").lower()
    if use_redis == 'n':
        existing_config['database']['use_redis'] = False
        print_warning("Redis缓存已禁用")
        return existing_config, None
    
    existing_config['database']['use_redis'] = True
    
    print("请输入Redis配置信息:")
    host = input(f"主机地址 (默认: {redis_config.get('host', '127.0.0.1')}): ").strip() or redis_config.get('host', '127.0.0.1')
    port = input(f"端口 (默认: {redis_config.get('port', 6379)}): ").strip() or redis_config.get('port', 6379)
    try:
        port = int(port)
    except ValueError:
        print_warning("端口必须是数字，使用默认值6379")
        port = 6379
    
    password = input(f"密码 (默认: {'保持现有密码' if redis_config.get('password') else '无'}): ")
    if not password and 'password' in redis_config:
        password = redis_config['password']
    
    db = input(f"数据库索引 (默认: {redis_config.get('db', 0)}): ").strip() or redis_config.get('db', 0)
    try:
        db = int(db)
    except ValueError:
        print_warning("数据库索引必须是数字，使用默认值0")
        db = 0
    
    # 测试Redis连接
    success, error_msg = test_redis_connection(host, port, password, db)
    if not success:
        print_error(f"无法连接到Redis: {error_msg}")
        print_warning("请确认Redis服务已启动且配置正确")
        retry = input("是否重试Redis配置? (y/n): ").lower()
        if retry == 'y':
            return configure_redis_settings(existing_config)
        elif retry == 'n':
            use_anyway = input("是否仍然使用这些设置? (y/n): ").lower()
            if use_anyway != 'y':
                existing_config['database']['use_redis'] = False
                print_warning("Redis缓存已禁用")
                return existing_config, None
    else:
        print_success("Redis连接测试成功")
    
    # 更新配置
    existing_config['redis'] = {
        'host': host,
        'port': port,
        'password': password,
        'db': db
    }
    
    redis_info = {
        'host': host,
        'port': port,
        'password': password,
        'db': db
    }
    
    return existing_config, redis_info

def setup_mysql_database(config, root_info):
    """设置MySQL数据库和用户"""
    print_step(5, "创建MySQL数据库和用户")
    
    db_config = config['database']
    db_name = db_config['database']
    db_user = db_config['username']
    db_password = db_config['password']
    
    import pymysql
    try:
        # 连接MySQL
        conn = pymysql.connect(
            host=root_info['host'],
            port=root_info['port'],
            user=root_info['user'],
            password=root_info['password']
        )
        
        with conn.cursor() as cursor:
            # 创建数据库
            print(f"创建数据库 {db_name}...")
            cursor.execute(f"CREATE DATABASE IF NOT EXISTS `{db_name}` CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci")
            
            # 创建用户并授权
            print(f"创建用户 {db_user}...")
            
            # 检查用户是否存在
            cursor.execute(f"SELECT 1 FROM mysql.user WHERE user = %s AND host = %s", (db_user, '%'))
            user_exists = cursor.fetchone()
            
            if user_exists:
                print_warning(f"用户 {db_user} 已存在，更新密码...")
                cursor.execute(f"ALTER USER '{db_user}'@'%' IDENTIFIED BY %s", (db_password,))
            else:
                cursor.execute(f"CREATE USER '{db_user}'@'%' IDENTIFIED BY %s", (db_password,))
            
            # 授权
            cursor.execute(f"GRANT ALL PRIVILEGES ON `{db_name}`.* TO '{db_user}'@'%'")
            cursor.execute("FLUSH PRIVILEGES")
        
        conn.close()
        print_success(f"数据库 {db_name} 和用户 {db_user} 创建成功")
        return True
    except Exception as e:
        print_error(f"设置MySQL数据库失败: {str(e)}")
        return False

def create_database_tables(config):
    """创建数据库表结构"""
    print_step(6, "创建数据库表结构")
    
    db_config = config['database']
    
    import pymysql
    try:
        # 连接MySQL
        conn = pymysql.connect(
            host=db_config['host'],
            port=db_config['port'],
            user=db_config['username'],
            password=db_config['password'],
            database=db_config['database']
        )
        
        with conn.cursor() as cursor:
            # 创建email_accounts表
            print("创建email_accounts表...")
            cursor.execute('''
                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)
                )
            ''')
            
            # 创建system_settings表
            print("创建system_settings表...")
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS system_settings (
                    id INT AUTO_INCREMENT PRIMARY KEY,
                    key_name VARCHAR(50) UNIQUE NOT NULL,
                    value TEXT,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP
                )
            ''')
            
            # 初始化system_settings
            cursor.execute(
                "INSERT INTO system_settings (key_name, value) VALUES ('auto_service_enabled', '0') "
                "ON DUPLICATE KEY UPDATE value = value"
            )
        
        conn.close()
        print_success("数据库表结构创建成功")
        return True
    except Exception as e:
        print_error(f"创建数据库表失败: {str(e)}")
        return False

def update_config_file(config):
    """更新配置文件"""
    print_step(7, "更新配置文件")
    
    config_path = "config.yaml"
    backup_path = f"config.yaml.bak.{int(time.time())}"
    
    # 备份现有配置
    if os.path.exists(config_path):
        try:
            with open(config_path, 'r', encoding='utf-8') as src:
                with open(backup_path, 'w', encoding='utf-8') as dst:
                    dst.write(src.read())
            print_success(f"现有配置已备份为 {backup_path}")
        except Exception as e:
            print_warning(f"备份配置文件失败: {str(e)}")
    
    # 写入新配置
    try:
        with open(config_path, 'w', encoding='utf-8') as f:
            yaml.dump(config, f, default_flow_style=False, sort_keys=False)
        print_success(f"配置已更新至 {config_path}")
        return True
    except Exception as e:
        print_error(f"更新配置文件失败: {str(e)}")
        return False

def finalize_setup():
    """完成设置并提供指导"""
    print_title("设置完成")
    
    print(f"{Colors.GREEN}Cursor自动化服务环境配置已完成!{Colors.ENDC}")
    print("\n接下来您可以:")
    
    print(f"{Colors.BOLD}1. 导入邮箱账号:{Colors.ENDC}")
    print("   python import_emails.py")
    
    print(f"\n{Colors.BOLD}2. 启动服务:{Colors.ENDC}")
    print("   python start.py")
    print("   或")
    print("   python auto_cursor_service.py")
    
    print(f"\n{Colors.BOLD}3. 设置为系统服务:{Colors.ENDC}")
    print("   请参考README文件中的系统服务设置说明")
    
    print(f"\n{Colors.BOLD}如需帮助:{Colors.ENDC}")
    print("   查看各README文件获取详细使用说明")
    
    print("\n" + "=" * 60)

def main():
    """主函数"""
    print_title("Cursor自动化服务环境配置")
    
    # 检查Python版本
    if sys.version_info < (3, 7):
        print_error("需要Python 3.7或更高版本")
        sys.exit(1)
    
    # 检查并安装依赖
    check_and_install_packages()
    
    # 读取现有配置
    existing_config = read_config()
    
    # 配置服务器设置
    config = configure_server_settings(existing_config)
    
    # 配置MySQL
    config, root_info = configure_mysql_settings(config)
    
    # 配置Redis
    config, redis_info = configure_redis_settings(config)
    
    # 设置MySQL数据库和用户
    if not setup_mysql_database(config, root_info):
        retry = input("MySQL设置失败，是否继续? (y/n): ").lower()
        if retry != 'y':
            sys.exit(1)
    
    # 创建数据库表
    if not create_database_tables(config):
        retry = input("创建表失败，是否继续? (y/n): ").lower()
        if retry != 'y':
            sys.exit(1)
    
    # 更新配置文件
    if not update_config_file(config):
        retry = input("更新配置文件失败，是否继续? (y/n): ").lower()
        if retry != 'y':
            sys.exit(1)
    
    # 完成设置
    finalize_setup()

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n\n设置已取消")
        sys.exit(1)