panel-vpn/internal/nodeinstall/ssh.go

package nodeinstall

import (
	"bytes"
	"fmt"
	"io"
	"net"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"golang.org/x/crypto/ssh"

	"github.com/orohi/vpn-panel/internal/models"
)

type LogFunc func(line string)

type SSHConfig struct {
	Host     string
	Port     int
	User     string
	AuthType string // password | key
	Secret   string
}

func (c SSHConfig) dial() (*ssh.Client, error) {
	var auth []ssh.AuthMethod
	switch c.AuthType {
	case "key":
		signer, err := ssh.ParsePrivateKey([]byte(c.Secret))
		if err != nil {
			return nil, fmt.Errorf("parse ssh key: %w", err)
		}
		auth = append(auth, ssh.PublicKeys(signer))
	default:
		auth = append(auth, ssh.Password(c.Secret))
	}

	cfg := &ssh.ClientConfig{
		User:            c.User,
		Auth:            auth,
		HostKeyCallback: ssh.InsecureIgnoreHostKey(), // bootstrap only; harden later with known_hosts
		Timeout:         20 * time.Second,
	}
	addr := net.JoinHostPort(c.Host, fmt.Sprintf("%d", c.Port))
	return ssh.Dial("tcp", addr, cfg)
}

func run(client *ssh.Client, cmd string, log LogFunc) (string, error) {
	session, err := client.NewSession()
	if err != nil {
		return "", err
	}
	defer session.Close()

	var stdout, stderr bytes.Buffer
	session.Stdout = &stdout
	session.Stderr = &stderr
	if log != nil {
		log("$ " + cmd)
	}
	err = session.Run(cmd)
	out := strings.TrimSpace(stdout.String() + "\n" + stderr.String())
	if log != nil && out != "" {
		log(out)
	}
	if err != nil {
		return out, fmt.Errorf("%w: %s", err, out)
	}
	return out, nil
}

func shellQuote(s string) string {
	return strconv.Quote(s)
}

// uploadViaStdin writes bytes to remotePath using `cat` over SSH stdin.
// Avoids legacy `scp -t`, which fails on many modern OpenSSH servers (status 1).
func uploadViaStdin(client *ssh.Client, remotePath string, data []byte, mode os.FileMode, log LogFunc) error {
	dir := filepath.ToSlash(filepath.Dir(remotePath))
	if _, err := run(client, "mkdir -p "+shellQuote(dir)+" 2>/dev/null || sudo mkdir -p "+shellQuote(dir), log); err != nil {
		return fmt.Errorf("mkdir %s: %w", dir, err)
	}

	session, err := client.NewSession()
	if err != nil {
		return err
	}
	defer session.Close()

	var stderr bytes.Buffer
	session.Stderr = &stderr

	stdin, err := session.StdinPipe()
	if err != nil {
		return err
	}

	// Prefer direct write; fall back to sudo tee if permission denied.
	cmd := fmt.Sprintf(
		`sh -c 'cat > %s && chmod %o %s' 2>/dev/null || sudo sh -c 'cat > %s && chmod %o %s'`,
		shellQuote(remotePath), mode, shellQuote(remotePath),
		shellQuote(remotePath), mode, shellQuote(remotePath),
	)

	if log != nil {
		log(fmt.Sprintf("upload %s (%d bytes) via ssh stdin", remotePath, len(data)))
	}

	if err := session.Start(cmd); err != nil {
		return fmt.Errorf("start upload: %w", err)
	}

	if _, err := io.Copy(stdin, bytes.NewReader(data)); err != nil {
		_ = stdin.Close()
		_ = session.Wait()
		return fmt.Errorf("write upload: %w (%s)", err, strings.TrimSpace(stderr.String()))
	}
	_ = stdin.Close()

	if err := session.Wait(); err != nil {
		msg := strings.TrimSpace(stderr.String())
		if msg == "" {
			msg = err.Error()
		}
		return fmt.Errorf("upload %s failed: %s", remotePath, msg)
	}
	return nil
}

func writeRemoteFile(client *ssh.Client, remotePath, content string, mode os.FileMode, log LogFunc) error {
	return uploadViaStdin(client, remotePath, []byte(content), mode, log)
}

func uploadBinary(client *ssh.Client, remotePath string, data []byte, log LogFunc) error {
	return uploadViaStdin(client, remotePath, data, 0o755, log)
}

// Provision installs Docker (if needed), uploads agent binary + compose, starts the node.
func Provision(n *models.Node, sshSecret string, agentBinary []byte, log LogFunc) error {
	if len(agentBinary) == 0 {
		return fmt.Errorf("agent binary is empty — build cmd/node for linux/amd64")
	}
	client, err := SSHConfig{
		Host: n.Host, Port: n.SSHPort, User: n.SSHUser,
		AuthType: n.SSHAuthType, Secret: sshSecret,
	}.dial()
	if err != nil {
		return fmt.Errorf("ssh connect: %w", err)
	}
	defer client.Close()

	if log != nil {
		log("SSH connected")
	}

	if _, err := run(client, `command -v docker >/dev/null 2>&1 || (curl -fsSL https://get.docker.com | sh)`, log); err != nil {
		return fmt.Errorf("install docker: %w", err)
	}
	if _, err := run(client, `sudo mkdir -p /opt/vpn-panel-node/data && (sudo chown -R "$(whoami):$(whoami)" /opt/vpn-panel-node || true)`, log); err != nil {
		return fmt.Errorf("prepare dir: %w", err)
	}

	if err := uploadBinary(client, "/opt/vpn-panel-node/vpn-node", agentBinary, log); err != nil {
		return err
	}
	if err := writeRemoteFile(client, "/opt/vpn-panel-node/docker-compose.yml", ComposeYAML(n), 0o644, log); err != nil {
		return err
	}

	if _, err := run(client, `cd /opt/vpn-panel-node && chmod +x vpn-node && (docker compose up -d --force-recreate || docker-compose up -d --force-recreate)`, log); err != nil {
		return fmt.Errorf("docker compose up: %w", err)
	}

	if log != nil {
		log(fmt.Sprintf("node agent started on :%d", n.NodePort))
		log("firewall: allow NODE_PORT only from panel IP")
	}
	return nil
}

func FindAgentBinary() (string, error) {
	candidates := []string{
		"/app/bin/vpn-node",
		"bin/vpn-node",
		"bin/vpn-node-linux-amd64",
		"vpn-node",
	}
	for _, c := range candidates {
		if st, err := os.Stat(c); err == nil && !st.IsDir() && st.Size() > 0 {
			return c, nil
		}
	}
	return "", fmt.Errorf("vpn-node binary not found (looked in %s)", strings.Join(candidates, ", "))
}

func LoadAgentBinary() ([]byte, error) {
	path, err := FindAgentBinary()
	if err != nil {
		return nil, err
	}
	return os.ReadFile(path)
}