Merge pull request #5 from rapid7/PluginFixes

Additional Fixes to underlying metasploit library

src/metasploit/msfrpc.py

@@ -3,6 +3,7 @@
 from http.client import HTTPConnection, HTTPSConnection
 import ssl
 from numbers import Number
+from .utils import convert_bytes_to_string
 
 from msgpack import packb, unpackb
 
@@ -228,14 +229,16 @@ def call(self, method, *args):
             self.client.request('POST', self.uri, packb(l), self._headers)
             r = self.client.getresponse()
             if r.status == 200:
-                return unpackb(r.read(), raw=False)
+                byte = r.read()
+                return convert_bytes_to_string(unpackb(byte, raw=False))
             raise MsfRpcError('An unknown error has occurred while logging in.')
         elif self.authenticated:
             l.insert(1, self.sessionid)
             self.client.request('POST', self.uri, packb(l), self._headers)
             r = self.client.getresponse()
             if r.status == 200:
-                result = unpackb(r.read())
+                byte = r.read()
+                result = convert_bytes_to_string(unpackb(byte, raw=False))
                 if 'error' in result:
                     raise MsfRpcError(result['error_message'])
                 return result
@@ -312,21 +315,9 @@ def login(self, username, password):
         if self.sessionid is None:
             r = self.call(MsfRpcMethod.AuthLogin, username, password)
             # in case r is actually encoded in bytes, this is a safe way to turn it back to string for the try/catch
-            str_data = {}
-            for key, val in r.items():
-                if isinstance(key, bytes):
-                    key_temp = key.decode()
-                else:
-                    key_temp = key
-                if isinstance(val, bytes):
-                    val_temp = val.decode()
-                else:
-                    val_temp = val
-                str_data[key_temp] = val_temp
-
             try:
-                if str_data['result'] == 'success':
-                    self.sessionid = str_data['token']
+                if r['result'] == 'success':
+                    self.sessionid = r['token']
             except KeyError:
                 raise MsfRpcError('Login failed.')
         else:
@@ -341,6 +332,12 @@ def logout(self):
         """
         self.call(MsfRpcMethod.AuthLogout, self.sessionid)
 
+    def close(self):
+        """
+        Close the current http client connection for memory management purposes
+        """
+        self.client.close()
+
 
 class MsfTable(object):
 

src/metasploit/utils.py

@@ -29,4 +29,44 @@ def parseargs():
         print('[-] Error: a password must be specified (-P)\n')
         p.print_help()
         exit(-1)
-    return o
+    return o
+
+
+def convert_bytes_to_string(bytes_dict: dict) -> dict:
+    """
+    :param bytes_dict: dictionary, where inner objects are strings, bytes, or collections of those types
+    :return: original dictionary with any bytes values converted to strings
+    """
+    str_data = {}
+    if isinstance(bytes_dict, dict):
+        for key, val in bytes_dict.items():
+            if isinstance(key, bytes):
+                key_temp = key.decode()
+            else:
+                key_temp = key
+            if isinstance(val, list) or isinstance(val, tuple):
+                val_temp = convert_val(val)
+            elif isinstance(val, dict):
+                val_temp = convert_bytes_to_string(val)
+            elif isinstance(val, bytes):
+                val_temp = val.decode()
+            else:
+                val_temp = val
+            str_data[key_temp] = val_temp
+    return str_data
+
+
+def convert_val(bytes_obj: iter):
+    """
+    :param bytes_obj: some object containing either string or byte values, can be collection or just bytes
+    :return: iterable object (tuple or list) containing original strings and converted bytes to strings
+    """
+    for index, item in enumerate(bytes_obj):
+        if isinstance(item, bytes):
+            bytes_obj[index] = item.decode('utf-8')
+        elif isinstance(item, dict):
+            item = convert_bytes_to_string(item)
+            bytes_obj[index] = item
+        elif isinstance(item, tuple) or isinstance(item, list):
+            bytes_obj[index] = convert_val(item)
+    return bytes_obj