gexpect.go

package gexpect

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"regexp"
	"time"
	"unicode/utf8"

	shell "github.com/coreos/etcd/Godeps/_workspace/src/github.com/kballard/go-shellquote"
	"github.com/coreos/etcd/Godeps/_workspace/src/github.com/kr/pty"
)

type ExpectSubprocess struct {
	Cmd          *exec.Cmd
	buf          *buffer
	outputBuffer []byte
}

type buffer struct {
	f       *os.File
	b       bytes.Buffer
	collect bool

	collection bytes.Buffer
}

func (buf *buffer) StartCollecting() {
	buf.collect = true
}

func (buf *buffer) StopCollecting() (result string) {
	result = string(buf.collection.Bytes())
	buf.collect = false
	buf.collection.Reset()
	return result
}

func (buf *buffer) Read(chunk []byte) (int, error) {
	nread := 0
	if buf.b.Len() > 0 {
		n, err := buf.b.Read(chunk)
		if err != nil {
			return n, err
		}
		if n == len(chunk) {
			return n, nil
		}
		nread = n
	}
	fn, err := buf.f.Read(chunk[nread:])
	return fn + nread, err
}

func (buf *buffer) ReadRune() (r rune, size int, err error) {
	l := buf.b.Len()

	chunk := make([]byte, utf8.UTFMax)
	if l > 0 {
		n, err := buf.b.Read(chunk)
		if err != nil {
			return 0, 0, err
		}
		if utf8.FullRune(chunk) {
			r, rL := utf8.DecodeRune(chunk)
			if n > rL {
				buf.PutBack(chunk[rL:n])
			}
			if buf.collect {
				buf.collection.WriteRune(r)
			}
			return r, rL, nil
		}
	}
	// else add bytes from the file, then try that
	for l < utf8.UTFMax {
		fn, err := buf.f.Read(chunk[l : l+1])
		if err != nil {
			return 0, 0, err
		}
		l = l + fn

		if utf8.FullRune(chunk) {
			r, rL := utf8.DecodeRune(chunk)
			if buf.collect {
				buf.collection.WriteRune(r)
			}
			return r, rL, nil
		}
	}
	return 0, 0, errors.New("File is not a valid UTF=8 encoding")
}

func (buf *buffer) PutBack(chunk []byte) {
	if len(chunk) == 0 {
		return
	}
	if buf.b.Len() == 0 {
		buf.b.Write(chunk)
		return
	}
	d := make([]byte, 0, len(chunk)+buf.b.Len())
	d = append(d, chunk...)
	d = append(d, buf.b.Bytes()...)
	buf.b.Reset()
	buf.b.Write(d)
}

func SpawnAtDirectory(command string, directory string) (*ExpectSubprocess, error) {
	expect, err := _spawn(command)
	if err != nil {
		return nil, err
	}
	expect.Cmd.Dir = directory
	return _start(expect)
}

func Command(command string) (*ExpectSubprocess, error) {
	expect, err := _spawn(command)
	if err != nil {
		return nil, err
	}
	return expect, nil
}

func (expect *ExpectSubprocess) Start() error {
	_, err := _start(expect)
	return err
}

func Spawn(command string) (*ExpectSubprocess, error) {
	expect, err := _spawn(command)
	if err != nil {
		return nil, err
	}
	return _start(expect)
}

func (expect *ExpectSubprocess) Close() error {
	return expect.Cmd.Process.Kill()
}

func (expect *ExpectSubprocess) AsyncInteractChannels() (send chan string, receive chan string) {
	receive = make(chan string)
	send = make(chan string)

	go func() {
		for {
			str, err := expect.ReadLine()
			if err != nil {
				close(receive)
				return
			}
			receive <- str
		}
	}()

	go func() {
		for {
			select {
			case sendCommand, exists := <-send:
				{
					if !exists {
						return
					}
					err := expect.Send(sendCommand)
					if err != nil {
						receive <- "gexpect Error: " + err.Error()
						return
					}
				}
			}
		}
	}()

	return
}

func (expect *ExpectSubprocess) ExpectRegex(regex string) (bool, error) {
	return regexp.MatchReader(regex, expect.buf)
}

func (expect *ExpectSubprocess) expectRegexFind(regex string, output bool) ([]string, string, error) {
	re, err := regexp.Compile(regex)
	if err != nil {
		return nil, "", err
	}
	expect.buf.StartCollecting()
	pairs := re.FindReaderSubmatchIndex(expect.buf)
	stringIndexedInto := expect.buf.StopCollecting()
	l := len(pairs)
	numPairs := l / 2
	result := make([]string, numPairs)
	for i := 0; i < numPairs; i += 1 {
		result[i] = stringIndexedInto[pairs[i*2]:pairs[i*2+1]]
	}
	// convert indexes to strings

	if len(result) == 0 {
		err = fmt.Errorf("ExpectRegex didn't find regex '%v'.", regex)
	}
	return result, stringIndexedInto, err
}

func (expect *ExpectSubprocess) expectTimeoutRegexFind(regex string, timeout time.Duration) (result []string, out string, err error) {
	t := make(chan bool)
	go func() {
		result, out, err = expect.ExpectRegexFindWithOutput(regex)
		t <- false
	}()
	go func() {
		time.Sleep(timeout)
		err = fmt.Errorf("ExpectRegex timed out after %v finding '%v'.\nOutput:\n%s", timeout, regex, expect.Collect())
		t <- true
	}()
	<-t
	return result, out, err
}

