mirror of
https://github.com/ziirish/burp-ui.git
synced 2026-05-21 06:45:24 -06:00
use external socket layer due to race conditions with gevent
This commit is contained in:
ziirish
parent
d51d03c9db
commit
341255815b
1 changed files with 129 additions and 135 deletions
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue