Repo iniciado
This commit is contained in:
@@ -0,0 +1,146 @@
|
||||
// Package ssh provides impure SSH command execution.
|
||||
package ssh
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"fmt"
|
||||
"net"
|
||||
"os"
|
||||
"time"
|
||||
|
||||
gossh "golang.org/x/crypto/ssh"
|
||||
|
||||
"github.com/enmanuel/agents/internal/config"
|
||||
"github.com/enmanuel/agents/pkg/tools"
|
||||
)
|
||||
|
||||
// Result holds the output of an SSH command execution.
|
||||
type Result struct {
|
||||
Stdout string
|
||||
Stderr string
|
||||
ExitCode int
|
||||
Err error
|
||||
}
|
||||
|
||||
// Executor runs SSH commands against configured targets.
|
||||
type Executor struct {
|
||||
cfg config.SSHCfg
|
||||
}
|
||||
|
||||
// NewExecutor creates an Executor from the SSH config section.
|
||||
func NewExecutor(cfg config.SSHCfg) *Executor {
|
||||
return &Executor{cfg: cfg}
|
||||
}
|
||||
|
||||
// Execute runs the SSH command described by spec. Impure.
|
||||
func (e *Executor) Execute(ctx context.Context, spec tools.SSHCommandSpec) Result {
|
||||
target, ok := e.cfg.Targets[spec.Target]
|
||||
if !ok {
|
||||
return Result{Err: fmt.Errorf("unknown SSH target: %s", spec.Target)}
|
||||
}
|
||||
|
||||
if len(target.Hosts) == 0 {
|
||||
return Result{Err: fmt.Errorf("no hosts for target: %s", spec.Target)}
|
||||
}
|
||||
|
||||
// Use first host (round-robin or load balancing can be added later)
|
||||
host := target.Hosts[0]
|
||||
user := target.User
|
||||
if user == "" {
|
||||
user = e.cfg.Defaults.User
|
||||
}
|
||||
port := target.Port
|
||||
if port == 0 {
|
||||
port = e.cfg.Defaults.Port
|
||||
}
|
||||
if port == 0 {
|
||||
port = 22
|
||||
}
|
||||
|
||||
keyEnv := target.KeyFileEnv
|
||||
if keyEnv == "" {
|
||||
keyEnv = e.cfg.Defaults.KeyFileEnv
|
||||
}
|
||||
|
||||
signer, err := loadSigner(keyEnv)
|
||||
if err != nil {
|
||||
return Result{Err: fmt.Errorf("load SSH key: %w", err)}
|
||||
}
|
||||
|
||||
sshCfg := &gossh.ClientConfig{
|
||||
User: user,
|
||||
Auth: []gossh.AuthMethod{gossh.PublicKeys(signer)},
|
||||
HostKeyCallback: gossh.InsecureIgnoreHostKey(), // TODO: use known_hosts
|
||||
Timeout: e.cfg.Defaults.Timeout,
|
||||
}
|
||||
if sshCfg.Timeout == 0 {
|
||||
sshCfg.Timeout = 10 * time.Second
|
||||
}
|
||||
|
||||
addr := fmt.Sprintf("%s:%d", host, port)
|
||||
conn, err := gossh.Dial("tcp", addr, sshCfg)
|
||||
if err != nil {
|
||||
return Result{Err: fmt.Errorf("ssh dial %s: %w", addr, err)}
|
||||
}
|
||||
defer conn.Close()
|
||||
|
||||
session, err := conn.NewSession()
|
||||
if err != nil {
|
||||
return Result{Err: fmt.Errorf("ssh session: %w", err)}
|
||||
}
|
||||
defer session.Close()
|
||||
|
||||
var stdout, stderr bytes.Buffer
|
||||
session.Stdout = &stdout
|
||||
session.Stderr = &stderr
|
||||
|
||||
// Respect context cancellation via a goroutine
|
||||
done := make(chan error, 1)
|
||||
go func() { done <- session.Run(spec.Command) }()
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
session.Signal(gossh.SIGTERM)
|
||||
return Result{Err: ctx.Err()}
|
||||
case err := <-done:
|
||||
code := 0
|
||||
if err != nil {
|
||||
var exitErr *gossh.ExitError
|
||||
if ok := asExitError(err, &exitErr); ok {
|
||||
code = exitErr.ExitStatus()
|
||||
} else {
|
||||
return Result{Err: err}
|
||||
}
|
||||
}
|
||||
return Result{
|
||||
Stdout: stdout.String(),
|
||||
Stderr: stderr.String(),
|
||||
ExitCode: code,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func loadSigner(keyFileEnv string) (gossh.Signer, error) {
|
||||
keyPath := os.Getenv(keyFileEnv)
|
||||
if keyPath == "" {
|
||||
return nil, fmt.Errorf("env var %s not set", keyFileEnv)
|
||||
}
|
||||
raw, err := os.ReadFile(keyPath)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return gossh.ParsePrivateKey(raw)
|
||||
}
|
||||
|
||||
// asExitError is a helper for type-asserting ssh.ExitError.
|
||||
func asExitError(err error, target **gossh.ExitError) bool {
|
||||
e, ok := err.(*gossh.ExitError)
|
||||
if ok {
|
||||
*target = e
|
||||
}
|
||||
return ok
|
||||
}
|
||||
|
||||
// Ensure net is used (for future jump host support)
|
||||
var _ = net.Dial
|
||||
Reference in New Issue
Block a user