blob: 0e3d36ed22cf2ed63fc8342ceb74c501f75711bb [file] [edit]
// Copyright 2016 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
// Package expect is a Go version of the classic TCL Expect.
package expect
import (
"bytes"
"errors"
"fmt"
"io"
"os/exec"
"regexp"
"strconv"
"strings"
"sync"
"syscall"
"time"
log "github.com/golang/glog"
stdlog "log"
"golang.org/x/crypto/ssh"
"google.golang.org/grpc/codes"
"github.com/google/goterm/term"
)
// DefaultTimeout is the default Expect timeout.
const DefaultTimeout = 60 * time.Second
const (
bufferSize = 8192 // bufferSize sets the size of the io buffers.
checkDuration = 2 * time.Second // checkDuration how often to check for new output.
)
type Status struct {
code codes.Code
msg string
}
func NewStatus(code codes.Code, msg string) *Status {
return &Status{code, msg}
}
func NewStatusf(code codes.Code, format string, a ...interface{}) *Status {
return NewStatus(code, fmt.Sprintf(fmt.Sprintf(format, a...)))
}
func (s *Status) Err() error {
if s == nil || s.code == codes.OK {
return nil
}
return s
}
func (s *Status) Error() string {
return s.msg
}
// Option represents one Expecter option.
type Option func(*GExpect) Option
// CheckDuration changes the default duration checking for new incoming data.
func CheckDuration(d time.Duration) Option {
return func(e *GExpect) Option {
prev := e.chkDuration
e.chkDuration = d
return CheckDuration(prev)
}
}
// NoCheck turns off the Expect alive checks.
func NoCheck() Option {
return changeChk(func(*GExpect) bool {
return true
})
}
// DebugCheck adds logging to the check function.
// The check function for the spawners are called at creation/timeouts and I/O so can
// be usable for printing current state during debugging.
func DebugCheck(l *stdlog.Logger) Option {
lg := log.Infof
if l != nil {
lg = l.Printf
}
return func(e *GExpect) Option {
prev := e.chk
e.chkMu.Lock()
e.chk = func(ge *GExpect) bool {
res := prev(ge)
ge.mu.Lock()
lg("chk: %t, ge: %v", res, ge)
ge.mu.Unlock()
return res
}
e.chkMu.Unlock()
return changeChk(prev)
}
}
// ChangeCheck changes the Expect check function.
func ChangeCheck(f func() bool) Option {
return changeChk(func(*GExpect) bool {
return f()
})
}
func changeChk(f func(*GExpect) bool) Option {
return func(e *GExpect) Option {
prev := e.chk
e.chkMu.Lock()
e.chk = f
e.chkMu.Unlock()
return changeChk(prev)
}
}
// BatchCommands.
const (
// BatchSend for invoking Send in a batch
BatchSend = iota
// BatchExpect for invoking Expect in a batch
BatchExpect
// BatchSwitchCase for invoking ExpectSwitchCase in a batch
BatchSwitchCase
)
// TimeoutError is the error returned by all Expect functions upon timer expiry.
type TimeoutError int
// Error implements the Error interface.
func (t TimeoutError) Error() string {
return fmt.Sprintf("expect: timer expired after %d seconds", time.Duration(t)/time.Second)
}
// BatchRes returned from ExpectBatch for every Expect command executed.
type BatchRes struct {
// Idx is used to match the result with the []Batcher commands sent in.
Idx int
// Out output buffer for the expect command at Batcher[Idx].
Output string
// Match regexp matches for expect command at Batcher[Idx].
Match []string
}
// Batcher interface is used to make it more straightforward and readable to create
// batches of Expects.
//
// var batch = []Batcher{
// &BExpT{"password",8},
// &BSnd{"password\n"},
// &BExp{"olakar@router>"},
// &BSnd{ "show interface description\n"},
// &BExp{ "olakar@router>"},
// }
//
// var batchSwCaseReplace = []Batcher{
// &BCasT{[]Caser{
// &BCase{`([0-9]) -- .*\(MASTER\)`, `\1` + "\n"}}, 1},
// &BExp{`prompt/>`},
// }
type Batcher interface {
// cmd returns the Batch command.
Cmd() int
// Arg returns the command argument.
Arg() string
// Timeout returns the timeout duration for the command , <0 gives default value.
Timeout() time.Duration
// Cases returns the Caser structure for SwitchCase commands.
Cases() []Caser
}
// BExp implements the Batcher interface for Expect commands using the default timeout.
type BExp struct {
// R contains the Expect command regular expression.
R string
}
// Cmd returns the Expect command (BatchExpect).
func (be *BExp) Cmd() int {
return BatchExpect
}
// Arg returns the Expect regular expression.
func (be *BExp) Arg() string {
return be.R
}
// Timeout always returns -1 which sets it to the value used to call the ExpectBatch function.
func (be *BExp) Timeout() time.Duration {
return -1
}
// Cases always returns nil for the Expect command.
func (be *BExp) Cases() []Caser {
return nil
}
// BExpT implements the Batcher interface for Expect commands adding a timeout option to the BExp
// type.
type BExpT struct {
// R contains the Expect command regular expression.
R string
// T holds the Expect command timeout in seconds.
T int
}
// Cmd returns the Expect command (BatchExpect).
func (bt *BExpT) Cmd() int {
return BatchExpect
}
// Timeout returns the timeout in seconds.
func (bt *BExpT) Timeout() time.Duration {
return time.Duration(bt.T) * time.Second
}
// Arg returns the Expect regular expression.
func (bt *BExpT) Arg() string {
return bt.R
}
// Cases always return nil for the Expect command.
func (bt *BExpT) Cases() []Caser {
return nil
}
// BSnd implements the Batcher interface for Send commands.
type BSnd struct {
S string
}
// Cmd returns the Send command(BatchSend).
func (bs *BSnd) Cmd() int {
return BatchSend
}
// Arg returns the data to be sent.
func (bs *BSnd) Arg() string {
return bs.S
}
// Timeout always returns 0 , Send doesn't have a timeout.
func (bs *BSnd) Timeout() time.Duration {
return 0
}
// Cases always returns nil , not used for Send commands.
func (bs *BSnd) Cases() []Caser {
return nil
}
// BCas implements the Batcher interface for SwitchCase commands.
type BCas struct {
// C holds the Caser array for the SwitchCase command.
C []Caser
}
// Cmd returns the SwitchCase command(BatchSwitchCase).
func (bc *BCas) Cmd() int {
return BatchSwitchCase
}
// Arg returns an empty string , not used for SwitchCase.
func (bc *BCas) Arg() string {
return ""
}
// Timeout returns -1 , setting it to the default value.
func (bc *BCas) Timeout() time.Duration {
return -1
}
// Cases returns the Caser structure.
func (bc *BCas) Cases() []Caser {
return bc.C
}
// BCasT implements the Batcher interfacs for SwitchCase commands, adding a timeout option
// to the BCas type.
type BCasT struct {
// Cs holds the Caser array for the SwitchCase command.
C []Caser
// Tout holds the SwitchCase timeout in seconds.
T int
}
// Timeout returns the timeout in seconds.
func (bct *BCasT) Timeout() time.Duration {
return time.Duration(bct.T) * time.Second
}
// Cmd returns the SwitchCase command(BatchSwitchCase).
func (bct *BCasT) Cmd() int {
return BatchSwitchCase
}
// Arg returns an empty string , not used for SwitchCase.
func (bct *BCasT) Arg() string {
return ""
}
// Cases returns the Caser structure.
func (bct *BCasT) Cases() []Caser {
return bct.C
}
// Tag represents the state for a Caser.
type Tag int32
const (
// OKTag marks the desired state was reached.
OKTag = Tag(iota)
// FailTag means reaching this state will fail the Switch/Case.
FailTag
// ContinueTag will recheck for matches.
ContinueTag
// NextTag skips match and continues to the next one.
NextTag
// NoTag signals no tag was set for this case.
NoTag
)
// OK returns the OK Tag and status.
func OK() func() (Tag, *Status) {
return func() (Tag, *Status) {
return OKTag, NewStatus(codes.OK, "state reached")
}
}
// Fail returns Fail Tag and status.
func Fail(s *Status) func() (Tag, *Status) {
return func() (Tag, *Status) {
return FailTag, s
}
}
// Continue returns the Continue Tag and status.
func Continue(s *Status) func() (Tag, *Status) {
return func() (Tag, *Status) {
return ContinueTag, s
}
}
// Next returns the Next Tag and status.
func Next() func() (Tag, *Status) {
return func() (Tag, *Status) {
return NextTag, NewStatus(codes.Unimplemented, "Next returns not implemented")
}
}
// LogContinue logs the message and returns the Continue Tag and status.
func LogContinue(msg string, s *Status) func() (Tag, *Status) {
return func() (Tag, *Status) {
log.Info(msg)
return ContinueTag, s
}
}
// Caser is an interface for ExpectSwitchCase and Batch to be able to handle
// both the Case struct and the more script friendly BCase struct.
type Caser interface {
// RE returns a compiled regexp
RE() (*regexp.Regexp, error)
// Send returns the send string
String() string
// Tag returns the Tag.
Tag() (Tag, *Status)
// Retry returns true if there are retries left.
Retry() bool
}
// Case used by the ExpectSwitchCase to take different Cases.
// Implements the Caser interface.
type Case struct {
// R is the compiled regexp to match.
R *regexp.Regexp
// S is the string to send if Regexp matches.
S string
// T is the Tag for this Case.
T func() (Tag, *Status)
// Rt specifies number of times to retry, only used for cases tagged with Continue.
Rt int
}
// Tag returns the tag for this case.
func (c *Case) Tag() (Tag, *Status) {
if c.T == nil {
return NoTag, NewStatus(codes.OK, "no Tag set")
}
return c.T()
}
// RE returns the compiled regular expression.
func (c *Case) RE() (*regexp.Regexp, error) {
return c.R, nil
}
// Retry decrements the Retry counter and checks if there are any retries left.
func (c *Case) Retry() bool {
defer func() { c.Rt-- }()
return c.Rt > 0
}
// Send returns the string to send if regexp matches
func (c *Case) String() string {
return c.S
}
// BCase with just a string is a bit more friendly to scripting.
// Implements the Caser interface.
type BCase struct {
// R contains the string regular expression.
R string
// S contains the string to be sent if R matches.
S string
// T contains the Tag.
T func() (Tag, *Status)
// Rt contains the number of retries.
Rt int
}
// RE returns the compiled regular expression.
func (b *BCase) RE() (*regexp.Regexp, error) {
if b.R == "" {
return nil, nil
}
return regexp.Compile(b.R)
}
// Send returns the string to send.
func (b *BCase) String() string {
return b.S
}
// Tag returns the BCase Tag.
func (b *BCase) Tag() (Tag, *Status) {
if b.T == nil {
return NoTag, NewStatus(codes.OK, "no Tag set")
}
return b.T()
}
// Retry decrements the Retry counter and checks if there are any retries left.
func (b *BCase) Retry() bool {
b.Rt--
return b.Rt > -1
}
// Expecter interface primarily to make testing easier.
type Expecter interface {
// Expect reads output from a spawned session and tries matching it with the provided regular expression.
// It returns all output found until match.
Expect(*regexp.Regexp, time.Duration) (string, []string, error)
// ExpectBatch takes an array of BatchEntries and runs through them in order. For every Expect
// command a BatchRes entry is created with output buffer and sub matches.
// Failure of any of the batch commands will stop the execution, returning the results up to the
// failure.
ExpectBatch([]Batcher, time.Duration) ([]BatchRes, error)
// ExpectSwitchCase makes it possible to Expect with multiple regular expressions and actions. Returns the
// full output and submatches of the commands together with an index for the matching Case.
ExpectSwitchCase([]Caser, time.Duration) (string, []string, int, error)
// Send sends data into the spawned session.
Send(string) error
// Close closes the spawned session and files.
Close() error
}
// GExpect implements the Expecter interface.
type GExpect struct {
// pty holds the virtual terminal used to interact with the spawned commands.
pty *term.PTY
// cmd contains the cmd information for the spawned process.
cmd *exec.Cmd
ssh *ssh.Session
// snd is the channel used by the Send command to send data into the spawned command.
snd chan string
// rcv is used to signal the Expect commands that new data arrived.
rcv chan struct{}
// chkMu lock protecting the check function.
chkMu sync.RWMutex
// chk contains the function to check if the spawned command is alive.
chk func(*GExpect) bool
// cls contains the function to close spawned command.
cls func(*GExpect) error
// timeout contains the default timeout for a spawned command.
timeout time.Duration
// chkDuration contains the duration between checks for new incoming data.
chkDuration time.Duration
// mu protects the output buffer. It must be held for any operations on out.
mu sync.Mutex
out bytes.Buffer
}
// String implements the stringer interface.
func (e *GExpect) String() string {
res := fmt.Sprintf("%p: ", e)
if e.pty != nil {
_, name := e.pty.PTSName()
res += fmt.Sprintf("pty: %s ", name)
}
switch {
case e.cmd != nil:
res += fmt.Sprintf("cmd: %s(%d) ", e.cmd.Path, e.cmd.Process.Pid)
case e.ssh != nil:
res += fmt.Sprint("ssh session ")
}
res += fmt.Sprintf("buf: %q", e.out.String())
return res
}
// ExpectBatch takes an array of BatchEntry and executes them in order filling in the BatchRes
// array for any Expect command executed.
func (e *GExpect) ExpectBatch(batch []Batcher, timeout time.Duration) ([]BatchRes, error) {
res := []BatchRes{}
for i, b := range batch {
switch b.Cmd() {
case BatchExpect:
re, err := regexp.Compile(b.Arg())
if err != nil {
return res, err
}
to := b.Timeout()
if to < 0 {
to = timeout
}
out, match, err := e.Expect(re, to)
res = append(res, BatchRes{i, out, match})
if err != nil {
return res, err
}
case BatchSend:
if err := e.Send(b.Arg()); err != nil {
return res, err
}
case BatchSwitchCase:
to := b.Timeout()
if to < 0 {
to = timeout
}
out, match, _, err := e.ExpectSwitchCase(b.Cases(), to)
res = append(res, BatchRes{i, out, match})
if err != nil {
return res, err
}
default:
return res, errors.New("unknown command:" + strconv.Itoa(b.Cmd()))
}
}
return res, nil
}
func (e *GExpect) check() bool {
e.chkMu.RLock()
defer e.chkMu.RUnlock()
return e.chk(e)
}
// ExpectSwitchCase checks each Case against the accumulated out buffer, sending specified
// string back. Leaving Send empty will Send nothing to the process.
// Substring expansion can be used eg.
// Case{`vf[0-9]{2}.[a-z]{3}[0-9]{2}\.net).*UP`,`show arp \1`}
// Given: vf11.hnd01.net UP 35 (4) 34 (4) CONNECTED 0 0/0
// Would send: show arp vf11.hnd01.net
func (e *GExpect) ExpectSwitchCase(cs []Caser, timeout time.Duration) (string, []string, int, error) {
// Compile all regexps
rs := make([]*regexp.Regexp, 0, len(cs))
for _, c := range cs {
re, err := c.RE()
if err != nil {
return "", []string{""}, -1, err
}
rs = append(rs, re)
}
// Setup timeouts
// timeout == 0 => Just dump the buffer and exit.
// timeout < 0 => Set default value.
if timeout < 0 {
timeout = e.timeout
}
timer := time.NewTimer(timeout)
check := e.chkDuration
// Check if any new data arrived every checkDuration interval.
// If timeout/4 is less than the checkout interval we set the checkout to
// timeout/4. If timeout ends up being 0 we bump it to one to keep the Ticker from
// panicking.
// All this b/c of the unreliable channel send setup in the read function,making it
// possible for Expect* functions to miss the rcv signal.
//
// from read():
// // Ping Expect function
// select {
// case e.rcv <- struct{}{}:
// default:
// }
//
// A signal is only sent if any Expect function is running. Expect could miss it
// while playing around with buffers and matching regular expressions.
if timeout>>2 < check {
check = timeout >> 2
if check <= 0 {
check = 1
}
}
chTicker := time.NewTicker(check)
// Read in current data and start actively check for matches.
var tbuf bytes.Buffer
if _, err := io.Copy(&tbuf, e); err != nil {
return tbuf.String(), nil, -1, fmt.Errorf("io.Copy failed: %v", err)
}
for {
L1:
for i, c := range cs {
if rs[i] == nil {
continue
}
match := rs[i].FindStringSubmatch(tbuf.String())
if match == nil {
continue
}
t, s := c.Tag()
if t == NextTag && !c.Retry() {
continue
}
// Clear the buffer directly after match.
o := tbuf.String()
tbuf.Reset()
st := c.String()
// Replace the submatches \[0-9]+ in the send string.
if len(match) > 1 && len(st) > 0 {
for i := 1; i < len(match); i++ {
// \(submatch) will be expanded in the Send string.
// To escape use \\(number).
si := strconv.Itoa(i)
r := strings.NewReplacer(`\\`+si, `\`+si, `\`+si, `\\`+si)
st = r.Replace(st)
st = strings.Replace(st, `\\`+si, match[i], -1)
}
}
// Don't send anything if string is empty.
if st != "" {
if err := e.Send(st); err != nil {
return o, match, i, fmt.Errorf("failed to send: %q err: %v", st, err)
}
}
// Tag handling.
switch t {
case OKTag, FailTag, NoTag:
return o, match, i, s.Err()
case ContinueTag:
if !c.Retry() {
return o, match, i, s.Err()
}
break L1
case NextTag:
break L1
default:
s = NewStatusf(codes.Unknown, "Tag: %d unknown, err: %v", t, s)
}
return o, match, i, s.Err()
}
if !e.check() {
nr, err := io.Copy(&tbuf, e)
if err != nil {
return tbuf.String(), nil, -1, fmt.Errorf("io.Copy failed: %v", err)
}
if nr == 0 {
return tbuf.String(), nil, -1, errors.New("expect: Process not running")
}
}
select {
case <-timer.C:
// Expect timeout.
nr, err := io.Copy(&tbuf, e)
if err != nil {
return tbuf.String(), nil, -1, fmt.Errorf("io.Copy failed: %v", err)
}
// If we got no new data we return otherwise give it another chance to match.
if nr == 0 {
return tbuf.String(), nil, -1, TimeoutError(timeout)
}
timer = time.NewTimer(timeout)
case <-chTicker.C:
// Periodical timer to make sure data is handled in case the <-e.rcv channel
// was missed.
if _, err := io.Copy(&tbuf, e); err != nil {
return tbuf.String(), nil, -1, fmt.Errorf("io.Copy failed: %v", err)
}
case <-e.rcv:
// Data to fetch.
if _, err := io.Copy(&tbuf, e); err != nil {
return tbuf.String(), nil, -1, fmt.Errorf("io.Copy failed: %v", err)
}
}
}
}
// GenOptions contains the options needed to set up a generic Spawner.
type GenOptions struct {
// In is where Expect Send messages will be written.
In io.WriteCloser
// Out will be read and matched by the expecter.
Out io.Reader
// Wait is used by expect to know when the session is over and cleanup of io Go routines should happen.
Wait func() error
// Close will be called as part of the expect Close, should normally include a Close of the In WriteCloser.
Close func() error
// Check is called everytime a Send or Expect function is called to makes sure the session is still running.
Check func() bool
}
// SpawnGeneric is used to write generic Spawners. It returns an Expecter. The returned channel will give the return
// status of the spawned session, in practice this means the return value of the provided Wait function.
func SpawnGeneric(opt *GenOptions, timeout time.Duration, opts ...Option) (*GExpect, <-chan error, error) {
switch {
case opt == nil:
return nil, nil, errors.New("GenOptions is <nil>")
case opt.In == nil:
return nil, nil, errors.New("In can't be <nil>")
case opt.Out == nil:
return nil, nil, errors.New("Out can't be <nil>")
case opt.Wait == nil:
return nil, nil, errors.New("Wait can't be <nil>")
case opt.Close == nil:
return nil, nil, errors.New("Close can't be <nil>")
case opt.Check == nil:
return nil, nil, errors.New("Check can't be <nil>")
}
if timeout < 1 {
timeout = DefaultTimeout
}
e := &GExpect{
rcv: make(chan struct{}),
snd: make(chan string),
timeout: timeout,
chkDuration: checkDuration,
cls: func(e *GExpect) error {
return opt.Close()
},
chk: func(e *GExpect) bool {
return opt.Check()
},
}
for _, o := range opts {
o(e)
}
errCh := make(chan error, 1)
go e.waitForSession(errCh, opt.Wait, opt.In, opt.Out, nil)
return e, errCh, nil
}
// SpawnFake spawns an expect.Batcher.
func SpawnFake(b []Batcher, timeout time.Duration, opt ...Option) (*GExpect, <-chan error, error) {
rr, rw := io.Pipe()
wr, ww := io.Pipe()
done := make(chan struct{})
srv, _, err := SpawnGeneric(&GenOptions{
In: ww,
Out: rr,
Wait: func() error {
<-done
return nil
},
Close: func() error {
return ww.Close()
},
Check: func() bool { return true },
}, timeout, opt...)
if err != nil {
return nil, nil, err
}
go func() {
res, err := srv.ExpectBatch(b, timeout)
if err != nil {
log.Warningf("ExpectBatch(%v,%v) failed: %v, out: %v", b, timeout, err, res)
}
close(done)
}()
return SpawnGeneric(&GenOptions{
In: rw,
Out: wr,
Close: func() error {
return rw.Close()
},
Check: func() bool { return true },
Wait: func() error {
<-done
return nil
},
}, timeout, opt...)
}
// Spawn starts a new process and collects the output. The error channel returns the result of the
// command Spawned when it finishes.
func Spawn(command string, timeout time.Duration, opts ...Option) (*GExpect, <-chan error, error) {
pty, err := term.OpenPTY()
if err != nil {
return nil, nil, err
}
var t term.Termios
t.Raw()
t.Set(pty.Slave)
if timeout < 1 {
timeout = DefaultTimeout
}
// Get the command up and running
args := strings.Fields(command)
cmd := exec.Command(args[0], args[1:]...)
// This ties the commands Stdin,Stdout & Stderr to the virtual terminal we created
cmd.Stdin, cmd.Stdout, cmd.Stderr = pty.Slave, pty.Slave, pty.Slave
// New process needs to be the process leader and control of a tty
cmd.SysProcAttr = &syscall.SysProcAttr{
Setsid: true,
Setctty: true}
e := &GExpect{
rcv: make(chan struct{}),
snd: make(chan string),
cmd: cmd,
timeout: timeout,
chkDuration: checkDuration,
pty: pty,
cls: func(e *GExpect) error {
if e.cmd != nil {
return e.cmd.Process.Kill()
}
return nil
},
chk: func(e *GExpect) bool {
if e.cmd.Process == nil {
return false
}
// Sending Signal 0 to a process returns nil if process can take a signal , something else if not.
return e.cmd.Process.Signal(syscall.Signal(0)) == nil
},
}
for _, o := range opts {
o(e)
}
res := make(chan error, 1)
go e.runcmd(res)
// Wait until command started
return e, res, <-res
}
// SpawnSSH starts an interactive SSH session,ties it to a PTY and collects the output. The returned channel sends the
// state of the SSH session after it finishes.
func SpawnSSH(sshClient *ssh.Client, timeout time.Duration, opts ...Option) (*GExpect, <-chan error, error) {
tios := term.Termios{}
tios.Raw()
tios.Wz.WsCol, tios.Wz.WsRow = sshTermWidth, sshTermHeight
return SpawnSSHPTY(sshClient, timeout, tios, opts...)
}
const (
sshTerm = "xterm"
sshTermWidth = 132
sshTermHeight = 43
)
// SpawnSSHPTY starts an interactive SSH session and ties it to a local PTY, optionally requests a remote PTY.
func SpawnSSHPTY(sshClient *ssh.Client, timeout time.Duration, term term.Termios, opts ...Option) (*GExpect, <-chan error, error) {
if sshClient == nil {
return nil, nil, errors.New("*ssh.Client is nil")
}
if timeout < 1 {
timeout = DefaultTimeout
}
// Setup interactive session
session, err := sshClient.NewSession()
if err != nil {
return nil, nil, err
}
e := &GExpect{
rcv: make(chan struct{}),
snd: make(chan string),
chk: func(e *GExpect) bool {
if e.ssh == nil {
return false
}
_, err := e.ssh.SendRequest("dummy", false, nil)
return err == nil
},
cls: func(e *GExpect) error {
if e.ssh != nil {
return e.ssh.Close()
}
return nil
},
ssh: session,
timeout: timeout,
chkDuration: checkDuration,
}
for _, o := range opts {
o(e)
}
if term.Wz.WsCol == 0 {
term.Wz.WsCol = sshTermWidth
}
if term.Wz.WsRow == 0 {
term.Wz.WsRow = sshTermHeight
}
if err := session.RequestPty(sshTerm, int(term.Wz.WsRow), int(term.Wz.WsCol), term.ToSSH()); err != nil {
return nil, nil, err
}
inPipe, err := session.StdinPipe()
if err != nil {
return nil, nil, err
}
outPipe, err := session.StdoutPipe()
if err != nil {
return nil, nil, err
}
errPipe, err := session.StderrPipe()
if err != nil {
return nil, nil, err
}
if err := session.Shell(); err != nil {
return nil, nil, err
}
// Shell started.
errCh := make(chan error, 1)
go e.waitForSession(errCh, session.Wait, inPipe, outPipe, errPipe)
return e, errCh, nil
}
func (e *GExpect) waitForSession(r chan error, wait func() error, sIn io.WriteCloser, sOut io.Reader, sErr io.Reader) {
chDone := make(chan struct{})
var wg sync.WaitGroup
wg.Add(1)
go func() {
defer wg.Done()
for {
select {
case <-chDone:
return
case sstr, ok := <-e.snd:
if !ok {
log.Infof("Send channel closed")
return
}
if _, err := sIn.Write([]byte(sstr)); err != nil || !e.check() {
log.Infof("Write failed: %v", err)
return
}
}
}
}()
rdr := func(out io.Reader) {
defer wg.Done()
buf := make([]byte, bufferSize)
for {
nr, err := out.Read(buf)
if err != nil || !e.check() {
if err == io.EOF {
log.V(2).Infof("read closing down: %v", err)
return
}
return
}
e.mu.Lock()
e.out.Write(buf[:nr])
e.mu.Unlock()
// Inform Expect (if it's currently running) that there's some new data to look through.
select {
case e.rcv <- struct{}{}:
default:
}
}
}
wg.Add(1)
go rdr(sOut)
if sErr != nil {
wg.Add(1)
go rdr(sErr)
}
err := wait()
close(chDone)
wg.Wait()
r <- err
}
// Close closes the Spawned session.
func (e *GExpect) Close() error {
return e.cls(e)
}
// Read implements the reader interface for the out buffer.
func (e *GExpect) Read(p []byte) (nr int, err error) {
e.mu.Lock()
defer e.mu.Unlock()
return e.out.Read(p)
}
// Send sends a string to spawned process.
func (e *GExpect) Send(in string) error {
if !e.check() {
return errors.New("expect: Process not running")
}
e.snd <- in
return nil
}
// runcmd executes the command and Wait for the return value.
func (e *GExpect) runcmd(res chan error) {
if err := e.cmd.Start(); err != nil {
res <- err
return
}
// Moving the go read/write functions here makes sure the command is started before first checking if it's running.
clean := make(chan struct{})
chDone := e.goIO(clean)
// Signal command started
res <- nil
cErr := e.cmd.Wait()
close(chDone)
e.pty.Slave.Close()
// make sure the read/send routines are done before closing the pty.
<-clean
res <- cErr
}
// goIO starts the io handlers.
func (e *GExpect) goIO(clean chan struct{}) (done chan struct{}) {
done = make(chan struct{})
var ptySync sync.WaitGroup
ptySync.Add(2)
go e.read(done, &ptySync)
go e.send(done, &ptySync)
go func() {
ptySync.Wait()
e.pty.Master.Close()
close(clean)
}()
return done
}
// Expect reads spawned processes output looking for input regular expression.
// Timeout set to 0 makes Expect return the current buffer.
// Negative timeout value sets it to Default timeout.
func (e *GExpect) Expect(re *regexp.Regexp, timeout time.Duration) (string, []string, error) {
out, match, _, err := e.ExpectSwitchCase([]Caser{&Case{re, "", nil, 0}}, timeout)
return out, match, err
}
// Options sets the specified options.
func (e *GExpect) Options(opts ...Option) (prev Option) {
for _, o := range opts {
prev = o(e)
}
return prev
}
// read reads from the PTY master and forwards to active Expect function.
func (e *GExpect) read(done chan struct{}, ptySync *sync.WaitGroup) {
defer ptySync.Done()
buf := make([]byte, bufferSize)
for {
nr, err := e.pty.Master.Read(buf)
if err != nil || !e.check() {
if err == io.EOF {
log.V(2).Infof("read closing down: %v", err)
return
}
return
}
// Add to buffer
e.mu.Lock()
e.out.Write(buf[:nr])
e.mu.Unlock()
// Ping Expect function
select {
case e.rcv <- struct{}{}:
default:
}
}
}
// send writes to the PTY master.
func (e *GExpect) send(done chan struct{}, ptySync *sync.WaitGroup) {
defer ptySync.Done()
for {
select {
case <-done:
return
case sstr, ok := <-e.snd:
if !ok {
return
}
if _, err := e.pty.Master.Write([]byte(sstr)); err != nil || !e.check() {
log.Infof("send failed: %v", err)
break
}
}
}
}