cli.go

package main

import (
	"context"
	"fmt"
	"strings"

	"github.com/golang/glog"
	"github.com/ziutek/telnet"
	"golang.org/x/net/trace"
)

type cliClient struct {
	conn *telnet.Conn

	username string
	password string

	loggedIn       bool
	promptHostname string
}

func newCliClient(c *telnet.Conn, username, password string) *cliClient {
	return &cliClient{
		conn:     c,
		username: username,
		password: password,
	}
}

func (c *cliClient) 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 *cliClient) 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 *cliClient) 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 *cliClient) 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 *cliClient) 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 *cliClient) 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
}