diff --git a/src/metasploit/msfrpc.py b/src/metasploit/msfrpc.py index d9b5874..bc39fff 100644 --- a/src/metasploit/msfrpc.py +++ b/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): diff --git a/src/metasploit/utils.py b/src/metasploit/utils.py index d241241..907a953 100644 --- a/src/metasploit/utils.py +++ b/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 \ No newline at end of file + 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