122 lines
2.7 KiB
Go
122 lines
2.7 KiB
Go
|
package ssh
|
||
|
|
|
||
|
|
import (
|
||
|
|
"fmt"
|
||
|
|
"io"
|
||
|
|
|
||
|
|
"golang.org/x/crypto/ssh"
|
||
|
|
)
|
||
|
|
|
||
|
|
// Client wraps the SSH connection
|
||
|
|
type Client struct {
|
||
|
|
config *ssh.ClientConfig
|
||
|
|
address string
|
||
|
|
conn *ssh.Client
|
||
|
|
}
|
||
|
|
|
||
|
|
// Config contains SSH connection parameters
|
||
|
|
type Config struct {
|
||
|
|
Host string
|
||
|
|
Port int
|
||
|
|
User string
|
||
|
|
Password string // TODO: Add key-based auth support
|
||
|
|
}
|
||
|
|
|
||
|
|
// NewClient creates a new SSH client
|
||
|
|
func NewClient(cfg Config) *Client {
|
||
|
|
config := &ssh.ClientConfig{
|
||
|
|
User: cfg.User,
|
||
|
|
Auth: []ssh.AuthMethod{
|
||
|
|
ssh.Password(cfg.Password),
|
||
|
|
},
|
||
|
|
HostKeyCallback: ssh.InsecureIgnoreHostKey(), // TODO: Implement proper host key verification
|
||
|
|
}
|
||
|
|
|
||
|
|
return &Client{
|
||
|
|
config: config,
|
||
|
|
address: fmt.Sprintf("%s:%d", cfg.Host, cfg.Port),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// Connect establishes the SSH connection
|
||
|
|
func (c *Client) Connect() error {
|
||
|
|
conn, err := ssh.Dial("tcp", c.address, c.config)
|
||
|
|
if err != nil {
|
||
|
|
return fmt.Errorf("failed to connect: %w", err)
|
||
|
|
}
|
||
|
|
c.conn = conn
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// Close closes the SSH connection
|
||
|
|
func (c *Client) Close() error {
|
||
|
|
if c.conn != nil {
|
||
|
|
return c.conn.Close()
|
||
|
|
}
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// RunCommand runs a command and returns stdout
|
||
|
|
func (c *Client) RunCommand(cmd string) (string, error) {
|
||
|
|
session, err := c.conn.NewSession()
|
||
|
|
if err != nil {
|
||
|
|
return "", fmt.Errorf("failed to create session: %w", err)
|
||
|
|
}
|
||
|
|
defer session.Close()
|
||
|
|
|
||
|
|
output, err := session.CombinedOutput(cmd)
|
||
|
|
if err != nil {
|
||
|
|
return string(output), fmt.Errorf("command failed: %w", err)
|
||
|
|
}
|
||
|
|
return string(output), nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// StreamCommand runs a command and streams stdout to a writer
|
||
|
|
func (c *Client) StreamCommand(cmd string, stdout io.Writer, stderr io.Writer) error {
|
||
|
|
session, err := c.conn.NewSession()
|
||
|
|
if err != nil {
|
||
|
|
return fmt.Errorf("failed to create session: %w", err)
|
||
|
|
}
|
||
|
|
defer session.Close()
|
||
|
|
|
||
|
|
session.Stdout = stdout
|
||
|
|
session.Stderr = stderr
|
||
|
|
|
||
|
|
if err := session.Run(cmd); err != nil {
|
||
|
|
return fmt.Errorf("command failed: %w", err)
|
||
|
|
}
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// StartCommand starts a command and returns pipes for streaming
|
||
|
|
func (c *Client) StartCommand(cmd string) (io.Reader, io.Reader, func() error, error) {
|
||
|
|
session, err := c.conn.NewSession()
|
||
|
|
if err != nil {
|
||
|
|
return nil, nil, nil, fmt.Errorf("failed to create session: %w", err)
|
||
|
|
}
|
||
|
|
|
||
|
|
stdout, err := session.StdoutPipe()
|
||
|
|
if err != nil {
|
||
|
|
session.Close()
|
||
|
|
return nil, nil, nil, fmt.Errorf("failed to get stdout: %w", err)
|
||
|
|
}
|
||
|
|
|
||
|
|
stderr, err := session.StderrPipe()
|
||
|
|
if err != nil {
|
||
|
|
session.Close()
|
||
|
|
return nil, nil, nil, fmt.Errorf("failed to get stderr: %w", err)
|
||
|
|
}
|
||
|
|
|
||
|
|
if err := session.Start(cmd); err != nil {
|
||
|
|
session.Close()
|
||
|
|
return nil, nil, nil, fmt.Errorf("failed to start command: %w", err)
|
||
|
|
}
|
||
|
|
|
||
|
|
cleanup := func() error {
|
||
|
|
session.Signal(ssh.SIGTERM)
|
||
|
|
return session.Close()
|
||
|
|
}
|
||
|
|
|
||
|
|
return stdout, stderr, cleanup, nil
|
||
|
|
}
|