#!/usr/bin/env python3
"""
Headscale 配置修复脚本
将域名从 kockol 修改为 wtbadm.lqqnw.cn
"""

import os
import re
import sys
import shutil
from datetime import datetime

# 配置
OLD_DOMAIN = "kockol"
NEW_DOMAIN = "wtbadm.lqqnw.cn"
SERVER_IP = "18.145.193.23"

# 可能的配置文件路径
CONFIG_PATHS = [
    "/etc/headscale/config.yaml",
    "/etc/headscale/config.yml",
    "/opt/headscale/config.yaml",
    "/opt/headscale/config.yml",
    "/var/lib/headscale/config.yaml",
    "/var/lib/headscale/config.yml",
    "./config.yaml",
    "./config.yml",
]

# Docker Compose 文件路径
COMPOSE_PATHS = [
    "/opt/headscale/docker-compose.yml",
    "/opt/headscale/docker-compose.yaml",
    "/etc/headscale/docker-compose.yml",
    "/etc/headscale/docker-compose.yaml",
    "./docker-compose.yml",
    "./docker-compose.yaml",
]

# Nginx/Traefik/Caddy 配置路径
PROXY_PATHS = [
    "/etc/nginx/conf.d/headscale.conf",
    "/etc/nginx/sites-available/headscale",
    "/etc/nginx/sites-enabled/headscale",
    "/opt/traefik/dynamic/headscale.yml",
    "/etc/caddy/Caddyfile",
    "/opt/caddy/Caddyfile",
]

def backup_file(filepath):
    """创建备份"""
    if os.path.exists(filepath):
        backup_path = f"{filepath}.backup.{datetime.now().strftime('%Y%m%d_%H%M%S')}"
        shutil.copy2(filepath, backup_path)
        print(f"[备份] {filepath} -> {backup_path}")
        return True
    return False

def replace_in_file(filepath, old_str, new_str):
    """在文件中替换字符串"""
    if not os.path.exists(filepath):
        return False
    
    with open(filepath, 'r', encoding='utf-8') as f:
        content = f.read()
    
    if old_str not in content:
        return False
    
    new_content = content.replace(old_str, new_str)
    
    with open(filepath, 'w', encoding='utf-8') as f:
        f.write(new_content)
    
    print(f"[修改] {filepath}: 已替换 '{old_str}' -> '{new_str}'")
    return True

def find_and_replace_domain(filepath):
    """查找并替换域名相关配置"""
    if not os.path.exists(filepath):
        return False
    
    with open(filepath, 'r', encoding='utf-8') as f:
        content = f.read()
    
    original = content
    changes = []
    
    # 替换各种可能的域名格式
    replacements = [
        (f"server_url: https://{OLD_DOMAIN}", f"server_url: https://{NEW_DOMAIN}"),
        (f"server_url: http://{OLD_DOMAIN}", f"server_url: https://{NEW_DOMAIN}"),
        (f"base_domain: {OLD_DOMAIN}", f"base_domain: {NEW_DOMAIN}"),
        (f"- {OLD_DOMAIN}", f"- {NEW_DOMAIN}"),
        (f"@{OLD_DOMAIN}", f"@{NEW_DOMAIN}"),
        (OLD_DOMAIN, NEW_DOMAIN),
    ]
    
    for old, new in replacements:
        if old in content:
            content = content.replace(old, new)
            changes.append(f"'{old}' -> '{new}'")
    
    if content != original:
        with open(filepath, 'w', encoding='utf-8') as f:
            f.write(content)
        print(f"[修改] {filepath}:")
        for change in changes:
            print(f"  - {change}")
        return True
    
    return False

def generate_config():
    """生成正确的 headscale 配置模板"""
    config = f"""# Headscale 配置文件
# 域名: {NEW_DOMAIN}
# 生成时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}

server_url: https://{NEW_DOMAIN}
listen_addr: 0.0.0.0:8080
metrics_listen_addr: 0.0.0.0:9090
grpc_listen_addr: 0.0.0.0:50443
grpc_allow_insecure: false

# 私有 IP 范围配置
ip_prefixes:
  - fd7a:115c:a1e0::/48
  - 100.64.0.0/10

derp:
  server:
    enabled: true
    region_id: 999
    region_code: "headscale"
    region_name: "Headscale Embedded DERP"
    stun_listen_addr: "0.0.0.0:3478"
    private_key_path: /var/lib/headscale/private.key
    automatically_add_embedded_derp_region: true
    ipv4:
      enabled: true
    ipv6:
      enabled: false

  urls:
    - https://controlplane.tailscale.com/derpmap/default

  paths: []

  auto_update_enabled: true
  update_frequency: 24h

disable_check_updates: false
ephemeral_node_inactivity_timeout: 30m
node_update_check_interval: 10s

db_type: sqlite3
db_path: /var/lib/headscale/db.sqlite

# ACL 配置
acl_policy_path: ""

dns_config:
  override_local_dns: true
  nameservers:
    - 223.5.5.5
    - 8.8.8.8
  domains: []
  magic_dns: true
  base_domain: {NEW_DOMAIN}

unix_socket: /var/run/headscale/headscale.sock
unix_socket_permission: "0770"

log:
  format: text
  level: info

# 认证配置
oauth2:
  issuer: ""
  client_id: ""
  client_secret: ""
  expiry: 1h

# 前缀配置
prefixes:
  v6: fd7a:115c:a1e0::/48
  v4: 100.64.0.0/10
  allocation: sequential
"""
    return config

