// OpsX VPN WireGuard probe and application proxy. The upstream Go userspace
// engine uses an in-process network stack and does not alter OS routes.
// The private key stays in the local 0600 device record and is never printed.
package main

import (
	"context"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net"
	"net/netip"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"
	"time"

	"golang.zx2c4.com/wireguard/conn"
	"golang.zx2c4.com/wireguard/device"
	"golang.zx2c4.com/wireguard/tun/netstack"
)

const version = "0.1.0"

type record struct {
	PrivateKey string `json:"private_key"`
	Config     struct {
		ClientAddress   string `json:"client_address"`
		ServerPublicKey string `json:"server_public_key"`
		PresharedKey    string `json:"preshared_key"`
		Endpoint        string `json:"endpoint"`
		AllowedIPs      string `json:"allowed_ips"`
		DNS             string `json:"dns"`
		MTU             uint32 `json:"mtu"`
		Keepalive       uint32 `json:"keepalive"`
	} `json:"config"`
}

type result struct {
	Connected bool   `json:"connected"`
	Detail    string `json:"detail"`
	Target    string `json:"target,omitempty"`
}

func readRecord(path string) (*record, error) {
	info, err := os.Lstat(path)
	if err != nil {
		return nil, fmt.Errorf("无法读取本机设备配置: %w", err)
	}
	if !info.Mode().IsRegular() || info.Mode().Perm()&0077 != 0 {
		return nil, errors.New("本机设备配置权限必须为 0600")
	}
	data, err := os.ReadFile(path)
	if err != nil {
		return nil, fmt.Errorf("无法读取本机设备配置: %w", err)
	}
	var r record
	if err := json.Unmarshal(data, &r); err != nil {
		return nil, errors.New("本机设备配置格式无效")
	}
	return &r, nil
}

func key(value string) (string, error) {
	b, err := base64.StdEncoding.DecodeString(value)
	if err != nil || len(b) != 32 {
		return "", errors.New("WireGuard 密钥格式无效")
	}
	return hex.EncodeToString(b), nil
}

func permittedRoute(value string, target netip.Addr) (bool, error) {
	for _, raw := range strings.Split(value, ",") {
		prefix, err := netip.ParsePrefix(strings.TrimSpace(raw))
		if err != nil {
			return false, errors.New("WireGuard 路由无效")
		}
		if prefix.Contains(target) {
			return true, nil
		}
	}
	return false, nil
}

func uapi(r *record, target netip.Addr) (string, []netip.Addr, error) {
	priv, err := key(r.PrivateKey)
	if err != nil {
		return "", nil, err
	}
	pub, err := key(r.Config.ServerPublicKey)
	if err != nil {
		return "", nil, err
	}
	psk, err := key(r.Config.PresharedKey)
	if err != nil {
		return "", nil, err
	}
	addr, err := netip.ParsePrefix(r.Config.ClientAddress)
	if err != nil {
		return "", nil, errors.New("客户端地址无效")
	}
	host, port, err := net.SplitHostPort(r.Config.Endpoint)
	if err != nil || host == "" || port == "" {
		return "", nil, errors.New("网关端点无效")
	}
	portNum, err := strconv.Atoi(port)
	if err != nil || portNum < 1 || portNum > 65535 {
		return "", nil, errors.New("网关端口无效")
	}
	if strings.ContainsAny(host, "\n\r= ") {
		return "", nil, errors.New("网关地址无效")
	}
	var allowed []string
	var permitted bool
	for _, raw := range strings.Split(r.Config.AllowedIPs, ",") {
		prefix, err := netip.ParsePrefix(strings.TrimSpace(raw))
		if err != nil {
			return "", nil, errors.New("WireGuard 路由无效")
		}
		allowed = append(allowed, "allowed_ip="+prefix.String())
		permitted = permitted || !target.IsValid() || prefix.Contains(target)
	}
	if !permitted {
		return "", nil, errors.New("探测目标不在授权路由内")
	}
	if r.Config.Keepalive > 65535 {
		return "", nil, errors.New("保活参数无效")
	}
	lines := []string{
		"private_key=" + priv, "replace_peers=true", "public_key=" + pub,
		"preshared_key=" + psk, "endpoint=" + r.Config.Endpoint,
		"persistent_keepalive_interval=" + strconv.FormatUint(uint64(r.Config.Keepalive), 10),
	}
	lines = append(lines, allowed...)
	return strings.Join(lines, "\n") + "\n", []netip.Addr{addr.Addr()}, nil
}

