1 files changed, 13 insertions, 5 deletions
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