1package main
  2
  3import (
  4	"context"
  5	"fmt"
  6	"io"
  7	"log/slog"
  8	"net"
  9	"os"
 10	"os/exec"
 11	"os/signal"
 12	"path/filepath"
 13	"strings"
 14	"syscall"
 15	"time"
 16
 17	"github.com/creack/pty"
 18	"golang.org/x/crypto/ssh"
 19	"golang.org/x/crypto/ssh/agent"
 20)
 21
 22// scrubEnvironment filters out sensitive environment variables (tokens, keys, secrets, credentials)
 23// while retaining essential system, terminal, and pici variables so interactive debugging works.
 24func scrubEnvironment(environ []string) []string {
 25	safePrefixes := []string{
 26		"PATH=", "USER=", "HOME=", "SHELL=", "TERM=", "LANG=", "LC_",
 27		"ZMX_", "PICI_", "PWD=", "TMPDIR=", "HOSTNAME=", "EDITOR=", "COLORTERM=",
 28		"SHLVL=", "LOGNAME=", "_=",
 29	}
 30
 31	sensitiveSubstrings := []string{
 32		"TOKEN", "SECRET", "PASSWORD", "PASSWD", "CREDENTIAL", "AUTH",
 33		"PRIVATE", "API_KEY", "ACCESS_KEY",
 34	}
 35
 36	var cleaned []string
 37	for _, env := range environ {
 38		parts := strings.SplitN(env, "=", 2)
 39		if len(parts) == 0 {
 40			continue
 41		}
 42		key := parts[0]
 43		keyUpper := strings.ToUpper(key)
 44
 45		// Explicitly scrub SSH private keys or pico auth keys
 46		if keyUpper == "PICO_SSH_KEY" || keyUpper == "SSH_PRIVATE_KEY" || keyUpper == "SSH_KEY" {
 47			continue
 48		}
 49
 50		// Check if it matches an explicitly preserved safe prefix
 51		isSafe := false
 52		for _, sp := range safePrefixes {
 53			if strings.HasPrefix(env, sp) {
 54				isSafe = true
 55				break
 56			}
 57		}
 58
 59		if isSafe {
 60			cleaned = append(cleaned, env)
 61			continue
 62		}
 63
 64		// Check if key contains sensitive words
 65		isSensitive := false
 66		for _, substr := range sensitiveSubstrings {
 67			if strings.Contains(keyUpper, substr) {
 68				isSensitive = true
 69				break
 70			}
 71		}
 72
 73		if !isSensitive {
 74			cleaned = append(cleaned, env)
 75		}
 76	}
 77	return cleaned
 78}
 79
 80// loadCertSigner parses the certificate file at certPath and wraps the key signer.
 81func loadCertSigner(signer ssh.Signer, certPath string) (ssh.Signer, error) {
 82	certBytes, err := os.ReadFile(certPath)
 83	if err != nil {
 84		return nil, fmt.Errorf("read cert file (%s): %w", certPath, err)
 85	}
 86	pubKey, _, _, _, err := ssh.ParseAuthorizedKey(certBytes)
 87	if err != nil {
 88		return nil, fmt.Errorf("parse cert (%s): %w", certPath, err)
 89	}
 90	cert, ok := pubKey.(*ssh.Certificate)
 91	if !ok {
 92		return nil, fmt.Errorf("public key at %s is not an SSH certificate", certPath)
 93	}
 94	certSigner, err := ssh.NewCertSigner(cert, signer)
 95	if err != nil {
 96		return nil, fmt.Errorf("create cert signer: %w", err)
 97	}
 98	return certSigner, nil
 99}
