lib/expectr.rb in expectr-0.9.1 vs lib/expectr.rb in expectr-1.0.0

@@ -1,290 +1,287 @@
-# = expectr.rb
-#
-#  Copyright (c) Chris Wuest <[e-mail]>
-#  Expectr is freely distributable under the terms of an MIT-style license.
-#  See COPYING or http://www.opensource.org/licenses/mit-license.php.
-
-begin
-	require 'pty'
-rescue LoadError
-	require 'popen4'
-end
-
+require 'pty'
 require 'timeout'
 require 'thread'
 
-# Fixes specifically for Ruby 1.8
-if RUBY_VERSION =~ /^1.8/
-	# Enforcing encoding is not needed in 1.8 (probably.)  So, we'll define
-	# String#encode! to do nothing, for interoperability.
-	class String #:nodoc:
-		def encode!(encoding)
-		end
-	end
+require 'expectr/error'
 
-	# In Ruby 1.8, we want to ignore SIGCHLD.  This is for two reasons:
-	# * SIGCHLD will be sent (and cause exceptions) for every Expectr object
-	#   created
-	# * As John Carter documented in his RExpect library, calls to files which
-	#   do not exist can cause odd and unexpected behavior.
-	trap 'CHLD', Proc.new { nil }
-end
-
-# == Description
-# Expectr is an implementation of the Expect library in ruby (see
-# http://expect.nist.gov).
+# Public: Expectr is an API to the functionality of Expect (see
+# http://expect.nist.gov) implemented in ruby.
 #
 # Expectr contrasts with Ruby's built-in Expect class by avoiding tying in
-# with IO and instead creating a new object entirely to allow for more
-# fine-grained control over the execution and display of the program being
+# with the IO class, instead creating a new object entirely to allow for more
+# grainular control over the execution and display of the program being
 # run.
 #
-# == Examples
-# === Simple task automation
+# Examples
 #
-# Connect via telnet to remote.example.com, run my_command, and return the
-# output
+#   # SSH Login to another machine
+#   exp = Expectr.new('ssh [e-mail]')
+#   exp.expect("Password:")
+#   exp.send('password')
+#   exp.interact!(blocking: true)
 #
-#   exp = Expectr.new "telnet remote.example.com"
-#   exp.expect "username:"
-#   exp.send "example\r"
-#   exp.expect "password:"
-#   exp.send "my_password\r"
-#   exp.expect "%"
-#   exp.send "my_command\r"
-#   exp.expect "%"
-#   exp.send "logout"
-#
-#   output = exp.discard
-#
-# === Interactive control
-# Silently connect via ssh to remote.example.com, log in automatically, then
-# relinquish control to the user.  Expect slow networking, so increase
-# timeout.
-#
-#   exp = Expectr.new "ssh remote.example.com", :timeout=>45, :flush_buffer=>false
-#
-#   match = exp.expect /password|yes\/no/
-#   case match.to_s
-#     when /password/
-#       exp.send "my_password\r"
-#     when /yes\/no/
-#       exp.send "yes\r"
-#       exp.expect /password/
-#       exp.send "my_password\r"
-#     else
-#       puts "Cannot connect to remote.example.com!"
-#       die
+#   # See if a web server is running on the local host, react accordingly
+#   exp = Expectr.new('netstat -ntl|grep ":80 " && echo "WEB"', timeout: 1)
+#   if exp.expeect("WEB")
+#     # Do stuff if we see 'WEB' in the output
+#   else
+#     # Do other stuff
 #   end
-#
-#   exp.expect "$"
-#   exp.interact
-#
 class Expectr
-	# Amount of time in seconds a call to +expect+ may last (default 30)
-	attr_accessor :timeout
-	# Size of buffer in bytes to attempt to read in at once (default 8 KiB)
-	attr_accessor :buffer_size
-	# Whether to flush program output to STDOUT (default true)
-	attr_accessor :flush_buffer
-	# PID of running process
-	attr_reader :pid
-	# Active buffer to match against
-	attr_reader :buffer
-	# Buffer passed since last +expect+ match
-	attr_reader :discard
+  # Public: Gets/sets the number of seconds a call to Expectr#expect may last
+  attr_accessor :timeout
+  # Public: Gets/sets the number of bytes to use for the internal buffer
+  attr_accessor :buffer_size
+  # Public: Gets/sets whether to constrain the buffer to the buffer size
+  attr_accessor :constrain
+  # Public: Gets/sets whether to flush program output to STDOUT
+  attr_accessor :flush_buffer
+  # Public: Returns the PID of the running process
+  attr_reader :pid
+  # Public: Returns the active buffer to match against
+  attr_reader :buffer
+  # Public: Returns the buffer discarded by the latest call to Expectr#expect
+  attr_reader :discard
 
