forked from hswaw/hscloud
Serge Bazanski
848db46bc0
Change-Id: Ieededb08a930d7b862575cc569d467cdd93e3e0d Reviewed-on: https://gerrit.hackerspace.pl/c/hscloud/+/1156 Reviewed-by: q3k <q3k@hackerspace.pl>
242 lines
5.7 KiB
Go
242 lines
5.7 KiB
Go
package cli
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"strings"
|
|
|
|
"github.com/ziutek/telnet"
|
|
"golang.org/x/net/trace"
|
|
)
|
|
|
|
type Client struct {
|
|
conn *telnet.Conn
|
|
|
|
username string
|
|
password string
|
|
|
|
loggedIn bool
|
|
promptHostname string
|
|
}
|
|
|
|
func NewClient(c *telnet.Conn, username, password string) *Client {
|
|
return &Client{
|
|
conn: c,
|
|
username: username,
|
|
password: password,
|
|
}
|
|
}
|
|
|
|
func (c *Client) readUntil(ctx context.Context, delims ...string) (string, error) {
|
|
chStr := make(chan string, 1)
|
|
chErr := make(chan error, 1)
|
|
go func() {
|
|
s, err := c.conn.ReadUntil(delims...)
|
|
if err != nil {
|
|
chErr <- err
|
|
return
|
|
}
|
|
chStr <- string(s)
|
|
}()
|
|
|
|
select {
|
|
case <-ctx.Done():
|
|
return "", fmt.Errorf("context done")
|
|
case err := <-chErr:
|
|
c.trace(ctx, "readUntil failed: %v", err)
|
|
return "", err
|
|
case s := <-chStr:
|
|
c.trace(ctx, "readUntil <- %q", s)
|
|
return s, nil
|
|
|
|
}
|
|
}
|
|
|
|
func (c *Client) readString(ctx context.Context, delim byte) (string, error) {
|
|
chStr := make(chan string, 1)
|
|
chErr := make(chan error, 1)
|
|
go func() {
|
|
s, err := c.conn.ReadString(delim)
|
|
if err != nil {
|
|
chErr <- err
|
|
return
|
|
}
|
|
chStr <- s
|
|
}()
|
|
|
|
select {
|
|
case <-ctx.Done():
|
|
return "", fmt.Errorf("context done")
|
|
case err := <-chErr:
|
|
c.trace(ctx, "readString failed: %v", err)
|
|
return "", err
|
|
case s := <-chStr:
|
|
c.trace(ctx, "readString <- %q", s)
|
|
return s, nil
|
|
|
|
}
|
|
}
|
|
|
|
func (c *Client) writeLine(ctx context.Context, s string) error {
|
|
n, err := c.conn.Write([]byte(s + "\n"))
|
|
if got, want := n, len(s)+1; got != want {
|
|
err = fmt.Errorf("wrote %d bytes out of %d", got, want)
|
|
}
|
|
if err != nil {
|
|
c.trace(ctx, "writeLine failed: %v", err)
|
|
return err
|
|
}
|
|
c.trace(ctx, "writeLine -> %q", s)
|
|
return nil
|
|
}
|
|
|
|
func (c *Client) trace(ctx context.Context, f string, parts ...interface{}) {
|
|
tr, ok := trace.FromContext(ctx)
|
|
if !ok {
|
|
//fmted := fmt.Sprintf(f, parts...)
|
|
//glog.Infof("[no trace] %s", fmted)
|
|
return
|
|
}
|
|
tr.LazyPrintf(f, parts...)
|
|
}
|
|
|
|
func (c *Client) logIn(ctx context.Context) error {
|
|
if c.loggedIn {
|
|
return nil
|
|
}
|
|
|
|
// Provide username.
|
|
prompt, err := c.readString(ctx, ':')
|
|
if err != nil {
|
|
return fmt.Errorf("could not read username prompt: %v", err)
|
|
}
|
|
if !strings.HasSuffix(prompt, "User:") {
|
|
return fmt.Errorf("invalid username prompt: %v", err)
|
|
}
|
|
if err := c.writeLine(ctx, c.username); err != nil {
|
|
return fmt.Errorf("could not write username: %v")
|
|
}
|
|
|
|
// Provide password.
|
|
prompt, err = c.readString(ctx, ':')
|
|
if err != nil {
|
|
return fmt.Errorf("could not read password prompt: %v", err)
|
|
}
|
|
if !strings.HasSuffix(prompt, "Password:") {
|
|
return fmt.Errorf("invalid password prompt: %v", err)
|
|
}
|
|
if err := c.writeLine(ctx, c.password); err != nil {
|
|
return fmt.Errorf("could not write password: %v")
|
|
}
|
|
|
|
// Get unprivileged prompt.
|
|
prompt, err = c.readString(ctx, '>')
|
|
if err != nil {
|
|
return fmt.Errorf("could not read unprivileged prompt: %v", err)
|
|
}
|
|
|
|
parts := strings.Split(prompt, "\r\n")
|
|
c.promptHostname = strings.TrimSuffix(parts[len(parts)-1], ">")
|
|
|
|
// Enable privileged mode.
|
|
|
|
if err := c.writeLine(ctx, "enable"); err != nil {
|
|
return fmt.Errorf("could not write enable: %v")
|
|
}
|
|
|
|
// Provide password (again)
|
|
prompt, err = c.readString(ctx, ':')
|
|
if err != nil {
|
|
return fmt.Errorf("could not read password prompt: %v", err)
|
|
}
|
|
if !strings.HasSuffix(prompt, "Password:") {
|
|
return fmt.Errorf("invalid password prompt: %v", err)
|
|
}
|
|
if err := c.writeLine(ctx, c.password); err != nil {
|
|
return fmt.Errorf("could not write password: %v")
|
|
}
|
|
|
|
// Get privileged prompt.
|
|
prompt, err = c.readString(ctx, '#')
|
|
if err != nil {
|
|
return fmt.Errorf("could not read privileged prompt: %v", err)
|
|
}
|
|
|
|
if !strings.HasSuffix(prompt, c.promptHostname+"#") {
|
|
return fmt.Errorf("unexpected privileged prompt: %v", prompt)
|
|
}
|
|
|
|
// Disable pager.
|
|
if err := c.writeLine(ctx, "terminal length 0"); err != nil {
|
|
return fmt.Errorf("could not diable pager: %v", err)
|
|
}
|
|
prompt, err = c.readString(ctx, '#')
|
|
if err != nil {
|
|
return fmt.Errorf("could not disable pager: %v", err)
|
|
}
|
|
if !strings.HasSuffix(prompt, c.promptHostname+"#") {
|
|
return fmt.Errorf("unexpected privileged prompt: %v", prompt)
|
|
}
|
|
|
|
// Success!
|
|
c.loggedIn = true
|
|
c.trace(ctx, "logged into %v", c.promptHostname)
|
|
return nil
|
|
}
|
|
|
|
func (c *Client) RunCommand(ctx context.Context, command string) ([]string, string, error) {
|
|
if err := c.logIn(ctx); err != nil {
|
|
return nil, "", fmt.Errorf("could not log in: %v", err)
|
|
}
|
|
|
|
// First, synchronize to prompt.
|
|
attempts := 3
|
|
for {
|
|
c.writeLine(ctx, "")
|
|
line, err := c.readString(ctx, '\n')
|
|
if err != nil {
|
|
return nil, "", fmt.Errorf("while synchronizing to prompt: %v", err)
|
|
}
|
|
line = strings.Trim(line, "\r\n")
|
|
if strings.HasSuffix(line, c.promptHostname+"#") {
|
|
break
|
|
}
|
|
|
|
attempts -= 1
|
|
if attempts == 0 {
|
|
return nil, "", fmt.Errorf("could not find prompt, last result %q", line)
|
|
}
|
|
}
|
|
|
|
// Send comand.
|
|
c.writeLine(ctx, command)
|
|
|
|
// First, read until prompt again.
|
|
if _, err := c.readUntil(ctx, c.promptHostname+"#"); err != nil {
|
|
return nil, "", fmt.Errorf("could not get command hostname echo: %v", err)
|
|
}
|
|
|
|
loopback, err := c.readUntil(ctx, "\r\n")
|
|
if err != nil {
|
|
return nil, "", fmt.Errorf("could not get command loopback: %v", err)
|
|
}
|
|
loopback = strings.Trim(loopback, "\r\n")
|
|
c.trace(ctx, "effective command: %q", loopback)
|
|
|
|
// Read until we have a standalone prompt with no newline afterwards.
|
|
data, err := c.readUntil(ctx, c.promptHostname+"#")
|
|
if err != nil {
|
|
return nil, "", fmt.Errorf("could not get command results: %v", err)
|
|
}
|
|
|
|
lines := []string{}
|
|
for _, line := range strings.Split(data, "\r\n") {
|
|
if line == c.promptHostname+"#" {
|
|
break
|
|
}
|
|
lines = append(lines, line)
|
|
}
|
|
c.trace(ctx, "command %q returned lines: %v", command, lines)
|
|
|
|
return lines, loopback, nil
|
|
}
|