wormhole-server/pkg/system/proxy.go

package system

import (
	"fmt"
	"os"
	"os/exec"
	"runtime"
	"strings"

	"github.com/sirupsen/logrus"
)

type SystemProxy struct {
	logger       *logrus.Logger
	proxyAddr    string
	backupConfig map[string]string
}

type Config struct {
	HTTPProxy  string
	HTTPSProxy string
	SOCKSProxy string
	NoProxy    []string
}

func NewSystemProxy(proxyAddr string, logger *logrus.Logger) *SystemProxy {
	return &SystemProxy{
		logger:       logger,
		proxyAddr:    proxyAddr,
		backupConfig: make(map[string]string),
	}
}

// SetSystemProxy 设置系统代理
func (sp *SystemProxy) SetSystemProxy(config Config) error {
	sp.logger.Info("Setting system proxy configuration")

	// 备份当前配置
	if err := sp.backupCurrentConfig(); err != nil {
		sp.logger.WithError(err).Warn("Failed to backup current proxy config")
	}

	switch runtime.GOOS {
	case "darwin":
		return sp.setMacOSProxy(config)
	case "linux":
		return sp.setLinuxProxy(config)
	case "windows":
		return sp.setWindowsProxy(config)
	default:
		return fmt.Errorf("unsupported operating system: %s", runtime.GOOS)
	}
}

// RestoreSystemProxy 恢复系统代理设置
func (sp *SystemProxy) RestoreSystemProxy() error {
	sp.logger.Info("Restoring system proxy configuration")

	switch runtime.GOOS {
	case "darwin":
		return sp.restoreMacOSProxy()
	case "linux":
		return sp.restoreLinuxProxy()
	case "windows":
		return sp.restoreWindowsProxy()
	default:
		return fmt.Errorf("unsupported operating system: %s", runtime.GOOS)
	}
}

// setMacOSProxy 设置macOS系统代理
func (sp *SystemProxy) setMacOSProxy(config Config) error {
	// 获取所有网络服务
	services, err := sp.getMacOSNetworkServices()
	if err != nil {
		return fmt.Errorf("failed to get network services: %w", err)
	}

	for _, service := range services {
		// 设置HTTP代理
		if config.HTTPProxy != "" {
			parts := strings.Split(config.HTTPProxy, ":")
			if len(parts) == 2 {
				if err := sp.runCommand("networksetup", "-setwebproxy", service, parts[0], parts[1]); err != nil {
					sp.logger.WithError(err).WithField("service", service).Warn("Failed to set HTTP proxy")
				}
				if err := sp.runCommand("networksetup", "-setwebproxystate", service, "on"); err != nil {
					sp.logger.WithError(err).WithField("service", service).Warn("Failed to enable HTTP proxy")
				}
			}
		}

		// 设置HTTPS代理
		if config.HTTPSProxy != "" {
			parts := strings.Split(config.HTTPSProxy, ":")
			if len(parts) == 2 {
				if err := sp.runCommand("networksetup", "-setsecurewebproxy", service, parts[0], parts[1]); err != nil {
					sp.logger.WithError(err).WithField("service", service).Warn("Failed to set HTTPS proxy")
				}
				if err := sp.runCommand("networksetup", "-setsecurewebproxystate", service, "on"); err != nil {
					sp.logger.WithError(err).WithField("service", service).Warn("Failed to enable HTTPS proxy")
				}
			}
		}

		// 设置SOCKS代理
		if config.SOCKSProxy != "" {
			parts := strings.Split(config.SOCKSProxy, ":")
			if len(parts) == 2 {
				if err := sp.runCommand("networksetup", "-setsocksfirewallproxy", service, parts[0], parts[1]); err != nil {
					sp.logger.WithError(err).WithField("service", service).Warn("Failed to set SOCKS proxy")
				}
				if err := sp.runCommand("networksetup", "-setsocksfirewallproxystate", service, "on"); err != nil {
					sp.logger.WithError(err).WithField("service", service).Warn("Failed to enable SOCKS proxy")
				}
			}
		}

		// 设置代理绕过列表
		if len(config.NoProxy) > 0 {
			bypassList := strings.Join(config.NoProxy, " ")
			if err := sp.runCommand("networksetup", "-setproxybypassdomains", service, bypassList); err != nil {
				sp.logger.WithError(err).WithField("service", service).Warn("Failed to set proxy bypass list")
			}
		}
	}

	return nil
}

