frpc-console/frp.go

package main

import (
	"bytes"
	"embed"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"sync"
	"text/template"

	"github.com/creack/pty"
)

//go:embed bin/*
var embeddedFrpc embed.FS

//go:embed frpc.tmpl
var FrpcTemplateContent string

var (
	cachedFrpcPath string
	frpcPathMutex  sync.Mutex
)

// getFrpcPath 获取 frpc 路径，优先级：
// 1. 当前目录 ./bin/<platform>   （便于用户手动管理 frpc 版本）
// 2. embed 解压到 /tmp           （便于开箱即用）
// 3. 系统 PATH 中的 frpc          （降级方案）
func getFrpcPath() (string, error) {
	frpcPathMutex.Lock()
	defer frpcPathMutex.Unlock()

	// 如果缓存有效，直接返回
	if cachedFrpcPath != "" {
		if _, err := os.Stat(cachedFrpcPath); err == nil {
			return cachedFrpcPath, nil
		}
		cachedFrpcPath = ""
	}

	// 确定当前平台的文件名
	var fileName string
	switch {
	case runtime.GOOS == "windows" && runtime.GOARCH == "amd64":
		fileName = "frpc_windows_amd64.exe"
	case runtime.GOOS == "windows" && runtime.GOARCH == "386":
		fileName = "frpc_windows_386.exe"
	case runtime.GOOS == "linux" && runtime.GOARCH == "amd64":
		fileName = "frpc_linux_amd64"
	case runtime.GOOS == "linux" && runtime.GOARCH == "arm64":
		fileName = "frpc_linux_arm64"
	case runtime.GOOS == "linux" && runtime.GOARCH == "arm":
		fileName = "frpc_linux_armv7"
	default:
		// 不支持的平台，尝试从 PATH 找
		path, err := exec.LookPath("frpc")
		if err == nil {
			cachedFrpcPath = path
			return path, nil
		}
		return "", fmt.Errorf("不支持的平台: %s/%s", runtime.GOOS, runtime.GOARCH)
	}

	// ===== 优先级 1：从当前目录 ./bin/ 查找 =====
	localPath := filepath.Join(".", "bin", fileName)
	if _, err := os.Stat(localPath); err == nil {
		cachedFrpcPath = localPath
		return localPath, nil
	}

	// ===== 优先级 2：从 embed 读取并解压到临时目录 =====
	data, err := embeddedFrpc.ReadFile("bin/" + fileName)
	if err == nil {
		tmpPath := filepath.Join(os.TempDir(), "frpc")
		if runtime.GOOS == "windows" {
			tmpPath += ".exe"
		}
		if err := os.WriteFile(tmpPath, data, 0755); err == nil {
			cachedFrpcPath = tmpPath
			return tmpPath, nil
		}
		// 写入失败时，检查是否已有可用的临时文件
		if _, statErr := os.Stat(tmpPath); statErr == nil {
			cachedFrpcPath = tmpPath
			return tmpPath, nil
		}
	}

	// ===== 优先级 3：降级到系统 PATH =====
	path, err := exec.LookPath("frpc")
	if err == nil {
		cachedFrpcPath = path
		return path, nil
	}

	return "", fmt.Errorf("未找到 frpc 文件 (本地 bin/、embed、PATH 均无)")
}

// GenerateFrpcConfig 生成 frpc.toml 配置文件
func GenerateFrpcConfig() error {
	cfg, err := GetGlobalConfig()
	if err != nil {
		return fmt.Errorf("读取全局配置失败: %w", err)
	}

	proxies, err := GetProxies()
	if err != nil {
		return fmt.Errorf("读取隧道列表失败: %w", err)
	}

	var activeProxies []Proxy
	for _, p := range proxies {
		if p.Enabled {
			activeProxies = append(activeProxies, p)
		}
	}

	data := struct {
		*GlobalConfig
		Proxies []Proxy
	}{
		GlobalConfig: cfg,
		Proxies:      activeProxies,
	}

	// 优先使用外部模板文件，不存在则用 embed
	var tmplContent string
	if _, err := os.Stat("frpc.tmpl"); err == nil {
		content, readErr := os.ReadFile("frpc.tmpl")
		if readErr == nil {
			tmplContent = string(content)
		} else {
			tmplContent = FrpcTemplateContent
		}
	} else {
		tmplContent = FrpcTemplateContent
	}

	tmpl, err := template.New("frpc").Parse(tmplContent)
	if err != nil {
		return fmt.Errorf("解析模板失败: %w", err)
	}

	var buf bytes.Buffer
	if err := tmpl.Execute(&buf, data); err != nil {
		return fmt.Errorf("渲染模板失败: %w", err)
	}

	if err := os.WriteFile("./frpc.toml", buf.Bytes(), 0644); err != nil {
		return fmt.Errorf("写入配置文件失败: %w", err)
	}

	return nil
}