100
101// findSSHAuthMethod discovers an available SSH authentication method from:
102// 1. Explicit keyLocation (and optional explicit certificateLocation)
103// 2. Explicit PICO_SSH_KEY environment variable (file path or raw key content)
104// 3. Active SSH agent ($SSH_AUTH_SOCK)
105// 4. Default user SSH keys (~/.ssh/id_ed25519, ~/.ssh/id_rsa)
106func findSSHAuthMethod(keyLocation string, certificateLocation string) (ssh.AuthMethod, error) {
107	// 1. Specified key path (-pk)
108	if keyLocation != "" {
109		keyBytes, err := os.ReadFile(keyLocation)
110		if err != nil {
111			return nil, fmt.Errorf("read key file: %w", err)
112		}
113		signer, err := ssh.ParsePrivateKey(keyBytes)
114		if err != nil {
115			return nil, fmt.Errorf("parse private key: %w", err)
116		}
117		if certificateLocation != "" {
118			certSigner, err := loadCertSigner(signer, certificateLocation)
119			if err != nil {
120				return nil, err
121			}
122			return ssh.PublicKeys(certSigner), nil
123		}
124		return ssh.PublicKeys(signer), nil
125	}
126
127	// 2. PICO_SSH_KEY environment variable (path or raw text)
128	if envKey := os.Getenv("PICO_SSH_KEY"); envKey != "" {
129		var keyBytes []byte
130		if _, err := os.Stat(envKey); err == nil {
131			keyBytes, _ = os.ReadFile(envKey)
132		} else {
133			keyBytes = []byte(envKey)
134		}
135		signer, err := ssh.ParsePrivateKey(keyBytes)
136		if err == nil {
137			if certificateLocation != "" {
138				certSigner, err := loadCertSigner(signer, certificateLocation)
139				if err != nil {
140					return nil, err
141				}
142				return ssh.PublicKeys(certSigner), nil
143			}
144			return ssh.PublicKeys(signer), nil
145		}
146	}
147
148	// 3. SSH Agent
149	if sock := os.Getenv("SSH_AUTH_SOCK"); sock != "" {
150		if agentConn, err := net.Dial("unix", sock); err == nil {
151			ag := agent.NewClient(agentConn)
152			return ssh.PublicKeysCallback(ag.Signers), nil
153		}
154	}
155
156	// 4. Standard ~/.ssh/ files (default fallback when no key is explicitly passed)
157	if homeDir, err := os.UserHomeDir(); err == nil {
158		for _, keyName := range []string{"id_ed25519", "id_rsa"} {
159			keyPath := filepath.Join(homeDir, ".ssh", keyName)
160			if keyBytes, err := os.ReadFile(keyPath); err == nil {
161				if signer, err := ssh.ParsePrivateKey(keyBytes); err == nil {
162					if certificateLocation != "" {
163						certSigner, err := loadCertSigner(signer, certificateLocation)
164						if err != nil {
165							return nil, err
166						}
167						return ssh.PublicKeys(certSigner), nil
168					}
169					return ssh.PublicKeys(signer), nil
170				}
171			}
172		}
173	}
174
175	return nil, fmt.Errorf("no SSH authentication method found (provide -pk or set PICO_SSH_KEY)")
176}
177
178type DebugBridgeConfig struct {
179	Host         string
180	Topic        string
181	KeyLocation  string
182	CertLocation string
183	WorkDir      string
184	Timeout      time.Duration
185	Logger       *slog.Logger
186	Stdout       io.Writer
187}
188
189// startDebugBridge establishes a PTY relay over pipe.pico.sh
190func startDebugBridge(ctx context.Context, cfg DebugBridgeConfig) error {
191	if cfg.Host == "" {
192		cfg.Host = "pipe.pico.sh:22"
193	}
194	if cfg.Timeout <= 0 {
195		cfg.Timeout = 15 * time.Minute
196	}
197	if cfg.Stdout == nil {
198		cfg.Stdout = os.Stdout
199	}
200
201	auth, err := findSSHAuthMethod(cfg.KeyLocation, cfg.CertLocation)
202	if err != nil {
203		return fmt.Errorf("ssh auth: %w", err)
204	}
205
206	sshConfig := &ssh.ClientConfig{
207		User:            "pici",
208		Auth:            []ssh.AuthMethod{auth},
209		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
210		Timeout:         5 * time.Second,
211	}
212
213	fmt.Fprintf(cfg.Stdout, "🔌 connecting debug bridge to %s (topic: %s)...\n", cfg.Host, cfg.Topic) //nolint:errcheck
214
215	dialer := net.Dialer{Timeout: 5 * time.Second}
216	conn, err := dialer.DialContext(ctx, "tcp", cfg.Host)
217	if err != nil {
218		return fmt.Errorf("dial ssh (%s): %w", cfg.Host, err)
219	}
220	defer func() { _ = conn.Close() }()
221
222	sshConn, chans, reqs, err := ssh.NewClientConn(conn, cfg.Host, sshConfig)
223	if err != nil {
224		return fmt.Errorf("ssh handshake (%s): %w", cfg.Host, err)
225	}
226	client := ssh.NewClient(sshConn, chans, reqs)
227	defer func() { _ = client.Close() }()
228
229	session, err := client.NewSession()
230	if err != nil {
231		return fmt.Errorf("ssh session: %w", err)
232	}
233	defer func() { _ = session.Close() }()
234
235	sshStdin, err := session.StdinPipe()
236	if err != nil {
237		return fmt.Errorf("stdin pipe: %w", err)
238	}
239	sshStdout, err := session.StdoutPipe()
240	if err != nil {
241		return fmt.Errorf("stdout pipe: %w", err)
242	}
243	session.Stderr = io.Discard
244
245	if err := session.Start(fmt.Sprintf("pipe %s", cfg.Topic)); err != nil {
246		return fmt.Errorf("start pipe command: %w", err)
247	}
248
249	shell := os.Getenv("SHELL")
250	if shell == "" {
251		shell = "/bin/bash"
252	}
253	cmd := exec.Command(shell, "-i")
254	if cfg.WorkDir != "" {
255		cmd.Dir = cfg.WorkDir
256	}
257	cmd.Env = scrubEnvironment(os.Environ())
258
259	ptmx, err := pty.Start(cmd)
260	if err != nil {
261		return fmt.Errorf("start pty: %w", err)
262	}
263	defer func() {
264		_ = ptmx.Close()
265		if cmd.Process != nil {
266			_ = cmd.Process.Kill()
267		}
268	}()
269
270	fmt.Fprintf(cfg.Stdout, "\n🔧 ========================================================\n")                 //nolint:errcheck
271	fmt.Fprintf(cfg.Stdout, "🔧 PICI REMOTE DEBUG BRIDGE ACTIVE\n")                                            //nolint:errcheck
272	fmt.Fprintf(cfg.Stdout, "🔧 Run this command in your terminal to attach:\n")                               //nolint:errcheck
273	fmt.Fprintf(cfg.Stdout, "🔧   ssh -t pipe.pico.sh pipe %s\n", cfg.Topic)                                   //nolint:errcheck
274	fmt.Fprintf(cfg.Stdout, "🔧 Session will close when shell exits or after %s of idle time.\n", cfg.Timeout) //nolint:errcheck
275	fmt.Fprintf(cfg.Stdout, "🔧 ========================================================\n\n")                 //nolint:errcheck
276
277	ctxWithTimeout, cancel := context.WithTimeout(ctx, cfg.Timeout)
278	defer cancel()
279
280	sigCh := make(chan os.Signal, 1)
281	signal.Notify(sigCh, os.Interrupt, syscall.SIGTERM)
282	defer signal.Stop(sigCh)
283
284	errCh := make(chan error, 3)
285
286	go func() {
287		_, err := io.Copy(ptmx, sshStdout)
288		errCh <- err
289	}()
290
291	go func() {
292		_, err := io.Copy(sshStdin, ptmx)
293		errCh <- err
294	}()
295
296	go func() {
297		errCh <- cmd.Wait()
298	}()
299
300	select {
301	case <-ctxWithTimeout.Done():
302		fmt.Fprintf(cfg.Stdout, "\n⏰ debug session timed out after %s\n", cfg.Timeout) //nolint:errcheck
303	case <-sigCh:
304		fmt.Fprintf(cfg.Stdout, "\n⏹ debug session aborted by signal\n") //nolint:errcheck
305	case <-errCh:
306		fmt.Fprintf(cfg.Stdout, "\n🔌 debug session ended\n") //nolint:errcheck
307	}
308
309	return nil
310}
311
312func runDebug(cfg *Cfg, target string) error {
313	if target == "" {
314		return fmt.Errorf("target session or repo/job_id required (e.g. pici debug myrepo/12345 or pici debug ci.myrepo.12345.runner)")
315	}
316
317	// Check if target is a full session name or short session prefix
318	if strings.Contains(target, "/") {
319		parts := strings.SplitN(target, "/", 2)
320		repo, jobID := parts[0], parts[1]
321		prefix := fmt.Sprintf("ci.%s.%s.", repo, jobID)
322		localPrefix := fmt.Sprintf("local.%s.%s.", repo, jobID)
323		listOutput, err := exec.Command("zmx", "list").CombinedOutput()
324		if err != nil {
325			return fmt.Errorf("list zmx sessions: %w", err)
326		}
327		sessions := parseZMXList(string(listOutput))
328		var matches []SessionInfo
329		for _, s := range sessions {
330			if strings.HasPrefix(s.Name, prefix) || strings.HasPrefix(s.Name, localPrefix) {
331				matches = append(matches, s)
332			}
333		}
334		if len(matches) == 0 {
335			return fmt.Errorf("no active zmx sessions found for %s", target)
336		}
337		fmt.Fprintf(os.Stdout, "Found sessions for %s:\n", target) //nolint:errcheck
338		for _, m := range matches {
339			fmt.Fprintf(os.Stdout, "  • %s (pid %s, exit: %s)\n", m.Name, m.PID, m.ExitCode) //nolint:errcheck
340		}
341		target = matches[0].Name
342		fmt.Fprintf(os.Stdout, "\nAttaching to %s...\n", target) //nolint:errcheck
343	}
344
345	cmd := exec.Command("zmx", "attach", target)
346	cmd.Stdin = os.Stdin
347	cmd.Stdout = os.Stdout
348	cmd.Stderr = os.Stderr
349	return cmd.Run()
350}