func startDevice(r *record, target netip.Addr) (*device.Device, *netstack.Net, error) {
	settings, ips, err := uapi(r, target)
	if err != nil {
		return nil, nil, err
	}
	mtu := int(r.Config.MTU)
	if mtu == 0 {
		mtu = 1420
	}
	if mtu < 576 || mtu > 9000 {
		return nil, nil, errors.New("MTU 无效")
	}
	var dns []netip.Addr
	for _, raw := range strings.Split(r.Config.DNS, ",") {
		if strings.TrimSpace(raw) == "" {
			continue
		}
		ip, err := netip.ParseAddr(strings.TrimSpace(raw))
		if err != nil {
			return nil, nil, errors.New("DNS 地址无效")
		}
		dns = append(dns, ip)
	}
	tun, stack, err := netstack.CreateNetTUN(ips, dns, mtu)
	if err != nil {
		return nil, nil, fmt.Errorf("创建 Go 网络栈失败: %w", err)
	}
	dev := device.NewDevice(tun, conn.NewDefaultBind(), device.NewLogger(device.LogLevelSilent, ""))
	if err := dev.IpcSet(settings); err != nil {
		dev.Close()
		return nil, nil, errors.New("WireGuard 连接参数无效")
	}
	if err := dev.Up(); err != nil {
		dev.Close()
		return nil, nil, fmt.Errorf("启动 Go 隧道失败: %w", err)
	}
	return dev, stack, nil
}

func probe(ctx context.Context, r *record, target string) error {
	addr, err := netip.ParseAddrPort(target)
	if err != nil {
		return errors.New("探测地址必须为 IP:端口")
	}
	dev, stack, err := startDevice(r, addr.Addr())
	if err != nil {
		return err
	}
	defer dev.Close()
	channel := make(chan error, 1)
	go func() {
		c, err := stack.DialContext(ctx, "tcp", target)
		if err == nil {
			_ = c.Close()
		}
		channel <- err
	}()
	select {
	case err := <-channel:
		if err != nil {
			return fmt.Errorf("VPN 出口 TCP 连接失败: %w", err)
		}
		return nil
	case <-ctx.Done():
		return errors.New("VPN 出口连接超时，请检查握手、网关 NAT 和防火墙")
	}
}

// A loopback-only SOCKS5 listener provides real, user-space TCP access without
// claiming to alter the operating system's routing table. DNS is resolved in
// the VPN netstack, and every destination is checked against AllowedIPs.
func serveSOCKS(r *record, address string) error {
	listenAddr, err := netip.ParseAddrPort(address)
	if err != nil || !listenAddr.Addr().IsLoopback() || listenAddr.Port() == 0 {
		return errors.New("代理只允许监听明确指定的本机回环端口")
	}
	dev, stack, err := startDevice(r, netip.Addr{})
	if err != nil {
		return err
	}
	defer dev.Close()
	listener, err := net.Listen("tcp", address)
	if err != nil {
		return fmt.Errorf("本机代理监听失败: %w", err)
	}
	defer listener.Close()
	parentPID := os.Getppid()
	go func() {
		ticker := time.NewTicker(2 * time.Second)
		defer ticker.Stop()
		for range ticker.C {
			if os.Getppid() != parentPID {
				_ = listener.Close()
				return
			}
		}
	}()
	// The launcher waits for this line before reporting the proxy ready.
	_ = json.NewEncoder(os.Stdout).Encode(map[string]string{"status": "ready", "listen": address})
	stop := make(chan os.Signal, 1)
	signal.Notify(stop, os.Interrupt, syscall.SIGTERM)
	defer signal.Stop(stop)
	go func() { <-stop; _ = listener.Close() }()
	clients := make(chan struct{}, 128)
	for {
		client, err := listener.Accept()
		if err != nil {
			return nil
		}
		select {
		case clients <- struct{}{}:
			go func() {
				defer func() { <-clients }()
				handleSOCKS(client, r, stack)
			}()
		default:
			_ = client.Close()
		}
	}
}

