#!/usr/bin/env python3
"""
Headscale 域名修复脚本 - 全自动版本
通过 SSH 连接服务器，修改 headscale 配置
"""

import paramiko
import sys
import re
import os
from datetime import datetime

# 配置
HOST = "18.145.193.23"
PORT = 22
USER = "hooper"
PASS = "HS7nFGxDbqQ"
OLD_DOMAIN = "kockol"
NEW_DOMAIN = "wtbadm.lqqnw.cn"

def ssh_exec(client, cmd, desc=""):
    print(f"  [执行] {desc or cmd}")
    stdin, stdout, stderr = client.exec_command(cmd)
    out = stdout.read().decode().strip()
    err = stderr.read().decode().strip()
    if out:
        print(f"    {out}")
    if err and "Warning" not in err:
        print(f"    ERR: {err}")
    return out, err

def find_config_paths(client):
    """查找所有 headscale 相关配置文件"""
    paths = []
    commands = [
        "find /etc /opt /var/lib /root /home -maxdepth 5 -name 'config.yaml' -path '*headscale*' 2>/dev/null",
        "find /etc /opt /var/lib /root /home -maxdepth 5 -name 'config.yml' -path '*headscale*' 2>/dev/null",
        "find /etc /opt /var/lib /root /home -maxdepth 5 -name 'docker-compose*.yml' 2>/dev/null | xargs grep -l headscale 2>/dev/null",
        "find /etc /opt /var/lib /root /home -maxdepth 5 -name 'Caddyfile' 2>/dev/null",
        "find /etc /opt /var/lib /root /home -maxdepth 5 -name '*.conf' -path '*nginx*' 2>/dev/null | xargs grep -l headscale 2>/dev/null",
    ]
    for cmd in commands:
        out, _ = ssh_exec(client, cmd, f"查找配置: {cmd[:50]}...")
        if out:
            paths.extend([p.strip() for p in out.split('\n') if p.strip()])
    return list(set(paths))

def backup_file(client, path):
    backup = f"{path}.backup.{datetime.now().strftime('%Y%m%d_%H%M%S')}"
    ssh_exec(client, f"sudo cp '{path}' '{backup}'", f"备份 {path}")
    return backup

def fix_file(client, path):
    """修复单个文件中的域名"""
    # 先读取内容
    out, _ = ssh_exec(client, f"cat '{path}'", f"读取 {path}")
    if not out:
        return False
    
    if OLD_DOMAIN not in out and NEW_DOMAIN not in out:
        print(f"    跳过（不包含 {OLD_DOMAIN}）")
        return False
    
    if NEW_DOMAIN in out and OLD_DOMAIN not in out:
        print(f"    跳过（已经是新域名）")
        return False
    
    # 备份
    backup_file(client, path)
    
    # 使用 sed 替换
    ssh_exec(client, f"sudo sed -i 's/{OLD_DOMAIN}/{NEW_DOMAIN}/g' '{path}'", f"替换域名")
    
    # 验证
    out2, _ = ssh_exec(client, f"grep -n '{NEW_DOMAIN}' '{path}'", "验证替换")
    if out2:
        print(f"    ✓ 修复成功")
        return True
    return False

def generate_headscale_config():
    return f"""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_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_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
"""

def generate_caddyfile():
    return f"""{NEW_DOMAIN} {{
    reverse_proxy headscale:8080
}}
"""

def main():
    print("=" * 60)
    print("Headscale 域名修复脚本")
    print(f"目标域名: {NEW_DOMAIN}")
    print(f"服务器: {HOST}")
    print("=" * 60)
    
    print("\n[1/5] 连接服务器...")
    client = paramiko.SSHClient()
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    client.connect(HOST, PORT, USER, PASS, timeout=30, banner_timeout=30)
    print("  ✓ SSH 连接成功")
    
    print("\n[2/5] 查找现有配置...")
    paths = find_config_paths(client)
    print(f"  找到 {len(paths)} 个相关文件:")
    for p in paths:
        print(f"    - {p}")
    
    print("\n[3/5] 修复现有配置...")
    fixed = []
    for path in paths:
        if fix_file(client, path):
            fixed.append(path)
    
    print(f"\n[4/5] 重启服务...")
    # 尝试多种重启方式
    restart_cmds = [
        ("docker compose -f /opt/headscale/docker-compose.yml restart", "Docker Compose 重启"),
        ("docker-compose -f /opt/headscale/docker-compose.yml restart", "Docker Compose 重启(旧版)"),
        ("sudo systemctl restart headscale", "systemd 重启"),
        ("sudo docker restart headscale 2>/dev/null || true", "Docker 直接重启"),
    ]
    for cmd, desc in restart_cmds:
        ssh_exec(client, cmd, desc)
    
    print("\n[5/5] 验证服务...")
    ssh_exec(client, f"curl -s -o /dev/null -w '%{{http_code}}' http://localhost:8080/api/v1 || echo '本地检查失败'", "本地服务检查")
    
    print("\n" + "=" * 60)
    print("修复完成!")
    print(f"已修改文件: {fixed}")
    print(f"\n客户端需要执行:")
    print(f"  sudo tailscale logout")
    print(f"  sudo tailscale up --login-server=https://{NEW_DOMAIN}")
    print("=" * 60)
    
    client.close()

if __name__ == "__main__":
    main()