-	# 
-	# === Synopsis
-	#
-	#   Expectr.new(cmd, args)
-	#
-	# === Arguments
-	# +cmd+::
-	#   Command to be executed (String or File)
-	# +args+::
-	#   Hash of modifiers for Expectr.  Meaningful values are:
-	# * :buffer_size::
-	#   Amount of data to read at a time.  Default 8 KiB
-	# * :flush_buffer::
-	#   Flush buffer to STDOUT during execution?  Default true
-	# * :timeout::
-	#   Timeout in seconds for each +expect+ call.  Default 30
-	#
-	# === Description
-	#
-	# Spawn +cmd+ and attach to STDIN and STDOUT for new process.  Fall back
-	# to using Open4 if PTY is not present (this is the case on Windows
-	# implementations of ruby.
+  # Public: Initialize a new Expectr object.
+  # Spawns a sub-process and attaches to STDIN and STDOUT for the new process.
   #
-	def initialize(cmd, args={})
-		raise ArgumentError, "String or File expected, was given #{cmd.class}" unless cmd.kind_of? String or cmd.kind_of? File
-		cmd = cmd.path if cmd.kind_of? File
+  # cmd  - A String or File referencing the application to launch
+  # args - A Hash used to specify options for the new object (default: {}):
+  #        :timeout      - Number of seconds that a call to Expectr#expect has
+  #                        to complete (default: 30)
+  #        :flush_buffer - Whether to flush output of the process to the
+  #                        console (default: true)
+  #        :buffer_size  - Number of bytes to attempt to read from sub-process
+  #                        at a time.  If :constrain is true, this will be the
+  #                        maximum size of the internal buffer as well.
+  #                        (default: 8192)
+  #        :constrain    - Whether to constrain the internal buffer from the
+  #                        sub-process to :buffer_size (default: false)
+  def initialize(cmd, args={})
+    unless cmd.kind_of? String or cmd.kind_of? File
+     raise ArgumentError, "String or File expected"
+    end
 
-		args[0] = {} unless args[0]
-		@buffer = String.new
-		@discard = String.new
-		@timeout = args[:timeout] || 30
-		@flush_buffer = args[:flush_buffer].nil? ? true : args[:flush_buffer]
-		@buffer_size = args[:buffer_size] || 8192
-		@out_mutex = Mutex.new
-		@out_update = false
+    cmd = cmd.path if cmd.kind_of? File
 
-		[@buffer, @discard].each {|x| x.encode! "UTF-8" }
+    @buffer = ''.encode("UTF-8")
+    @discard = ''.encode("UTF-8")
 
-		if defined? PTY
-			@stdout,@stdin,@pid = PTY.spawn cmd
-		else
-			cmd << " 2>&1" if cmd[/2\s*>/].nil?
-			@pid, @stdin, @stdout, stderr = Open4::popen4 cmd
-		end
+    @timeout = args[:timeout] || 30
+    @flush_buffer = args[:flush_buffer].nil? ? true : args[:flush_buffer]
+    @buffer_size = args[:buffer_size] || 8192
+    @constrain = args[:constrain] || false
 
-		Thread.new do
-			while @pid > 0
-				unless select([@stdout], nil, nil, @timeout).nil?
-					buf = ''
+    @out_mutex = Mutex.new
+    @out_update = false
+    @interact = false
 
-					begin
-						@stdout.sysread(@buffer_size, buf)
-					rescue Errno::EIO #Application went away.
-						@pid = 0
-						break
-					end
+    @stdout,@stdin,@pid = PTY.spawn(cmd)
 
-					buf.encode! "UTF-8"
-					print_buffer buf
+    Thread.new do
+      while @pid > 0
+        unless select([@stdout], nil, nil, @timeout).nil?
+          buf = ''.encode("UTF-8")
 