func (expect *ExpectSubprocess) ExpectRegexFind(regex string) ([]string, error) {
	result, _, err := expect.expectRegexFind(regex, false)
	return result, err
}

func (expect *ExpectSubprocess) ExpectTimeoutRegexFind(regex string, timeout time.Duration) ([]string, error) {
	result, _, err := expect.expectTimeoutRegexFind(regex, timeout)
	return result, err
}

func (expect *ExpectSubprocess) ExpectRegexFindWithOutput(regex string) ([]string, string, error) {
	return expect.expectRegexFind(regex, true)
}

func (expect *ExpectSubprocess) ExpectTimeoutRegexFindWithOutput(regex string, timeout time.Duration) ([]string, string, error) {
	return expect.expectTimeoutRegexFind(regex, timeout)
}

func buildKMPTable(searchString string) []int {
	pos := 2
	cnd := 0
	length := len(searchString)

	var table []int
	if length < 2 {
		length = 2
	}

	table = make([]int, length)
	table[0] = -1
	table[1] = 0

	for pos < len(searchString) {
		if searchString[pos-1] == searchString[cnd] {
			cnd += 1
			table[pos] = cnd
			pos += 1
		} else if cnd > 0 {
			cnd = table[cnd]
		} else {
			table[pos] = 0
			pos += 1
		}
	}
	return table
}

func (expect *ExpectSubprocess) ExpectTimeout(searchString string, timeout time.Duration) (e error) {
	result := make(chan error)
	go func() {
		result <- expect.Expect(searchString)
	}()
	select {
	case e = <-result:
	case <-time.After(timeout):
		e = fmt.Errorf("Expect timed out after %v waiting for '%v'.\nOutput:\n%s", timeout, searchString, expect.Collect())
	}
	return e
}

func (expect *ExpectSubprocess) Expect(searchString string) (e error) {
	chunk := make([]byte, len(searchString)*2)
	target := len(searchString)
	if expect.outputBuffer != nil {
		expect.outputBuffer = expect.outputBuffer[0:]
	}
	m := 0
	i := 0
	// Build KMP Table
	table := buildKMPTable(searchString)

	for {
		n, err := expect.buf.Read(chunk)

		if err != nil {
			return err
		}
		if expect.outputBuffer != nil {
			expect.outputBuffer = append(expect.outputBuffer, chunk[:n]...)
		}
		offset := m + i
		for m+i-offset < n {
			if searchString[i] == chunk[m+i-offset] {
				i += 1
				if i == target {
					unreadIndex := m + i - offset
					if len(chunk) > unreadIndex {
						expect.buf.PutBack(chunk[unreadIndex:])
					}
					return nil
				}
			} else {
				m += i - table[i]
				if table[i] > -1 {
					i = table[i]
				} else {
					i = 0
				}
			}
		}
	}
}

func (expect *ExpectSubprocess) Send(command string) error {
	_, err := io.WriteString(expect.buf.f, command)
	return err
}

func (expect *ExpectSubprocess) Capture() {
	if expect.outputBuffer == nil {
		expect.outputBuffer = make([]byte, 0)
	}
}

func (expect *ExpectSubprocess) Collect() []byte {
	collectOutput := make([]byte, len(expect.outputBuffer))
	copy(collectOutput, expect.outputBuffer)
	expect.outputBuffer = nil
	return collectOutput
}

func (expect *ExpectSubprocess) SendLine(command string) error {
	_, err := io.WriteString(expect.buf.f, command+"\r\n")
	return err
}

func (expect *ExpectSubprocess) Interact() {
	defer expect.Cmd.Wait()
	io.Copy(os.Stdout, &expect.buf.b)
	go io.Copy(os.Stdout, expect.buf.f)
	go io.Copy(expect.buf.f, os.Stdin)
}

func (expect *ExpectSubprocess) ReadUntil(delim byte) ([]byte, error) {
	join := make([]byte, 1, 512)
	chunk := make([]byte, 255)

	for {

		n, err := expect.buf.Read(chunk)

		if err != nil {
			return join, err
		}

		for i := 0; i < n; i++ {
			if chunk[i] == delim {
				if len(chunk) > i+1 {
					expect.buf.PutBack(chunk[i+1:])
				}
				return join, nil
			} else {
				join = append(join, chunk[i])
			}
		}
	}
}

func (expect *ExpectSubprocess) Wait() error {
	return expect.Cmd.Wait()
}

func (expect *ExpectSubprocess) ReadLine() (string, error) {
	str, err := expect.ReadUntil('\n')
	if err != nil {
		return "", err
	}
	return string(str), nil
}

func _start(expect *ExpectSubprocess) (*ExpectSubprocess, error) {
	f, err := pty.Start(expect.Cmd)
	if err != nil {
		return nil, err
	}
	expect.buf.f = f

	return expect, nil
}

func _spawn(command string) (*ExpectSubprocess, error) {
	wrapper := new(ExpectSubprocess)

	wrapper.outputBuffer = nil

	splitArgs, err := shell.Split(command)
	if err != nil {
		return nil, err
	}
	numArguments := len(splitArgs) - 1
	if numArguments < 0 {
		return nil, errors.New("gexpect: No command given to spawn")
	}
	path, err := exec.LookPath(splitArgs[0])
	if err != nil {
		return nil, err
	}

	if numArguments >= 1 {
		wrapper.Cmd = exec.Command(path, splitArgs[1:]...)
	} else {
		wrapper.Cmd = exec.Command(path)
	}
	wrapper.buf = new(buffer)

	return wrapper, nil
}