func handleSOCKS(client net.Conn, r *record, stack *netstack.Net) {
	defer client.Close()
	_ = client.SetDeadline(time.Now().Add(15 * time.Second))
	var greeting [2]byte
	if _, err := io.ReadFull(client, greeting[:]); err != nil || greeting[0] != 5 {
		return
	}
	methods := make([]byte, int(greeting[1]))
	if _, err := io.ReadFull(client, methods); err != nil {
		return
	}
	noAuth := false
	for _, method := range methods {
		if method == 0 {
			noAuth = true
		}
	}
	if !noAuth {
		_, _ = client.Write([]byte{5, 255})
		return
	}
	if _, err := client.Write([]byte{5, 0}); err != nil {
		return
	}
	var header [4]byte
	if _, err := io.ReadFull(client, header[:]); err != nil || header[0] != 5 {
		return
	}
	if header[1] != 1 || header[2] != 0 {
		socksReply(client, 7)
		return
	}
	var host string
	switch header[3] {
	case 1:
		var ip [4]byte
		if _, err := io.ReadFull(client, ip[:]); err != nil {
			return
		}
		host = netip.AddrFrom4(ip).String()
	case 4:
		var ip [16]byte
		if _, err := io.ReadFull(client, ip[:]); err != nil {
			return
		}
		host = netip.AddrFrom16(ip).String()
	case 3:
		var size [1]byte
		if _, err := io.ReadFull(client, size[:]); err != nil || size[0] == 0 {
			return
		}
		name := make([]byte, size[0])
		if _, err := io.ReadFull(client, name); err != nil {
			return
		}
		host = string(name)
	default:
		socksReply(client, 8)
		return
	}
	var port [2]byte
	if _, err := io.ReadFull(client, port[:]); err != nil {
		return
	}
	portNum := int(port[0])<<8 | int(port[1])
	if portNum == 0 {
		socksReply(client, 1)
		return
	}
	addresses := []string{host}
	if _, err := netip.ParseAddr(host); err != nil {
		addresses, err = stack.LookupHost(host)
		if err != nil {
			socksReply(client, 4)
			return
		}
	}
	for _, address := range addresses {
		ip, err := netip.ParseAddr(address)
		if err != nil {
			continue
		}
		allowed, err := permittedRoute(r.Config.AllowedIPs, ip)
		if err != nil || !allowed {
			continue
		}
		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
		upstream, err := stack.DialContext(ctx, "tcp", net.JoinHostPort(address, strconv.Itoa(portNum)))
		cancel()
		if err != nil {
			continue
		}
		defer upstream.Close()
		if err := socksReply(client, 0); err != nil {
			return
		}
		_ = client.SetDeadline(time.Time{})
		go func() { _, _ = io.Copy(upstream, client); _ = upstream.Close() }()
		_, _ = io.Copy(client, upstream)
		return
	}
	socksReply(client, 4)
}

func socksReply(conn net.Conn, code byte) error {
	_, err := conn.Write([]byte{5, code, 0, 1, 127, 0, 0, 1, 0, 0})
	return err
}

func main() {
	if len(os.Args) == 2 && os.Args[1] == "--version" {
		fmt.Println("opsx-tunnel-go " + version)
		return
	}
	if len(os.Args) != 4 || (os.Args[1] != "probe" && os.Args[1] != "proxy") {
		fmt.Fprintln(os.Stderr, "usage: opsx-tunnel-go <probe|proxy> <local-record> <IP:port|loopback:port>")
		os.Exit(2)
	}
	// stdout only contains a non-secret, machine-readable result. Never echo the
	// configuration, keys, token, or arbitrary upstream error strings.
	r, err := readRecord(os.Args[2])
	if err == nil && os.Args[1] == "proxy" {
		err = serveSOCKS(r, os.Args[3])
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
		return
	}
	if err == nil {
		ctx, cancel := context.WithTimeout(context.Background(), 12*time.Second)
		defer cancel()
		err = probe(ctx, r, os.Args[3])
	}
	res := result{Connected: err == nil, Target: os.Args[3]}
	if err == nil {
		res.Detail = "WireGuard 出口 TCP 已连通"
	} else {
		res.Detail = err.Error()
	}
	_ = json.NewEncoder(os.Stdout).Encode(res)
	if err != nil {
		os.Exit(1)
	}
}