ProxyAuto/app.py

"""
ProxyAuto - 自动 IP 更换与 Cloudflare DNS 同步工具

Flask 版本
"""
import re
import sys
import random
from datetime import datetime, timedelta
from pathlib import Path
from functools import wraps
from zoneinfo import ZoneInfo

# 确保当前目录在路径中
sys.path.insert(0, str(Path(__file__).parent))

from flask import Flask, render_template, request, jsonify, redirect, url_for, session

from database import (
    Config,
    ProxyMachine,
    User,
    ensure_singleton_config,
    ensure_admin_user,
    get_session,
    mask_secret,
    LOG_DIR,
)
from services.scheduler import (
    get_scheduler_status,
    get_scheduler,
    start_machine_auto,
    stop_machine_auto,
    update_all_schedulers,
    get_machine_scheduler_status,
)
from services.ip_change import run_ip_change_for_machine
from services.aws_region import is_valid_aws_region, normalize_aws_region
from services.traffic_monitor import (
    get_machine_traffic,
    get_current_month_traffic,
    get_current_day_traffic,
    get_all_time_traffic,
    format_bytes,
)
from services.traffic_alert import check_all_traffic_alerts, reset_machine_alert
from services.email_service import send_email

SHANGHAI_TZ = ZoneInfo("Asia/Shanghai")

app = Flask(__name__)
app.secret_key = 'proxyauto-secret-key-change-in-production'

# 正则验证
_INSTANCE_ID_RE = re.compile(r"^i-[0-9a-f]{8,17}$", re.IGNORECASE)
_LIGHTSAIL_INSTANCE_NAME_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.-]{0,127}$")

# 初始化
db_session = get_session()
ensure_admin_user(db_session)
db_session.close()

# 初始化调度器
update_all_schedulers()


def login_required(f):
    """登录验证装饰器"""
    @wraps(f)
    def decorated_function(*args, **kwargs):
        if 'user_id' not in session:
            return redirect(url_for('login'))
        return f(*args, **kwargs)
    return decorated_function


def generate_captcha():
    """生成简单的数学验证码"""
    a = random.randint(1, 20)
    b = random.randint(1, 20)
    op = random.choice(['+', '-', 'x'])
    if op == '+':
        answer = a + b
        question = f"{a} + {b} = ?"
    elif op == '-':
        if a < b:
            a, b = b, a
        answer = a - b
        question = f"{a} - {b} = ?"
    else:
        answer = a * b
        question = f"{a} x {b} = ?"
    return question, answer


@app.route('/login', methods=['GET', 'POST'])
def login():
    """登录页面"""
    if request.method == 'GET':
        question, answer = generate_captcha()
        session['captcha_answer'] = answer
        return render_template('login.html', captcha_question=question)

    # POST 处理登录
    data = request.json or request.form
    username = data.get('username', '').strip()
    password = data.get('password', '')
    captcha = data.get('captcha', '').strip()

    # 验证码检查
    correct_answer = session.get('captcha_answer')
    if correct_answer is None or str(captcha) != str(correct_answer):
        return jsonify({"ok": False, "message": "验证码错误"}), 400

    db = get_session()
    try:
        user = db.query(User).filter_by(username=username).first()
        if not user or not user.check_password(password):
            # 重新生成验证码
            question, answer = generate_captcha()
            session['captcha_answer'] = answer
            return jsonify({"ok": False, "message": "用户名或密码错误", "new_captcha": question}), 401

        session['user_id'] = user.id
        session['username'] = user.username
        session.pop('captcha_answer', None)
        return jsonify({"ok": True, "message": "登录成功"})
    finally:
        db.close()


@app.route('/logout')
def logout():
    """注销"""
    session.clear()
    return redirect(url_for('login'))


@app.route('/api/refresh-captcha', methods=['POST'])
def refresh_captcha():
    """刷新验证码"""
    question, answer = generate_captcha()
    session['captcha_answer'] = answer
    return jsonify({"ok": True, "question": question})