// setLinuxProxy 设置Linux系统代理（通过环境变量）
func (sp *SystemProxy) setLinuxProxy(config Config) error {
	envVars := make(map[string]string)

	if config.HTTPProxy != "" {
		envVars["http_proxy"] = "http://" + config.HTTPProxy
		envVars["HTTP_PROXY"] = "http://" + config.HTTPProxy
	}

	if config.HTTPSProxy != "" {
		envVars["https_proxy"] = "https://" + config.HTTPSProxy
		envVars["HTTPS_PROXY"] = "https://" + config.HTTPSProxy
	}

	if len(config.NoProxy) > 0 {
		noProxy := strings.Join(config.NoProxy, ",")
		envVars["no_proxy"] = noProxy
		envVars["NO_PROXY"] = noProxy
	}

	// 写入到用户的shell配置文件
	return sp.writeLinuxProxyConfig(envVars)
}

// setWindowsProxy 设置Windows系统代理
func (sp *SystemProxy) setWindowsProxy(config Config) error {
	// Windows代理设置通过注册表
	if config.HTTPProxy != "" {
		// 启用代理
		if err := sp.runCommand("reg", "add", `HKCU\Software\Microsoft\Windows\CurrentVersion\Internet Settings`, "/v", "ProxyEnable", "/t", "REG_DWORD", "/d", "1", "/f"); err != nil {
			return fmt.Errorf("failed to enable proxy: %w", err)
		}

		// 设置代理服务器
		if err := sp.runCommand("reg", "add", `HKCU\Software\Microsoft\Windows\CurrentVersion\Internet Settings`, "/v", "ProxyServer", "/t", "REG_SZ", "/d", config.HTTPProxy, "/f"); err != nil {
			return fmt.Errorf("failed to set proxy server: %w", err)
		}
	}

	// 设置代理绕过列表
	if len(config.NoProxy) > 0 {
		bypassList := strings.Join(config.NoProxy, ";")
		if err := sp.runCommand("reg", "add", `HKCU\Software\Microsoft\Windows\CurrentVersion\Internet Settings`, "/v", "ProxyOverride", "/t", "REG_SZ", "/d", bypassList, "/f"); err != nil {
			return fmt.Errorf("failed to set proxy bypass list: %w", err)
		}
	}

	return nil
}

func (sp *SystemProxy) getMacOSNetworkServices() ([]string, error) {
	output, err := exec.Command("networksetup", "-listallnetworkservices").Output()
	if err != nil {
		return nil, err
	}

	lines := strings.Split(string(output), "\n")
	var services []string

	for _, line := range lines {
		line = strings.TrimSpace(line)
		if line != "" && !strings.HasPrefix(line, "*") && line != "An asterisk (*) denotes that a network service is disabled." {
			services = append(services, line)
		}
	}

	return services, nil
}

func (sp *SystemProxy) backupCurrentConfig() error {
	switch runtime.GOOS {
	case "darwin":
		return sp.backupMacOSConfig()
	case "linux":
		return sp.backupLinuxConfig()
	case "windows":
		return sp.backupWindowsConfig()
	}
	return nil
}

func (sp *SystemProxy) backupMacOSConfig() error {
	services, err := sp.getMacOSNetworkServices()
	if err != nil {
		return err
	}

	for _, service := range services {
		// 备份HTTP代理设置
		output, _ := exec.Command("networksetup", "-getwebproxy", service).Output()
		sp.backupConfig[fmt.Sprintf("http_%s", service)] = string(output)

		// 备份HTTPS代理设置
		output, _ = exec.Command("networksetup", "-getsecurewebproxy", service).Output()
		sp.backupConfig[fmt.Sprintf("https_%s", service)] = string(output)

		// 备份SOCKS代理设置
		output, _ = exec.Command("networksetup", "-getsocksfirewallproxy", service).Output()
		sp.backupConfig[fmt.Sprintf("socks_%s", service)] = string(output)
	}

	return nil
}

func (sp *SystemProxy) backupLinuxConfig() error {
	// 备份环境变量
	sp.backupConfig["http_proxy"] = os.Getenv("http_proxy")
	sp.backupConfig["https_proxy"] = os.Getenv("https_proxy")
	sp.backupConfig["no_proxy"] = os.Getenv("no_proxy")
	return nil
}

