readysite / pkg / platform / server.go
3.6 KB
server.go
package platform

import (
	"fmt"
	"os"
	"os/exec"
	"strings"
	"time"
)

// Server represents a cloud virtual machine
type Server struct {
	ID     string
	Name   string
	IP     string
	Size   string // e.g., "s-1vcpu-2gb"
	Region string // e.g., "nyc1"
	Status string // creating, active, off
}

// sshOptions returns common SSH options with connection multiplexing.
// ControlMaster reuses connections to avoid rate limiting.
func (s *Server) sshOptions() []string {
	return []string{
		"-o", "StrictHostKeyChecking=accept-new",
		"-o", "ControlMaster=auto",
		"-o", fmt.Sprintf("ControlPath=/tmp/ssh-readysite-%s", s.IP),
		"-o", "ControlPersist=60",
	}
}

// SSH executes a command on the server and returns output.
// SECURITY: Arguments are shell-escaped to prevent command injection.
func (s *Server) SSH(args ...string) (string, error) {
	cmdArgs := append(s.sshOptions(), "root@"+s.IP)

	// Shell-escape each argument to prevent command injection
	escapedArgs := make([]string, len(args))
	for i, arg := range args {
		escapedArgs[i] = shellEscape(arg)
	}
	cmdArgs = append(cmdArgs, strings.Join(escapedArgs, " "))

	cmd := exec.Command("ssh", cmdArgs...)
	output, err := cmd.CombinedOutput()
	if err != nil {
		return string(output), fmt.Errorf("ssh: %w: %s", err, output)
	}
	return strings.TrimSpace(string(output)), nil
}

// shellEscape escapes a string for safe use in shell commands.
func shellEscape(s string) string {
	// If the string is simple (alphanumeric, dash, underscore, dot, slash, colon),
	// no escaping needed
	safe := true
	for _, c := range s {
		if !((c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') ||
			(c >= '0' && c <= '9') || c == '-' || c == '_' ||
			c == '.' || c == '/' || c == ':' || c == '=' || c == ',') {
			safe = false
			break
		}
	}
	if safe && s != "" {
		return s
	}
	// Wrap in single quotes, escaping any single quotes within
	return "'" + strings.ReplaceAll(s, "'", "'\"'\"'") + "'"
}

// Copy uploads a file to the server via SCP
func (s *Server) Copy(local, remote string) error {
	cmdArgs := append(s.sshOptions(), local, "root@"+s.IP+":"+remote)
	cmd := exec.Command("scp", cmdArgs...)
	output, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("scp: %w: %s", err, output)
	}
	return nil
}

// Interactive opens an interactive SSH shell
func (s *Server) Interactive() error {
	cmdArgs := append(s.sshOptions(), "root@"+s.IP)
	cmd := exec.Command("ssh", cmdArgs...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	return cmd.Run()
}

// WaitForSSH waits until the server accepts SSH connections
func (s *Server) WaitForSSH(timeout time.Duration) error {
	deadline := time.Now().Add(timeout)
	for time.Now().Before(deadline) {
		_, err := s.SSH("echo", "ready")
		if err == nil {
			return nil
		}
		time.Sleep(5 * time.Second)
	}
	return ErrTimeout
}

// Connect executes a command with stdin/stdout/stderr connected
func (s *Server) Connect(args ...string) error {
	cmdArgs := append(s.sshOptions(), "root@"+s.IP)
	cmdArgs = append(cmdArgs, args...)

	cmd := exec.Command("ssh", cmdArgs...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	return cmd.Run()
}

// Write writes data to a file on the server
func (s *Server) Write(remotePath string, data []byte, executable bool) error {
	tmp, err := os.CreateTemp("", "readysite-*")
	if err != nil {
		return err
	}
	defer os.Remove(tmp.Name())

	if _, err := tmp.Write(data); err != nil {
		return err
	}
	tmp.Close()

	if err := s.Copy(tmp.Name(), remotePath); err != nil {
		return err
	}

	if executable {
		_, err := s.SSH("chmod", "+x", remotePath)
		return err
	}
	return nil
}
← Back