@app.route('/')
@login_required
def dashboard():
    """控制台页面"""
    db = get_session()
    try:
        machines = db.query(ProxyMachine).order_by(ProxyMachine.name.asc()).all()
        # 获取每台机器的调度状态
        machines_status = []
        for m in machines:
            status = get_machine_scheduler_status(m.id)
            machines_status.append({
                "machine": m,
                "scheduler": status
            })
        return render_template('dashboard.html', machines_status=machines_status)
    finally:
        db.close()


@app.route('/api/run-ip-change/<int:machine_id>', methods=['POST'])
@login_required
def api_run_ip_change(machine_id):
    """执行指定机器的 IP 更换"""
    result = run_ip_change_for_machine(machine_id)
    return jsonify(result)


@app.route('/api/toggle-auto/<int:machine_id>', methods=['POST'])
@login_required
def api_toggle_auto(machine_id):
    """切换指定机器的自动任务状态"""
    db = get_session()
    try:
        machine = db.query(ProxyMachine).filter_by(id=machine_id).first()
        if not machine:
            return jsonify({"ok": False, "message": "机器不存在"}), 404

        if machine.auto_enabled:
            machine.auto_enabled = False
            db.commit()
            stop_machine_auto(machine_id)
            return jsonify({"ok": True, "running": False, "message": f"已停止 {machine.name} 的自动任务"})
        else:
            machine.auto_enabled = True
            db.commit()
            start_machine_auto(machine_id, machine.change_interval_seconds)
            return jsonify({"ok": True, "running": True, "message": f"已启动 {machine.name} 的自动任务"})
    finally:
        db.close()


@app.route('/settings')
@login_required
def settings():
    """配置页面"""
    db = get_session()
    try:
        config = ensure_singleton_config(db)
        return render_template('settings.html', config=config, mask_secret=mask_secret)
    finally:
        db.close()


@app.route('/api/settings', methods=['POST'])
@login_required
def api_save_settings():
    """保存全局配置"""
    db = get_session()
    try:
        config = ensure_singleton_config(db)
        data = request.json

        # AWS 凭证
        if data.get('aws_access_key', '').strip():
            config.aws_access_key = data['aws_access_key'].strip()
        if data.get('aws_secret_key', '').strip():
            config.aws_secret_key = data['aws_secret_key'].strip()

        # Cloudflare 设置
        config.cloudflare_auth_type = data.get('cloudflare_auth_type', 'api_token')

        if config.cloudflare_auth_type == 'api_token':
            if data.get('cf_api_token', '').strip():
                config.cf_api_token = data['cf_api_token'].strip()
        else:
            config.cf_email = data.get('cf_email', '').strip() or None
            if data.get('cf_api_key', '').strip():
                config.cf_api_key = data['cf_api_key'].strip()

        config.release_old_eip = data.get('release_old_eip', True)

        # SMTP 配置
        if data.get('smtp_host', '').strip():
            config.smtp_host = data['smtp_host'].strip()
        config.smtp_port = data.get('smtp_port', 587)
        if data.get('smtp_user', '').strip():
            config.smtp_user = data['smtp_user'].strip()
        if data.get('smtp_password', '').strip():
            config.smtp_password = data['smtp_password'].strip()
        config.smtp_use_tls = data.get('smtp_use_tls', True)
        if data.get('alert_email', '').strip():
            config.alert_email = data['alert_email'].strip()

        db.commit()
        return jsonify({"ok": True, "message": "配置已更新"})
    except Exception as e:
        db.rollback()
        return jsonify({"ok": False, "message": str(e)}), 400
    finally:
        db.close()


@app.route('/machines')
@login_required
def machines():
    """节点管理页面"""
    db = get_session()
    try:
        machines_list = db.query(ProxyMachine).order_by(ProxyMachine.name.asc()).all()
        return render_template('machines.html', machines=machines_list)
    finally:
        db.close()