func (sp *SystemProxy) backupWindowsConfig() error {
	// 备份Windows注册表设置
	output, _ := exec.Command("reg", "query", `HKCU\Software\Microsoft\Windows\CurrentVersion\Internet Settings`, "/v", "ProxyEnable").Output()
	sp.backupConfig["ProxyEnable"] = string(output)

	output, _ = exec.Command("reg", "query", `HKCU\Software\Microsoft\Windows\CurrentVersion\Internet Settings`, "/v", "ProxyServer").Output()
	sp.backupConfig["ProxyServer"] = string(output)

	return nil
}

func (sp *SystemProxy) restoreMacOSProxy() error {
	services, err := sp.getMacOSNetworkServices()
	if err != nil {
		return err
	}

	for _, service := range services {
		// 禁用所有代理
		sp.runCommand("networksetup", "-setwebproxystate", service, "off")
		sp.runCommand("networksetup", "-setsecurewebproxystate", service, "off")
		sp.runCommand("networksetup", "-setsocksfirewallproxystate", service, "off")
	}

	return nil
}

func (sp *SystemProxy) restoreLinuxProxy() error {
	// 清除环境变量（这里简化处理）
	os.Unsetenv("http_proxy")
	os.Unsetenv("https_proxy")
	os.Unsetenv("no_proxy")
	os.Unsetenv("HTTP_PROXY")
	os.Unsetenv("HTTPS_PROXY")
	os.Unsetenv("NO_PROXY")
	return nil
}

func (sp *SystemProxy) restoreWindowsProxy() error {
	// 禁用代理
	return sp.runCommand("reg", "add", `HKCU\Software\Microsoft\Windows\CurrentVersion\Internet Settings`, "/v", "ProxyEnable", "/t", "REG_DWORD", "/d", "0", "/f")
}

func (sp *SystemProxy) writeLinuxProxyConfig(envVars map[string]string) error {
	// 写入到 ~/.bashrc 和 ~/.profile
	homeDir, err := os.UserHomeDir()
	if err != nil {
		return err
	}

	configFiles := []string{
		homeDir + "/.bashrc",
		homeDir + "/.profile",
	}

	proxyLines := []string{"\n# Wormhole SOCKS5 Proxy Configuration"}
	for key, value := range envVars {
		proxyLines = append(proxyLines, fmt.Sprintf("export %s=%s", key, value))
	}
	proxyLines = append(proxyLines, "# End Wormhole Configuration\n")

	for _, configFile := range configFiles {
		file, err := os.OpenFile(configFile, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
		if err != nil {
			sp.logger.WithError(err).WithField("file", configFile).Warn("Failed to open config file")
			continue
		}

		for _, line := range proxyLines {
			if _, err := file.WriteString(line + "\n"); err != nil {
				sp.logger.WithError(err).WithField("file", configFile).Warn("Failed to write to config file")
			}
		}

		file.Close()
	}

	return nil
}

func (sp *SystemProxy) runCommand(name string, args ...string) error {
	cmd := exec.Command(name, args...)
	output, err := cmd.CombinedOutput()
	if err != nil {
		sp.logger.WithFields(logrus.Fields{
			"command": fmt.Sprintf("%s %s", name, strings.Join(args, " ")),
			"output":  string(output),
		}).WithError(err).Debug("Command execution failed")
		return err
	}
	return nil
}

// GetCurrentConfig 获取当前系统代理配置
func (sp *SystemProxy) GetCurrentConfig() (Config, error) {
	config := Config{
		NoProxy: []string{},
	}

	switch runtime.GOOS {
	case "linux":
		config.HTTPProxy = os.Getenv("http_proxy")
		config.HTTPSProxy = os.Getenv("https_proxy")
		noProxy := os.Getenv("no_proxy")
		if noProxy != "" {
			config.NoProxy = strings.Split(noProxy, ",")
		}
	}

	return config, nil
}

// IsProxySet 检查是否已设置代理
func (sp *SystemProxy) IsProxySet() bool {
	switch runtime.GOOS {
	case "linux":
		return os.Getenv("http_proxy") != "" || os.Getenv("https_proxy") != ""
	case "darwin":
		// 检查macOS代理设置（简化）
		return false
	case "windows":
		// 检查Windows代理设置（简化）
		return false
	}
	return false
}