-					@out_mutex.synchronize do
-						@buffer << buf
-						@out_update = true
-					end
-				end
-			end
-		end
+          begin
+            @stdout.sysread(@buffer_size, buf)
+          rescue Errno::EIO #Application went away.
+            @pid = 0
+            break
+          end
 
-		Thread.new do
-			Process.wait @pid
-			@pid = 0
-		end
-	end
+          print_buffer(buf)
 
-	# 
-	# Clear output buffer
-	#
-	def clear_buffer
-		@out_mutex.synchronize do
-			@buffer = ''
-			@out_update = false
-		end
-	end
+          @out_mutex.synchronize do
+            @buffer << buf
+            if @buffer.length > @buffer_size && @constrain
+              @buffer = @buffer[-@buffer_size..-1]
+            end
+            @out_update = true
+          end
+        end
+      end
+    end
 
-	# 
-	# === Synopsis
-	#
-	#   Expectr#interact
-	# 
-	# === Description
-	#
-	# Relinquish control of the running process to the controlling terminal,
-	# acting simply as a pass-through for the life of the process.
-	#
-	# Interrupts should be caught and sent to the application.
-	#
-	def interact
-		oldtrap = trap 'INT' do
-				send "\C-c"
-		end
+    Thread.new do
+      Process.wait @pid
+      @pid = 0
+    end
+  end
 
-		@flush_buffer = true
-		old_tty = `stty -g`
-		`stty -icanon min 1 time 0 -echo`
-		
-		in_thread = Thread.new do
-			input = ''
-			while @pid > 0
-				if select([STDIN], nil, nil, 1)
-					send STDIN.getc.chr
-				end
-			end
-		end
+  # Public: Relinquish control of the running process to the controlling
+  # terminal, acting as a pass-through for the life of the process.  SIGINT
+  # will be caught and sent to the application as "\C-c".
+  #
+  # args - A Hash used to specify options to be used for interaction (default:
+  #        {}):
+  #        :flush_buffer - explicitly set @flush_buffer to the value specified
+  #        :blocking     - Whether to block on this call or allow code
+  #                        execution to continue (default: false)
+  #
+  # Returns the interaction Thread
+  def interact!(args = {})
+    raise ProcessError if @interact
 
-		in_thread.join
-		trap 'INT', oldtrap
-		`stty #{old_tty}`
-		return nil
-	end
-	alias :interact! :interact
+    blocking = args[:blocking] || false
+    @flush_buffer = args[:flush_buffer].nil? ? true : args[:flush_buffer]
+    @interact = true
 
-	#
-	# Send +str+ to application
-	#
-	def send(str)
-		begin
-			@stdin.syswrite str
-		rescue Errno::EIO #Application went away.
-			@pid = 0
-		end
-		raise ArgumentError unless @pid > 0
-	end
+    # Save our old tty settings and set up our new environment
+    old_tty = `stty -g`
+    `stty -icanon min 1 time 0 -echo`
 
-	# 
-	# === Synopsis
-	#
-	#   Expectr#expect /regexp/, recoverable=false
-	#   Expectr#expect "String", recoverable=true
-	# 
-	# === Arguments
-	#
-	# +pattern+::
-	#   String or regexp to match against
-	# +recoverable+::
-	#   Determines if execution can continue after a timeout
-	#
-	# === Description
-	#
-	# Wait +timeout+ seconds to match +pattern+ in +buffer+.  If timeout is
-	# reached, raise an error unless +recoverable+ is true.
-	#
-	def expect(pattern, recoverable = false)
-		match = nil
+    # SIGINT should be set along to the program
+    oldtrap = trap 'INT' do
+      send "\C-c"
+    end
+    
+    interact = Thread.new do
+      input = ''.encode("UTF-8")
+      while @pid > 0 && @interact
+        if select([STDIN], nil, nil, 1)
+          c = STDIN.getc.chr
+          send c unless c.nil?
+        end
+      end
 
-		case pattern
-			when String
-				pattern = Regexp.new(Regexp.quote(pattern))
-			when Regexp
-			else raise TypeError, "Pattern class should be String or Regexp, passed: #{pattern.class}"
-		end
+      trap 'INT', oldtrap
+      `stty #{old_tty}`
+      @interact = false
+    end
 