@app.route('/api/machines', methods=['POST'])
@login_required
def api_add_machine():
    """添加节点"""
    db = get_session()
    try:
        data = request.json
        errors = []

        name = data.get('name', '').strip()
        aws_service = data.get('aws_service', 'ec2')
        aws_region = data.get('aws_region', '').strip()
        instance_id = data.get('instance_id', '').strip()
        note = data.get('note', '').strip()
        enabled = data.get('enabled', True)

        # 域名配置
        cf_zone_id = data.get('cf_zone_id', '').strip()
        cf_record_name = data.get('cf_record_name', '').strip()
        cf_proxied = data.get('cf_proxied', False)

        # 时间配置
        change_interval_minutes = int(data.get('change_interval_minutes', 60))
        auto_enabled = data.get('auto_enabled', False)

        if not name:
            errors.append("节点名称不能为空")
        if not is_valid_aws_region(normalize_aws_region(aws_region)):
            errors.append("AWS 区域格式不正确")
        if not instance_id:
            errors.append("实例 ID 不能为空")
        elif aws_service == "ec2" and not _INSTANCE_ID_RE.match(instance_id):
            errors.append("EC2 Instance ID 格式不正确")
        elif aws_service == "lightsail" and not _LIGHTSAIL_INSTANCE_NAME_RE.match(instance_id):
            errors.append("Lightsail 实例名格式不正确")

        if errors:
            return jsonify({"ok": False, "errors": errors}), 400

        machine = ProxyMachine(
            name=name,
            aws_service=aws_service,
            aws_region=normalize_aws_region(aws_region),
            aws_instance_id=instance_id,
            note=note or None,
            enabled=enabled,
            cf_zone_id=cf_zone_id or None,
            cf_record_name=cf_record_name or None,
            cf_proxied=cf_proxied,
            change_interval_seconds=max(60, change_interval_minutes * 60),
            auto_enabled=auto_enabled,
            traffic_alert_enabled=data.get('traffic_alert_enabled', False),
            traffic_alert_limit_gb=data.get('traffic_alert_limit_gb'),
        )
        db.add(machine)
        db.commit()

        # 如果启用了自动任务，启动调度
        if auto_enabled:
            start_machine_auto(machine.id, machine.change_interval_seconds)

        return jsonify({"ok": True, "message": "节点添加成功", "id": machine.id})
    except Exception as e:
        db.rollback()
        return jsonify({"ok": False, "message": str(e)}), 400
    finally:
        db.close()


