route/vendor/github.com/ThomasRooney/gexpect/gexpect.go

450 lines
9.3 KiB
Go
Raw Normal View History

// +build !windows
package gexpect
import (
"bytes"
"errors"
"fmt"
"io"
"os"
"os/exec"
"regexp"
"time"
"unicode/utf8"
shell "github.com/kballard/go-shellquote"
"github.com/kr/pty"
)
var (
ErrEmptySearch = errors.New("empty search string")
)
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[:n]) {
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[:l]) {
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 {
if err := expect.Cmd.Process.Kill(); err != nil {
return err
}
if err := expect.buf.f.Close(); err != nil {
return err
}
return nil
}
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)
} else {
// The number in pairs[1] is an index of a first
// character outside the whole match
putBackIdx := pairs[1]
if len(stringIndexedInto) > putBackIdx {
stringToPutBack := stringIndexedInto[putBackIdx:]
stringIndexedInto = stringIndexedInto[:putBackIdx]
expect.buf.PutBack([]byte(stringToPutBack))
}
}
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) {
target := len(searchString)
if target < 1 {
return ErrEmptySearch
}
chunk := make([]byte, target*2)
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 n == 0 && 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:n])
}
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, 0, 512)
chunk := make([]byte, 255)
for {
n, err := expect.buf.Read(chunk)
for i := 0; i < n; i++ {
if chunk[i] == delim {
if len(chunk) > i+1 {
expect.buf.PutBack(chunk[i+1:n])
}
return join, nil
} else {
join = append(join, chunk[i])
}
}
if err != nil {
return join, err
}
}
}
func (expect *ExpectSubprocess) Wait() error {
return expect.Cmd.Wait()
}
func (expect *ExpectSubprocess) ReadLine() (string, error) {
str, err := expect.ReadUntil('\n')
return string(str), err
}
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
}