diff --git a/burpui/misc/backend/multi.py b/burpui/misc/backend/multi.py index fd802b57..e56870ff 100644 --- a/burpui/misc/backend/multi.py +++ b/burpui/misc/backend/multi.py @@ -15,18 +15,6 @@ from ...exceptions import BUIserverException from ..._compat import pickle from ...utils import implement -try: - from gevent.coros import RLock - lock = RLock() -except ImportError: - class DummyLock(): - def __enter__(self): - pass - - def __exit__(self, type, value, traceback): - pass - - lock = DummyLock() INTERFACE_METHODS = BUIbackend.__abstractmethods__ PARSER_INTERFACE_METHODS = BUIparser.__abstractmethods__ @@ -286,6 +274,56 @@ class Burp(BUIbackend): return self.servers[agent].get_server_version() +class Gsocket(): + def __init__(self, host, port, ssl=False, timeout=5, notimeout=False): + self.host = host + self.port = port + self.ssl = ssl + self.timeout = timeout + self.notimeout = notimeout + + def conn(self): + if self.ssl: + import ssl + s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + if not self.notimeout: + s.settimeout(self.timeout) + s.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) + ret = ssl.wrap_socket(s, cert_reqs=ssl.CERT_NONE, ssl_version=ssl.PROTOCOL_SSLv23) + ret.connect((self.host, self.port)) + else: + if not self.notimeout: + ret = socket.create_connection((self.host, self.port), timeout=self.timeout) + else: + ret = socket.create_connection((self.host, self.port)) + ret.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) + self.sock = ret + self.connected = True + + def __enter__(self): + self.conn() + return self.sock, self + + def __exit__(self, type, value, traceback): + if self.connected: + self.sock.close() + self.connected = False + + def recvall(self, length=1024): + """Read the answer of the agent""" + buf = b'' + bsize = 1024 + received = 0 + if length < bsize: + bsize = length + while received < length: + newbuf = self.sock.recv(bsize) + if not newbuf: + return None + buf += newbuf + received += len(newbuf) + return buf + class NClient(BUIbackend): """The :class:`burpui.misc.backend.multi.NClient` class provides a consistent backend to interact with ``agents``. @@ -320,7 +358,6 @@ class NClient(BUIbackend): self.port = port self.password = password self.ssl = ssl - self.connected = False self.app = app self.timeout = timeout or 5 @@ -346,139 +383,96 @@ class NClient(BUIbackend): return func return object.__getattribute__(self, name) - def conn(self, notimeout=False): - """Connects to the agent if needed""" - try: - if self.connected: - return - self.sock = self.do_conn(notimeout) - self.connected = True - self.logger.debug('OK, connected to agent %s:%s', self.host, self.port) - except Exception as e: - self.connected = False - self.logger.error('Could not connect to %s:%s => %s', self.host, self.port, str(e)) - - def do_conn(self, notimeout=False): - """Do the actual connection to the agent""" - ret = None - if self.ssl: - import ssl - s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - if not notimeout: - s.settimeout(self.timeout) - s.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) - ret = ssl.wrap_socket(s, cert_reqs=ssl.CERT_NONE, ssl_version=ssl.PROTOCOL_SSLv23) - try: - ret.connect((self.host, self.port)) - except Exception as e: - self.logger.error('ERROR: %s', str(e)) - raise e - else: - if not notimeout: - ret = socket.create_connection((self.host, self.port), timeout=self.timeout) - else: - ret = socket.create_connection((self.host, self.port)) - ret.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) - return ret - def ping(self): """Check if we are connected to the agent""" - self.conn() - res = self.connected + res = False + try: + with Gsocket(self.host, self.port, self.ssl, self.timeout): + res = True + except Exception: + pass return res + def setup(self, sock, gsock, data): + length = len(data) + sock.sendall(struct.pack('!Q', length)) + sock.sendall(data.encode('UTF-8')) + self.logger.debug("Sending: {}".format(data)) + tmp = sock.recv(2).decode('UTF-8') + self.logger.debug("recv: '{}'".format(tmp)) + if 'ER' == tmp: + lengthbuf = sock.recv(8) + length, = struct.unpack('!Q', lengthbuf) + err = gsock.recvall(length).decode('UTF-8') + raise BUIserverException(err) + if 'OK' != tmp: + self.logger.debug('Ooops, unsuccessful!') + return False + self.logger.debug("Data sent successfully") + return True + def do_command(self, data=None, restarted=False): """Send a command to the remote agent""" - with lock: - res = '[]' - err = None - if not data: - raise BUIserverException('Missing data') - try: - data['password'] = self.password - # manage long running operations - if data['func'] in ['restore_files', 'get_file', 'del_file']: - self.connected = False - self.conn(notimeout=True) - else: - self.conn() - if not self.connected: - raise BUIserverException('Failed to connect to agent') + res = '[]' + err = None + notimeout = False + if not data: + raise BUIserverException('Missing data') + data['password'] = self.password + # manage long running operations + if data['func'] in ['restore_files', 'get_file', 'del_file']: + notimeout = True + try: + # don't need a context manager here + if data['func'] == 'get_file': + gsock = Gsocket(self.host, self.port, self.ssl, notimeout=True) + gsock.conn() raw = json.dumps(data) - length = len(raw) - self.sock.sendall(struct.pack('!Q', length)) - self.sock.sendall(raw.encode('UTF-8')) - self.logger.debug("Sending: {}".format(raw)) - tmp = self.sock.recv(2).decode('UTF-8') - self.logger.debug("recv: '{}'".format(tmp)) - if 'ER' == tmp: - lengthbuf = self.sock.recv(8) - length, = struct.unpack('!Q', lengthbuf) - err = self.recvall(length).decode('UTF-8') - raise BUIserverException(err) - if 'OK' != tmp: - self.logger.debug('Ooops, unsuccessful!') + if not self.setup(gsock.sock, gsock, raw): return res - self.logger.debug("Data sent successfully") - if data['func'] == 'get_file': - self.connected = False - return self.sock - lengthbuf = self.sock.recv(8) - length, = struct.unpack('!Q', lengthbuf) - res = self.recvall(length).decode('UTF-8') - except IOError as e: - if not restarted and e.errno == errno.EPIPE: - self.connected = False - self.logger.warning('Broken pipe, restarting the request') - return self.do_command(data, True) - elif e.errno == errno.ECONNRESET: - self.connected = False - self.logger.error('!!! {} !!!\nPlease check your SSL configuration on both sides!'.format(str(e))) - else: - self.logger.error('!!! {} !!!\n{}'.format(str(e), traceback.format_exc())) - raise e - except socket.timeout as e: - if self.app.gunicorn and not restarted: - self.connected = False - self.logger.warning('Socket timed-out, restarting the request') - return self.do_command(data, True) - self.logger.error('!!! {} !!!\n{}'.format(str(e), traceback.format_exc())) - raise e - # catch all - except Exception as e: - self.logger.error('!!! {} !!!\n{}'.format(str(e), traceback.format_exc())) - if data['func'] == 'restore_files': - err = str(e) - elif isinstance(e, BUIserverException): + return gsock.sock + with Gsocket(self.host, self.port, self.ssl, self.timeout, notimeout) as (sock, gsock): + try: + raw = json.dumps(data) + if not self.setup(gsock.sock, gsock, raw): + return res + lengthbuf = sock.recv(8) + length, = struct.unpack('!Q', lengthbuf) + res = gsock.recvall(length).decode('UTF-8') + except IOError as e: + if not restarted and e.errno == errno.EPIPE: + self.logger.warning('Broken pipe, restarting the request') + return self.do_command(data, True) + elif e.errno == errno.ECONNRESET: + self.logger.error('!!! {} !!!\nPlease check your SSL configuration on both sides!'.format(str(e))) + else: + self.logger.error('!!! {} !!!\n{}'.format(str(e), traceback.format_exc())) raise e - else: - raise BUIserverException(str(e)) - finally: - if self.connected: - self.sock.close() - self.connected = False + except socket.timeout as e: + if self.app.gunicorn and not restarted: + self.logger.warning('Socket timed-out, restarting the request') + return self.do_command(data, True) + self.logger.error('!!! {} !!!\n{}'.format(str(e), traceback.format_exc())) + raise e + # catch all + except Exception as e: + self.logger.error('!!! {} !!!\n{}'.format(str(e), traceback.format_exc())) + if data['func'] == 'restore_files': + err = str(e) + elif isinstance(e, BUIserverException): + raise e + else: + raise BUIserverException(str(e)) + except Exception as e: + self.logger.error('!!! {} !!!\n{}'.format(str(e), traceback.format_exc())) + raise BUIserverException(str(e)) - if data['func'] == 'restore_files': - if err: - res = None - return res, err + if data['func'] == 'restore_files': + if err: + res = None + return res, err - return res - - def recvall(self, length=1024): - """Read the answer of the agent""" - buf = b'' - bsize = 1024 - received = 0 - if length < bsize: - bsize = length - while received < length: - newbuf = self.sock.recv(bsize) - if not newbuf: - return None - buf += newbuf - received += len(newbuf) - return buf + return res """ Utilities functions