Merge branch 'wchen-r7-fix_7959' into upstream-master

lib/msf/core/rpc/v10/client.rb

@@ -28,6 +28,9 @@ class Client
   # @option info [String] :token A token used by the client.
   # @return [void]
   def initialize(info={})
+    @user = nil
+    @pass = nil
+
     self.info = {
       :host => '127.0.0.1',
       :port => 3790,
@@ -41,14 +44,16 @@ class Client
   end
 
 
-  # Logs in by calling the 'auth.login' API. The authentication token will expire 5 minutes
-  # after the last request was made.
+  # Logs in by calling the 'auth.login' API. The authentication token will expire after 5
+  # minutes, but will automatically be rewnewed when you make a new RPC request.
   #
   # @param [String] user Username.
   # @param [String] pass Password.
   # @raise RuntimeError Indicating a failed authentication.
   # @return [TrueClass] Indicating a successful login.
   def login(user,pass)
+    @user = user
+    @pass = pass
     res = self.call("auth.login", user, pass)
     unless (res && res['result'] == "success")
       raise RuntimeError, "authentication failed"
@@ -58,6 +63,14 @@ class Client
   end
 
 
+  # Attempts to login again with the last known user name and password.
+  #
+  # @return [TrueClass] Indicating a successful login.
+  def re_login
+    login(@user, @pass)
+  end
+
+
   # Calls an API.
   #
   # @param [String] meth The RPC API to call.
@@ -75,6 +88,10 @@ class Client
   #  # {"version"=>"4.11.0-dev", "ruby"=>"2.1.5 x86_64-darwin14.0 2014-11-13", "api"=>"1.0"}
   #  rpc.call('core.version')
   def call(meth, *args)
+    if meth == 'auth.logout'
+      do_logout_cleanup
+    end
+
     unless meth == "auth.login"
       unless self.token
         raise RuntimeError, "client not authenticated"
@@ -84,6 +101,36 @@ class Client
 
     args.unshift(meth)
 
+    begin
+      send_rpc_request(args)
+    rescue Msf::RPC::ServerException => e
+      if e.message =~ /Invalid Authentication Token/i && meth != 'auth.login' && @user && @pass
+        re_login
+        args[1] = self.token
+        retry
+      else
+        raise e
+      end
+    ensure
+      @cli.close if @cli
+    end
+
+  end
+
+
+  # Closes the client.
+  #
+  # @return [void]
+  def close
+    if @cli && @cli.conn?
+      @cli.close
+    end
+    @cli = nil
+  end
+
+  private
+
+  def send_rpc_request(args)
     unless @cli
       @cli = Rex::Proto::Http::Client.new(info[:host], info[:port], info[:context], info[:ssl], info[:ssl_version])
       @cli.set_config(
@@ -101,7 +148,6 @@ class Client
     )
 
     res = @cli.send_recv(req)
-    @cli.close
 
     if res && [200, 401, 403, 500].include?(res.code)
       resp = MessagePack.unpack(res.body)
@@ -118,18 +164,11 @@ class Client
     end
   end
 
-
-  # Closes the client.
-  #
-  # @return [void]
-  def close
-    if @cli && @cli.conn?
-      @cli.close
-    end
-    @cli = nil
+  def do_logout_cleanup
+    @user = nil
+    @pass = nil
   end
 
 end
 end
 end
-