450 lines
9.3 KiB
Go
450 lines
9.3 KiB
Go
|
// +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
|
||
|
}
|