@app.route('/api/machines/<int:machine_id>', methods=['PUT'])
@login_required
def api_update_machine(machine_id):
    """更新节点"""
    db = get_session()
    try:
        machine = db.query(ProxyMachine).filter_by(id=machine_id).first()
        if not machine:
            return jsonify({"ok": False, "message": "节点不存在"}), 404

        data = request.json

        machine.name = data.get('name', machine.name).strip()
        machine.aws_service = data.get('aws_service', machine.aws_service)
        machine.aws_region = normalize_aws_region(data.get('aws_region', machine.aws_region))
        machine.aws_instance_id = data.get('instance_id', machine.aws_instance_id).strip()
        machine.note = data.get('note', '').strip() or None
        machine.enabled = data.get('enabled', machine.enabled)

        # 域名配置
        machine.cf_zone_id = data.get('cf_zone_id', '').strip() or None
        machine.cf_record_name = data.get('cf_record_name', '').strip() or None
        machine.cf_proxied = data.get('cf_proxied', machine.cf_proxied)

        # 时间配置
        change_interval_minutes = int(data.get('change_interval_minutes', machine.change_interval_seconds // 60))
        machine.change_interval_seconds = max(60, change_interval_minutes * 60)

        old_auto_enabled = machine.auto_enabled
        machine.auto_enabled = data.get('auto_enabled', machine.auto_enabled)

        # 流量预警配置
        machine.traffic_alert_enabled = data.get('traffic_alert_enabled', machine.traffic_alert_enabled)
        machine.traffic_alert_limit_gb = data.get('traffic_alert_limit_gb')

        db.commit()

        # 更新调度器
        if machine.auto_enabled and not old_auto_enabled:
            start_machine_auto(machine.id, machine.change_interval_seconds)
        elif not machine.auto_enabled and old_auto_enabled:
            stop_machine_auto(machine.id)
        elif machine.auto_enabled:
            # 更新间隔
            stop_machine_auto(machine.id)
            start_machine_auto(machine.id, machine.change_interval_seconds)

        return jsonify({"ok": True, "message": "节点已更新"})
    except Exception as e:
        db.rollback()
        return jsonify({"ok": False, "message": str(e)}), 400
    finally:
        db.close()


@app.route('/api/machines/<int:machine_id>', methods=['DELETE'])
@login_required
def api_delete_machine(machine_id):
    """删除节点"""
    db = get_session()
    try:
        machine = db.query(ProxyMachine).filter_by(id=machine_id).first()
        if not machine:
            return jsonify({"ok": False, "message": "节点不存在"}), 404

        # 停止调度
        stop_machine_auto(machine_id)

        db.delete(machine)
        db.commit()

        return jsonify({"ok": True, "message": "节点已删除"})
    except Exception as e:
        db.rollback()
        return jsonify({"ok": False, "message": str(e)}), 400
    finally:
        db.close()


@app.route('/logs')
@login_required
def logs():
    """日志查看页面"""
    return render_template('logs.html')


@app.route('/api/logs')
@login_required
def api_get_logs():
    """获取日志内容"""
    lines_count = request.args.get('lines', 100, type=int)
    lines_count = max(50, min(500, lines_count))

    log_file = LOG_DIR / "proxyauto.log"

    if not log_file.exists():
        return jsonify({"ok": True, "content": "", "empty": True})

    try:
        with open(log_file, "r", encoding="utf-8", errors="replace") as f:
            lines = f.read().splitlines()

        display_lines = lines[-lines_count:]
        return jsonify({"ok": True, "content": "\n".join(display_lines), "empty": len(display_lines) == 0})
    except Exception as e:
        return jsonify({"ok": False, "message": str(e)}), 500


@app.route('/api/status')
@login_required
def api_status():
    """获取系统状态"""
    db = get_session()
    try:
        machines = db.query(ProxyMachine).all()
        machines_data = []
        for m in machines:
            status = get_machine_scheduler_status(m.id)
            machines_data.append({
                "id": m.id,
                "name": m.name,
                "enabled": m.enabled,
                "auto_enabled": m.auto_enabled,
                "current_ip": m.current_ip,
                "last_run_at": m.last_run_at.isoformat() if m.last_run_at else None,
                "last_success": m.last_success,
                "scheduler_running": status["running"],
                "next_run_time": status["next_run_time"],
            })
        return jsonify({"ok": True, "machines": machines_data})
    finally:
        db.close()


# ============ 流量监控相关路由 ============

@app.route('/traffic')
@login_required
def traffic():
    """流量监控页面"""
    db = get_session()
    try:
        machines_list = db.query(ProxyMachine).filter_by(enabled=True).order_by(ProxyMachine.name.asc()).all()
        return render_template('traffic.html', machines=machines_list)
    finally:
        db.close()


@app.route('/api/traffic/<int:machine_id>')
@login_required
def api_get_traffic(machine_id):
    """获取指定机器的流量数据"""
    db = get_session()
    try:
        machine = db.query(ProxyMachine).filter_by(id=machine_id).first()
        if not machine:
            return jsonify({"ok": False, "message": "机器不存在"}), 404

        config = ensure_singleton_config(db)
        if not config.aws_access_key or not config.aws_secret_key:
            return jsonify({"ok": False, "message": "AWS 凭证未配置"}), 400

        # 解析日期参数
        start_str = request.args.get('start')
        end_str = request.args.get('end')

        if start_str and end_str:
            start_time = datetime.strptime(start_str, '%Y-%m-%d').replace(tzinfo=SHANGHAI_TZ)
            end_time = datetime.strptime(end_str, '%Y-%m-%d').replace(hour=23, minute=59, second=59, tzinfo=SHANGHAI_TZ)
        else:
            # 默认查询当月
            now = datetime.now(SHANGHAI_TZ)
            start_time = now.replace(day=1, hour=0, minute=0, second=0, microsecond=0)
            end_time = now

        # 根据时间跨度决定数据点间隔
        days_diff = (end_time - start_time).days
        if days_diff <= 1:
            period = 300  # 5分钟
        elif days_diff <= 7:
            period = 3600  # 1小时
        else:
            period = 86400  # 1天

        traffic_data = get_machine_traffic(
            aws_service=machine.aws_service,
            region=machine.aws_region,
            instance_id=machine.aws_instance_id,
            aws_access_key=config.aws_access_key,
            aws_secret_key=config.aws_secret_key,
            start_time=start_time,
            end_time=end_time,
            period=period,
        )

        return jsonify(traffic_data)
    except ValueError as e:
        return jsonify({"ok": False, "message": f"日期格式错误: {str(e)}"}), 400
    except Exception as e:
        return jsonify({"ok": False, "message": str(e)}), 500
    finally:
        db.close()


@app.route('/api/traffic/<int:machine_id>/current')
@login_required
def api_get_current_traffic(machine_id):
    """获取当月流量摘要"""
    db = get_session()
    try:
        machine = db.query(ProxyMachine).filter_by(id=machine_id).first()
        if not machine:
            return jsonify({"ok": False, "message": "机器不存在"}), 404

        config = ensure_singleton_config(db)
        if not config.aws_access_key or not config.aws_secret_key:
            return jsonify({"ok": False, "message": "AWS 凭证未配置"}), 400

        traffic_data = get_current_month_traffic(
            aws_service=machine.aws_service,
            region=machine.aws_region,
            instance_id=machine.aws_instance_id,
            aws_access_key=config.aws_access_key,
            aws_secret_key=config.aws_secret_key,
        )

        if traffic_data.get("ok"):
            traffic_data["network_in_formatted"] = format_bytes(traffic_data["network_in"])
            traffic_data["network_out_formatted"] = format_bytes(traffic_data["network_out"])
            traffic_data["total_formatted"] = format_bytes(traffic_data["total"])

        return jsonify(traffic_data)
    finally:
        db.close()


@app.route('/api/traffic/summary')
@login_required
def api_get_traffic_summary():
    """获取所有机器的流量汇总数据"""
    db = get_session()
    try:
        machines = db.query(ProxyMachine).filter_by(enabled=True).all()
        if not machines:
            return jsonify({
                "ok": True,
                "month": {"network_in": 0, "network_out": 0, "total": 0},
                "day": {"network_in": 0, "network_out": 0, "total": 0},
                "all_time": {"network_in": 0, "network_out": 0, "total": 0},
            })

        config = ensure_singleton_config(db)
        if not config.aws_access_key or not config.aws_secret_key:
            return jsonify({"ok": False, "message": "AWS 凭证未配置"}), 400

        # 初始化汇总数据
        month_total = {"network_in": 0, "network_out": 0, "total": 0}
        day_total = {"network_in": 0, "network_out": 0, "total": 0}
        all_time_total = {"network_in": 0, "network_out": 0, "total": 0}

        for machine in machines:
            # 获取当月流量
            month_data = get_current_month_traffic(
                aws_service=machine.aws_service,
                region=machine.aws_region,
                instance_id=machine.aws_instance_id,
                aws_access_key=config.aws_access_key,
                aws_secret_key=config.aws_secret_key,
            )
            if month_data.get("ok"):
                month_total["network_in"] += month_data.get("network_in", 0)
                month_total["network_out"] += month_data.get("network_out", 0)
                month_total["total"] += month_data.get("total", 0)

            # 获取当日流量
            day_data = get_current_day_traffic(
                aws_service=machine.aws_service,
                region=machine.aws_region,
                instance_id=machine.aws_instance_id,
                aws_access_key=config.aws_access_key,
                aws_secret_key=config.aws_secret_key,
            )
            if day_data.get("ok"):
                day_total["network_in"] += day_data.get("network_in", 0)
                day_total["network_out"] += day_data.get("network_out", 0)
                day_total["total"] += day_data.get("total", 0)

            # 获取历史总流量
            all_time_data = get_all_time_traffic(
                aws_service=machine.aws_service,
                region=machine.aws_region,
                instance_id=machine.aws_instance_id,
                aws_access_key=config.aws_access_key,
                aws_secret_key=config.aws_secret_key,
                created_at=machine.created_at,
            )
            if all_time_data.get("ok"):
                all_time_total["network_in"] += all_time_data.get("network_in", 0)
                all_time_total["network_out"] += all_time_data.get("network_out", 0)
                all_time_total["total"] += all_time_data.get("total", 0)

        return jsonify({
            "ok": True,
            "month": {
                "network_in": month_total["network_in"],
                "network_out": month_total["network_out"],
                "total": month_total["total"],
                "network_in_formatted": format_bytes(month_total["network_in"]),
                "network_out_formatted": format_bytes(month_total["network_out"]),
                "total_formatted": format_bytes(month_total["total"]),
            },
            "day": {
                "network_in": day_total["network_in"],
                "network_out": day_total["network_out"],
                "total": day_total["total"],
                "network_in_formatted": format_bytes(day_total["network_in"]),
                "network_out_formatted": format_bytes(day_total["network_out"]),
                "total_formatted": format_bytes(day_total["total"]),
            },
            "all_time": {
                "network_in": all_time_total["network_in"],
                "network_out": all_time_total["network_out"],
                "total": all_time_total["total"],
                "network_in_formatted": format_bytes(all_time_total["network_in"]),
                "network_out_formatted": format_bytes(all_time_total["network_out"]),
                "total_formatted": format_bytes(all_time_total["total"]),
            },
        })
    except Exception as e:
        return jsonify({"ok": False, "message": str(e)}), 500
    finally:
        db.close()


@app.route('/api/traffic/check-alerts', methods=['POST'])
@login_required
def api_check_traffic_alerts():
    """手动触发流量预警检查"""
    result = check_all_traffic_alerts()
    return jsonify(result)


@app.route('/api/machines/<int:machine_id>/reset-alert', methods=['POST'])
@login_required
def api_reset_machine_alert(machine_id):
    """重置机器的流量预警状态"""
    result = reset_machine_alert(machine_id)
    return jsonify(result)


@app.route('/api/test-email', methods=['POST'])
@login_required
def api_test_email():
    """发送测试邮件"""
    db = get_session()
    try:
        config = ensure_singleton_config(db)
        if not config.smtp_host or not config.alert_email:
            return jsonify({"ok": False, "message": "请先配置 SMTP 服务器和接收邮箱"}), 400

        result = send_email(
            smtp_host=config.smtp_host,
            smtp_port=config.smtp_port or 587,
            smtp_user=config.smtp_user,
            smtp_password=config.smtp_password,
            use_tls=config.smtp_use_tls,
            to_email=config.alert_email,
            subject="[ProxyAuto] 测试邮件",
            body="这是一封来自 ProxyAuto Pro 的测试邮件。\n\n如果您收到此邮件，说明邮件配置正确。",
        )
        return jsonify(result)
    finally:
        db.close()


# 启动流量预警定时检查任务
def _traffic_alert_job():
    """流量预警检查任务"""
    check_all_traffic_alerts()


def init_traffic_alert_scheduler():
    """初始化流量预警定时任务"""
    scheduler = get_scheduler()
    scheduler.add_job(
        _traffic_alert_job,
        "interval",
        minutes=10,  # 每10分钟检查一次
        id="traffic_alert_check",
        replace_existing=True,
        max_instances=1,
        coalesce=True,
    )


# 初始化流量预警调度
init_traffic_alert_scheduler()


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=8501, debug=True)