def generate_docker_compose():
    """生成 Docker Compose 配置"""
    compose = f"""version: "3.8"

services:
  headscale:
    image: headscale/headscale:latest
    container_name: headscale
    restart: unless-stopped
    ports:
      - "8080:8080"
      - "9090:9090"
    volumes:
      - ./config.yaml:/etc/headscale/config.yaml
      - ./data:/var/lib/headscale
    command: serve
    networks:
      - headscale-net

  # 可选: Caddy 反向代理
  caddy:
    image: caddy:2-alpine
    container_name: caddy
    restart: unless-stopped
    ports:
      - "80:80"
      - "443:443"
    volumes:
      - ./Caddyfile:/etc/caddy/Caddyfile
      - caddy_data:/data
      - caddy_config:/config
    networks:
      - headscale-net
    depends_on:
      - headscale

networks:
  headscale-net:
    driver: bridge

volumes:
  caddy_data:
  caddy_config:
"""
    return compose

def generate_caddyfile():
    """生成 Caddy 配置"""
    caddyfile = f"""# Caddyfile for Headscale
{NEW_DOMAIN} {{
    reverse_proxy headscale:8080
}}
"""
    return caddyfile

def main():
    print("=" * 60)
    print("Headscale 域名修复脚本")
    print(f"目标域名: {NEW_DOMAIN}")
    print(f"服务器IP: {SERVER_IP}")
    print("=" * 60)
    print()
    
    modified_files = []
    
    # 1. 查找并修改现有配置
    print("[1/4] 查找并修改现有配置文件...")
    all_paths = CONFIG_PATHS + COMPOSE_PATHS + PROXY_PATHS
    
    for path in all_paths:
        if os.path.exists(path):
            backup_file(path)
            if find_and_replace_domain(path):
                modified_files.append(path)
    
    if modified_files:
        print(f"\\n成功修改 {len(modified_files)} 个文件:")
        for f in modified_files:
            print(f"  - {f}")
    else:
        print("未找到需要修改的现有配置文件。")
    
    # 2. 生成新配置文件
    print("\\n[2/4] 生成新的配置文件模板...")
    
    # headscale 配置
    config_path = "/tmp/headscale-config.yaml"
    with open(config_path, 'w') as f:
        f.write(generate_config())
    print(f"[生成] Headscale 配置: {config_path}")
    
    # docker-compose
    compose_path = "/tmp/docker-compose-headscale.yml"
    with open(compose_path, 'w') as f:
        f.write(generate_docker_compose())
    print(f"[生成] Docker Compose: {compose_path}")
    
    # Caddyfile
    caddy_path = "/tmp/Caddyfile"
    with open(caddy_path, 'w') as f:
        f.write(generate_caddyfile())
    print(f"[生成] Caddyfile: {caddy_path}")
    
    # 3. 输出操作指南
    print("\\n[3/4] 操作指南:")
    print("""
# 将生成的配置文件上传到服务器:
scp /tmp/headscale-config.yaml hooper@{SERVER_IP}:/tmp/
scp /tmp/Caddyfile hooper@{SERVER_IP}:/tmp/

# SSH 到服务器后执行:
sudo cp /tmp/headscale-config.yaml /etc/headscale/config.yaml
sudo cp /tmp/Caddyfile /opt/headscale/Caddyfile

# 重启 headscale 服务
docker compose -f /opt/headscale/docker-compose.yml restart
# 或
sudo systemctl restart headscale

# 验证配置
curl https://{NEW_DOMAIN}/api/v1
""".format(SERVER_IP=SERVER_IP, NEW_DOMAIN=NEW_DOMAIN))
    
    # 4. 生成客户端重置命令
    print("[4/4] 客户端重置命令 (在每台客户端执行):")
    print(f"""
sudo tailscale logout
sudo tailscale up --login-server=https://{NEW_DOMAIN}
""")
    
    print("\\n" + "=" * 60)
    print("完成!")
    print("=" * 60)

if __name__ == "__main__":
    main()