-		begin
-			Timeout::timeout(@timeout) do
-				while match.nil?
-					if @out_update
-						@out_mutex.synchronize do
-							match = pattern.match @buffer
-							@out_update = false
-						end
-					end
-					sleep 0.1
-				end
-			end
+    blocking ? interact.join : interact
+  end
 
-			@out_mutex.synchronize do
-				@discard = @buffer[0..match.begin(0)-1]
-				@buffer = @buffer[match.end(0)..-1]
-				@out_update = true
-			end
-		rescue Timeout::Error => details
-			raise details unless recoverable
-		end
+  # Public: Report whether or not current Expectr object is in interact mode
+  #
+  # Returns true or false
+  def interact?
+    @interact
+  end
 
-		return match
-	end
+  # Public: Cause the current Expectr object to drop out of interact mode
+  #
+  # Returns nothing.
+  def leave!
+    @interact=false
+  end
 
-	#
-	# Print buffer to STDOUT only if +flush_buffer+ is true
-	#
-	def print_buffer(buf)
-		print buf if @flush_buffer
-		STDOUT.flush
-	end
+  # Public: Kill the running process, raise ProcessError if the pid isn't > 1
+  #
+  # signal - Symbol, String, or Fixnum representing the signal to send to the
+  #          running process. (default: :TERM)
+  #
+  # Returns true if the process was successfully killed, false otherwise
+  def kill!(signal=:TERM)
+    raise ProcessError unless @pid > 0
+    (Process::kill(signal.to_sym, @pid) == 1)
+  end
+
+  # Public: Send input to the active process
+  #
+  # str - String to be sent to the active process
+  #
+  # Returns nothing.
+  # Raises Expectr::ProcessError if the process isn't running
+  def send(str)
+    begin
+      @stdin.syswrite str
+    rescue Errno::EIO #Application went away.
+      @pid = 0
+    end
+    raise Expectr::ProcessError unless @pid > 0
+  end
+
+  # Public: Wraps Expectr#send, appending a newline to the end of the string
+  #
+  # str - String to be sent to the active process (default: '')
+  #
+  # Returns nothing.
+  def puts(str = '')
+    send str + "\n"
+  end
+
+  # Public: Begin a countdown and search for a given String or Regexp in the
+  # output buffer.
+  #
+  # pattern     - String or Regexp representing what we want to find
+  # recoverable - Denotes whether failing to match the pattern should cause the
+  #               method to raise an exception (default: false)
+  #
+  # Examples
+  #
+  #   exp.expect("this should exist")
+  #   # => MatchData
+  #
+  #   exp.expect("this should exist") do
+  #     # ...
+  #   end
+  #
+  #   exp.expect(/not there/)
+  #   # Raises Timeout::Error
+  #
+  #   exp.expect(/not there/, true)
+  #   # => nil
+  # 
+  # Returns a MatchData object once a match is found if no block is given
+  # Yields the MatchData object representing the match
+  # Raises TypeError if something other than a String or Regexp is given
+  # Raises Timeout::Error if a match isn't found in time, unless recoverable
+  def expect(pattern, recoverable = false)
+    match = nil
+
+    case pattern
+    when String
+      pattern = Regexp.new(Regexp.quote(pattern))
+    when Regexp
+    else
+      raise TypeError, "Pattern class should be String or Regexp"
+    end
+
+    begin
+      Timeout::timeout(@timeout) do
+        while match.nil?
+          if @out_update
+            @out_mutex.synchronize do
+              match = pattern.match @buffer
+              @out_update = false
+            end
+          end
+          sleep 0.1
+        end
+      end
+
+      @out_mutex.synchronize do
+        @discard = @buffer[0..match.begin(0)-1]
+        @buffer = @buffer[match.end(0)..-1]
+        @out_update = true
+      end
+    rescue Timeout::Error => details
+      raise details unless recoverable
+    end
+
+    block_given? ? yield(match) : match
+  end
+
+  # Public: Clear output buffer
+  #
+  # Returns nothing.
+  def clear_buffer!
+    @out_mutex.synchronize do
+      @buffer = ''.encode("UTF-8")
+      @out_update = false
+    end
+  end
+
+  # Internal: Print buffer to STDOUT if @flush_buffer is true
+  #
+  # buf - String to be printed to STDOUT
+  #
+  # Returns nothing.
+  def print_buffer(buf)
+    print buf if @flush_buffer
+    STDOUT.flush unless STDOUT.sync
+  end
 end