from flask import Flask, request, jsonify, Response
import random, json, os, uuid, threading

app = Flask(__name__)

# === 配置项 ===
PORT_FILE = "allocated_ports.json"
CLIENTS_FILE = "clients.json"
PORT_RANGE = range(20000, 21000)
AUTH_TOKEN = "Dong200206"
FRPS_ADDR = "lxiaodong.eu.org"
FRPS_PORT = 7123

# === 全局变量 & 锁 ===
lock = threading.Lock()

# === 加载端口记录 ===
if os.path.exists(PORT_FILE):
    with open(PORT_FILE, "r") as f:
        allocated_ports = set(json.load(f))
else:
    allocated_ports = set()

# === 加载客户端记录 ===
if os.path.exists(CLIENTS_FILE):
    with open(CLIENTS_FILE, "r") as f:
        client_ports = json.load(f)
else:
    client_ports = {}

def save_ports():
    with open(PORT_FILE, "w") as f:
        json.dump(list(allocated_ports), f)

def save_clients():
    with open(CLIENTS_FILE, "w") as f:
        json.dump(client_ports, f)

# === token 校验装饰器 ===
def require_token(func):
    def wrapper(*args, **kwargs):
        if request.args.get("token") != AUTH_TOKEN:
            return jsonify({"error": "unauthorized"}), 403
        return func(*args, **kwargs)
    wrapper.__name__ = func.__name__
    return wrapper

# === 工具函数：获取可用端口 ===
def get_available_port():
    with lock:
        available = list(set(PORT_RANGE) - allocated_ports)
        if not available:
            return None
        port = random.choice(available)
        allocated_ports.add(port)
        save_ports()
        return port

# === API: 获取端口 ===
@app.route("/get_port", methods=["GET"])
@require_token
def get_port():
    port = get_available_port()
    if not port:
        return jsonify({"error": "no port available"}), 500
    return jsonify({"port": port})

# === API: 释放端口 ===
@app.route("/release_port", methods=["POST"])
@require_token
def release_port():
    data = request.get_json()
    port = data.get("port")
    client_id = data.get("client_id")

    with lock:
        if client_id and client_id in client_ports:
            port = client_ports.pop(client_id)
            save_clients()

        if port in allocated_ports:
            allocated_ports.remove(port)
            save_ports()
            return jsonify({"status": "released", "port": port})
    return jsonify({"error": "port not found"}), 404

# === API: 生成彩色提示+错误检测的一键安装脚本 ===
@app.route("/install_frpc.sh", methods=["GET"])
@require_token
def install_frpc():
    client_id = request.args.get("client_id")

    with lock:
        if not client_id:
            client_id = str(uuid.uuid4())
        if client_id in client_ports:
            port = client_ports[client_id]
        else:
            port = get_available_port()
            if not port:
                return Response("echo '[ERROR] No port available'", mimetype="text/plain")
            client_ports[client_id] = port
            save_clients()

    script_template = f"""#!/bin/bash
set -e

# ====== 颜色定义 ======
RED='\\033[31m'
GREEN='\\033[32m'
YELLOW='\\033[33m'
BLUE='\\033[34m'
NC='\\033[0m'

echo -e "${{BLUE}}[INFO] 分配到端口: {port}${{NC}}"
echo -e "${{BLUE}}[INFO] client_id={client_id}${{NC}}"

if [ "$(id -u)" != "0" ]; then
    echo -e "${{RED}}[ERROR] 此脚本需要以 root 权限运行${{NC}}"
    exit 1
fi

FRP_VERSION="0.64.0"
ARCH=$(uname -m)
case "$ARCH" in
    x86_64) FRP_ARCH="amd64" ;;
    aarch64) FRP_ARCH="arm64" ;;
    *) echo -e "${{RED}}[ERROR] 不支持的架构: $ARCH${{NC}}"; exit 1 ;;
esac

echo -e "${{YELLOW}}[INFO] 下载 frp v$FRP_VERSION...${{NC}}"
wget -qO /tmp/frp.tar.gz http://{FRPS_ADDR}:8002/frp/frp_${{FRP_VERSION}}_linux_${{FRP_ARCH}}.tar.gz \
    || {{ echo -e "${{RED}}[ERROR] 下载 frp 失败${{NC}}"; exit 1; }}

tar -xzf /tmp/frp.tar.gz -C /tmp \
    || {{ echo -e "${{RED}}[ERROR] 解压 frp 失败${{NC}}"; exit 1; }}

mv /tmp/frp_${{FRP_VERSION}}_linux_${{FRP_ARCH}}/frpc /usr/local/bin/ \
    || {{ echo -e "${{RED}}[ERROR] 移动 frpc 失败${{NC}}"; exit 1; }}

chmod +x /usr/local/bin/frpc
rm -rf /tmp/frp*

mkdir -p /etc/frp
cat <<EOF | tee /etc/frp/frpc.toml
serverAddr = "{FRPS_ADDR}"
serverPort = {FRPS_PORT}
auth.method = "token"
auth.token = "{AUTH_TOKEN}"

[[proxies]]
name = "pve_{port}"
type = "tcp"
localIP = "127.0.0.1"
localPort = 8006
remotePort = {port}
EOF

cat <<EOF | tee /etc/systemd/system/frpc.service
[Unit]
Description=FRP Client Service
After=network.target

[Service]
ExecStart=/usr/local/bin/frpc -c /etc/frp/frpc.toml
Restart=always
User=nobody
LimitNOFILE=1048576

[Install]
WantedBy=multi-user.target
EOF

systemctl daemon-reexec \
    || {{ echo -e "${{RED}}[ERROR] systemctl daemon-reexec 失败${{NC}}"; exit 1; }}
systemctl enable frpc \
    || {{ echo -e "${{RED}}[ERROR] systemctl enable frpc 失败${{NC}}"; exit 1; }}
systemctl restart frpc \
    || {{ echo -e "${{RED}}[ERROR] systemctl restart frpc 失败${{NC}}"; exit 1; }}

echo -e "${{GREEN}}[SUCCESS] frpc 已安装并启动${{NC}}"
echo -e "${{GREEN}}访问地址: http://{FRPS_ADDR}:{port}${{NC}}"
"""
    return Response(script_template, mimetype="text/plain")

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=8888)