// StartFrpc 启动 frpc

func StartFrpc() error {
	frpcPath, err := getFrpcPath()
	if err != nil {
		return fmt.Errorf("获取 frpc 路径失败: %w", err)
	}

	if _, err := os.Stat("./frpc.toml"); os.IsNotExist(err) {
		if err := GenerateFrpcConfig(); err != nil {
			return fmt.Errorf("生成配置文件失败: %w", err)
		}
	}

	running, _ := GetFrpcStatus()
	if running {
		return nil
	}

	cmd := exec.Command(frpcPath, "-c", "./frpc.toml")
	setWindowHide(cmd)

	// ✅ Linux 下分配伪终端，让 frpc 认为自己在真实终端中运行
	if runtime.GOOS != "windows" {
		f, err := pty.Start(cmd)
		if err != nil {
			return fmt.Errorf("启动伪终端失败: %w", err)
		}
		// 可选：丢弃 pty 输出，或写入日志文件
		go io.Copy(io.Discard, f)
	} else {
		if err := cmd.Start(); err != nil {
			return fmt.Errorf("启动 frpc 失败: %w", err)
		}
	}

	return nil
}

// ReloadFrpc 热加载，失败时自动降级为重启
func ReloadFrpc() error {
	running, err := GetFrpcStatus()
	if err != nil {
		return fmt.Errorf("检查 frpc 状态失败: %w", err)
	}
	if !running {
		return StartFrpc()
	}

	frpcPath, err := getFrpcPath()
	if err != nil {
		return fmt.Errorf("获取 frpc 路径失败: %w", err)
	}

	cmd := exec.Command(frpcPath, "reload", "-c", "./frpc.toml")
	_, err = cmd.CombinedOutput()
	if err != nil {
		log.Printf("⚠️  热加载失败 (%v)，自动降级为重启 frpc", err)
		if stopErr := StopFrpc(); stopErr != nil {
			return fmt.Errorf("停止 frpc 失败: %w", stopErr)
		}
		if startErr := StartFrpc(); startErr != nil {
			return fmt.Errorf("启动 frpc 失败: %w", startErr)
		}
		return nil
	}
	return nil
}

// StopFrpc 停止 frpc
func StopFrpc() error {
	if runtime.GOOS == "windows" {
		cmd := exec.Command("taskkill", "/F", "/IM", "frpc.exe")
		if err := cmd.Run(); err != nil && !strings.Contains(err.Error(), "not found") {
			return fmt.Errorf("停止 frpc 失败: %w", err)
		}
		return nil
	}
	cmd := exec.Command("pkill", "-f", "frpc")
	if err := cmd.Run(); err != nil && !strings.Contains(err.Error(), "no process") {
		return fmt.Errorf("停止 frpc 失败: %w", err)
	}
	return nil
}

// GetFrpcStatus 检查 frpc 是否在运行
func GetFrpcStatus() (bool, error) {
	if runtime.GOOS == "windows" {
		cmd := exec.Command("tasklist", "/FI", "IMAGENAME eq frpc.exe")
		output, err := cmd.CombinedOutput()
		if err != nil {
			return false, nil
		}
		return strings.Contains(string(output), "frpc.exe"), nil
	}
	cmd := exec.Command("pgrep", "-f", "frpc")
	output, err := cmd.Output()
	if err != nil {
		return false, nil
	}
	return len(output) > 0, nil
}