Add Model#expect, support block at post/pre config

Now we can deal with pager and additional PW prompts, such as 'enable'

Examples in IOS model how to use.

The Telnet implementation is particularly fugly, I just need one line in
'waitfor' to handle pager while waiting for prompt, but couldn't figure
out clean way to do it, so needed to rewrit whole Telnet#waitfor just to
add that line.

6 files changed, 173 insertions, 21 deletions

diff --git a/lib/oxidized/config/bootstrap.rb b/lib/oxidized/config/bootstrap.rb
index 6d865b5..bae2b70 100644
--- a/lib/oxidized/config/bootstrap.rb
+++ b/lib/oxidized/config/bootstrap.rb
@@ -4,13 +4,16 @@ module Oxidized
   CFG.username = 'username'
   CFG.password = 'password'
   CFG.model    = 'junos'
-  CFG.interval = 60
+  CFG.interval = 3600
   CFG.log      = File.join Config::Root, 'log'
   CFG.debug    = false
   CFG.threads  = 30
   CFG.timeout  = 5
-  CFG.prompt   = /^([\w\.\-@]{3,30}[#>]\s?)$/ 
+  CFG.prompt   = /^([\w.@-]+[#>]\s?)$/
   CFG.rest     = 8888
+  CFG.vars     = {
+    :enable  => 'enablePW',
+  }
   CFG.input    = {
     :default => 'ssh, telnet',
   }
diff --git a/lib/oxidized/input/cli.rb b/lib/oxidized/input/cli.rb
index 1d58e85..8daddac 100644
--- a/lib/oxidized/input/cli.rb
+++ b/lib/oxidized/input/cli.rb
@@ -8,18 +8,26 @@ module Oxidized
       end
 
       def get
-        @post_login.each { |command| cmd command }
+        @post_login.each { |command, block| block ? block.call : (cmd command) }
         d = @node.model.get
         disconnect
         d
       end
+
+      def disconnect_cli
+        @pre_logout.each { |command, block| block ? block.call : (cmd command) }
+      end
   
-      def post_login _post_login
-        @post_login << _post_login unless @exec
+      def post_login _post_login=nil, &block
+        unless @exec
+          @post_login << [_post_login, block]
+        end
       end
   
-      def pre_logout _pre_logout
-        @pre_logout << _pre_logout unless @exec
+      def pre_logout _pre_logout=nil, &block
+        unless @exec
+          @pre_logout <<  [_pre_logout, block]
+        end
       end
     end
   end
diff --git a/lib/oxidized/input/ssh.rb b/lib/oxidized/input/ssh.rb
index 3471eea..3d14a57 100644
--- a/lib/oxidized/input/ssh.rb
+++ b/lib/oxidized/input/ssh.rb
@@ -15,7 +15,7 @@ module Oxidized
       rescue Timeout::Error, Net::SSH::Disconnect, Errno::ECONNREFUSED
         return false
       end
-      @ses = open_shell @ssh unless @exec
+      open_shell @ssh unless @exec
       not @ssh.closed?
     end
 
@@ -28,11 +28,15 @@ module Oxidized
       end
     end
 
+    def send data
+      @ses.send_data data
+    end
+
     private
 
     def disconnect
       begin
-        @pre_logout.each { |command| cmd command }
+        disconnect_cli
         @ssh.loop
         @ssh.close if not @ssh.closed?
       rescue Net::SSH::Disconnect
@@ -40,9 +44,10 @@ module Oxidized
     end
 
     def open_shell ssh
-      ses = ssh.open_channel do |ch|
+      @ses = ssh.open_channel do |ch|
         ch.on_data do |ch, data|
           @output << data
+          @output = @node.model.expects @output
         end
         ch.request_pty do |ch, success|
           raise NoShell, "Can't get PTY" unless success
@@ -52,7 +57,6 @@ module Oxidized
         end
       end
       expect @node.prompt
-      ses
     end
 
     def exec state=nil
diff --git a/lib/oxidized/input/telnet.rb b/lib/oxidized/input/telnet.rb
index 0ec08f9..c1afde2 100644
--- a/lib/oxidized/input/telnet.rb
+++ b/lib/oxidized/input/telnet.rb
@@ -10,7 +10,8 @@ module Oxidized
       @timeout = CFG.timeout
       @node.model.cfg['telnet'].each { |cb| instance_exec &cb }
       begin
-        @telnet  = Net::Telnet.new 'Host' => @node.ip, 'Waittime' => @timeout
+        @telnet  = Net::Telnet.new 'Host' => @node.ip, 'Waittime' => @timeout,
+                                   'Model' => @node.model
         expect username
         @telnet.puts @node.auth[:username]
         expect password
@@ -32,6 +33,10 @@ module Oxidized
       end
     end
 
+    def send data
+      @telnet.write data
+    end
+
     private
 
     def expect re
@@ -40,7 +45,7 @@ module Oxidized
 
     def disconnect
       begin
-        @pre_logout.each { |command| cmd(command, nil) }
+        disconnect_cli
         @telnet.close
       rescue Errno::ECONNRESET
       end
@@ -56,3 +61,89 @@ module Oxidized
 
   end
 end
+
+
+class Net::Telnet
+  ## FIXME: we just need 'line = model.expects line' to handle pager
+  ## how to do this, without redefining the whole damn thing
+  def waitfor(options) # :yield: recvdata
+    time_out = @options["Timeout"]
+    waittime = @options["Waittime"]
+    fail_eof = @options["FailEOF"]
+    model    = @options["Model"]
+
+    if options.kind_of?(Hash)
+      prompt   = if options.has_key?("Match")
+                   options["Match"]
+                 elsif options.has_key?("Prompt")
+                   options["Prompt"]
+                 elsif options.has_key?("String")
+                   Regexp.new( Regexp.quote(options["String"]) )
+                 end
+      time_out = options["Timeout"]  if options.has_key?("Timeout")
+      waittime = options["Waittime"] if options.has_key?("Waittime")
+      fail_eof = options["FailEOF"]  if options.has_key?("FailEOF")
+    else
+      prompt = options
+    end
+
+    if time_out == false
+      time_out = nil
+    end
+
+    line = ''
+    buf = ''
+    rest = ''
+    until(prompt === line and not IO::select([@sock], nil, nil, waittime))
+      unless IO::select([@sock], nil, nil, time_out)
+        raise Net::ReadTimeout, "timed out while waiting for more data"
+      end
+      begin
+        c = @sock.readpartial(1024 * 1024)
+        @dumplog.log_dump('<', c) if @options.has_key?("Dump_log")
+        if @options["Telnetmode"]
+          c = rest + c
+          if Integer(c.rindex(/#{IAC}#{SE}/no) || 0) <
+             Integer(c.rindex(/#{IAC}#{SB}/no) || 0)
+            buf = preprocess(c[0 ... c.rindex(/#{IAC}#{SB}/no)])
+            rest = c[c.rindex(/#{IAC}#{SB}/no) .. -1]
+          elsif pt = c.rindex(/#{IAC}[^#{IAC}#{AO}#{AYT}#{DM}#{IP}#{NOP}]?\z/no) ||
+                     c.rindex(/\r\z/no)
+            buf = preprocess(c[0 ... pt])
+            rest = c[pt .. -1]
+          else
+            buf = preprocess(c)
+            rest = ''
+          end
+       else
+         # Not Telnetmode.
+         #
+         # We cannot use preprocess() on this data, because that
+         # method makes some Telnetmode-specific assumptions.
+         buf = rest + c
+         rest = ''
+         unless @options["Binmode"]
+           if pt = buf.rindex(/\r\z/no)
+             buf = buf[0 ... pt]
+             rest = buf[pt .. -1]
+           end
+           buf.gsub!(/#{EOL}/no, "\n")
+         end
+        end
+        @log.print(buf) if @options.has_key?("Output_log")
+        line += buf
+        line = model.expects line
+        line = yield line if block_given?
+        yield buf if block_given?
+      rescue EOFError # End of file reached
+        raise if fail_eof
+        if line == ''
+          line = nil
+          yield nil if block_given?
+        end
+        break
+      end
+    end
+    line
+  end
+end
diff --git a/lib/oxidized/model/ios.rb b/lib/oxidized/model/ios.rb
index c3027e2..7da5ed2 100644
--- a/lib/oxidized/model/ios.rb
+++ b/lib/oxidized/model/ios.rb
@@ -2,7 +2,22 @@ class IOS < Oxidized::Model
 
   comment  '! '
 
+  # example how to handle pager
+  #expect /^\s--More--\s+.*$/ do |data, re|
+  #  send ' '
+  #  data.sub re, ''
+  #end
+
+  # non-preferred way to handle additional PW prompt
+  #expect /^[\w.]+>$/ do |data|
+  #  send "enable\n"
+  #  send CFG.passwords[:enable] + "\n"
+  #  data
+  #end
+
   cmd :all do |cfg|
+    #cfg.gsub! /\cH+\s{8}/, ''         # example how to handle pager
+    #cfg.gsub! /\cH+/, ''              # example how to handle pager
     cfg.each_line.to_a[1..-3].join
   end
 
@@ -24,6 +39,11 @@ class IOS < Oxidized::Model
   cfg :telnet, :ssh do
     post_login 'terminal length 0'
     post_login 'terminal width 0'
+    # preferred way to handle additional passwords
+    #post_login do
+    #  send "enable\n"
+    #  send CFG.passwords[:enable] + "\n"
+    #end
     pre_logout 'exit'
   end
 
diff --git a/lib/oxidized/model/model.rb b/lib/oxidized/model/model.rb
index 8f7b14f..2a6cbc4 100644
--- a/lib/oxidized/model/model.rb
+++ b/lib/oxidized/model/model.rb
@@ -2,22 +2,24 @@ module Oxidized
   class Model
     class << self
       def inherited klass
-        klass.instance_variable_set '@cmd', Hash.new { |h,k| h[k] = [] }
-        klass.instance_variable_set '@cfg', Hash.new { |h,k| h[k] = [] }
+        klass.instance_variable_set '@cmd',    Hash.new { |h,k| h[k] = [] }
+        klass.instance_variable_set '@cfg',    Hash.new { |h,k| h[k] = [] }
+        klass.instance_variable_set '@expect', []
+        klass.const_set :CFG, CFG
         Oxidized.mgr.loader = { :class => klass }
       end
       def comment _comment='# '
         return @comment if @comment
         @comment = block_given? ? yield : _comment
       end
+      def prompt _prompt=nil
+        @prompt or @prompt = _prompt
+      end
       def cfg *methods, &block
         [methods].flatten.each do |method|
           @cfg[method.to_s] << block
         end
       end
-      def prompt _prompt=nil
-        @prompt or @prompt = _prompt
-      end
       def cfgs
         @cfg
       end
@@ -31,8 +33,11 @@ module Oxidized
       def cmds
         @cmd
       end
-      def post_login &block
-        @post_login or @post_login = block
+      def expect re, &block
+        @expect << [re, block]
+      end
+      def expects
+        @expect
       end
     end
 
@@ -48,7 +53,15 @@ module Oxidized
       out
     end
 
-    def  cfg
+    def send data
+      @input.send data
+    end
+
+    def expect re, &block
+      self.class.expect re, &block
+    end
+
+    def cfg
       self.class.cfgs
     end
 
@@ -56,6 +69,19 @@ module Oxidized
       self.class.prompt
     end
 
+    def expects data
+      self.class.expects.each do |re, cb|
+        if data.match re
+          if cb.arity == 2
+            data = instance_exec [data, re], &cb
+          else
+            data = instance_exec data, &cb
+          end
+        end
+      end
+      data
+    end
+
     def get
       data = ''
       self.class.cmds[:cmd].each do |command, block|