class SocketFactory:
"""
Abstracts out the socket creation, specifically the issues with
transforming ports into zmq addresses.
Also sets up some default protocol and address handeling.
For example, if everything is being run locally on one machine, and the
default is to bind all addresses `*` and the transport protocol
`tcp` means tcp communication
"""
def __init__(self,
address: str,
protocol: str = 'tcp',
context: 'zmq.Context' = None,
logger: 'logging.Logger' = None,
loop=None,
auth_whitelist: list = None,
using_auth: bool = True):
self.address = address
self.protocol = protocol
self.context = context or _zmq.Context.instance()
if logger is None:
logger = logging.getLogger(__name__)
self.logger = logger
self._server_certs = (None, None)
self.using_auth = using_auth
if self.using_auth:
self._setup_auth(auth_whitelist, loop)
def _setup_auth(self, auth_whitelist: list, loop):
if loop is None:
self.auth = ThreadAuthenticator(self.context)
else:
self.auth = IOLoopAuthenticator(self.context, io_loop=loop)
# allow all local host
self.logger.debug('Auth whitelist %s', auth_whitelist)
if auth_whitelist is not None:
self.auth.allow(auth_whitelist)
self._base_filepath = get_certificate_filepath()
public_key_location = path.join(self._base_filepath, 'public_keys')
self.auth.configure_curve(domain='*', location=public_key_location)
self.auth.start()
def _set_server_certs(self, any_=True):
secret_filepath, secret = get_vexbot_certificate_filepath()
if not secret and not any_:
err = ('Could not find the vexbot certificates. Generate them '
'using `vexbot_generate_certificates` from command line')
raise FileNotFoundError(err)
self._server_certs = _zmq.auth.load_certificate(secret_filepath)
def _create_using_auth(self, socket, address, bind, on_error, socket_name):
if not any(self._server_certs):
self._set_server_certs(bind)
if bind:
if self._server_certs[1] is None:
err = ('Server Secret File Not Found! Generate using '
'`vexbot_generate_certificates` from the command line')
raise FileNotFoundError(err)
public, private = self._server_certs
else:
# NOTE: This raises a `FileNotFoundError` if not found
secret_filepath = get_client_certificate_filepath()
public, private = _zmq.auth.load_certificate(secret_filepath)
socket.curve_secretkey = private
socket.curve_publickey = public
if bind:
socket.curve_server = True
try:
socket.bind(address)
except _zmq.error.ZMQError:
self._handle_error(on_error, address, socket_name)
socket = None
else: # connect the socket
socket.curve_serverkey = self._server_certs[0]
socket.connect(address)
def _create_no_auth(self, socket, address, bind, on_error, socket_name):
if bind:
try:
socket.bind(address)
except _zmq.error.ZMQError:
self._handle_error(on_error, address, socket_name)
socket = None
else: # connect the socket
socket.connect(address)
def create_n_connect(self,
socket_type,
address: str,
bind=False,
on_error='log',
socket_name=''):
"""
Creates and connects or binds the sockets
on_error:
'log': will log error
'exit': will exit the program
socket_name:
used for troubleshooting/logging
"""
self.logger.debug('create and connect: %s %s %s', socket_type,
socket_name, address)
socket = self.context.socket(socket_type)
if self.using_auth:
self._create_using_auth(socket, address, bind, on_error,
socket_name)
else:
self._create_no_auth(socket, address, bind, on_error, socket_name)
return socket
def multiple_create_n_connect(self,
socket_type,
addresses: tuple,
bind=False,
on_error='log',
socket_name=''):
# TODO: Refactor this to remove code duplication with create_n_connect
# Or at least make better sense
socket = self.context.socket(socket_type)
for address in addresses:
if self.auth:
self._create_using_auth(socket, address, bind, on_error,
socket_name)
else:
self._create_no_auth(socket, address, bind, on_error,
socket_name)
return socket
def to_address(self, port: str):
"""
transforms the ports into addresses.
Will fall through if the port looks like an address
"""
# check to see if user passed in a string
# if they did, they want to use that instead
if isinstance(port, str) and len(port) > 6:
return port
zmq_address = '{}://{}:{}'
return zmq_address.format(self.protocol, self.address, port)
def iterate_multiple_addresses(self, ports: (str, list, tuple)):
"""
transforms an iterable, or the expectation of an iterable
into zmq addresses
"""
if isinstance(ports, (str, int)):
# TODO: verify this works.
ports = tuple(ports, )
result = []
for port in ports:
result.append(self.to_address(port))
return result
def _handle_error(self, how_to: str, address: str, socket_name: str):
if socket_name == '':
socket_name = 'unknown type'
if how_to == 'exit':
self._handle_bind_error_by_exit(address, socket_name)
else:
self._handle_bind_error_by_log(address, socket_name)
def _handle_bind_error_by_log(self, address: str, socket_name: str):
if self.logger is None:
return
s = 'Address bind attempt fail. Address tried: {}'
s = s.format(address)
self.logger.error(s)
self.logger.error('socket type: {}'.format(socket_name))
def _handle_bind_error_by_exit(self, address: str, socket_type: str):
if self.logger is not None:
s = 'Address bind attempt fail. Alredy in use? Address tried: {}'
s = s.format(address)
self.logger.error(s)
self.logger.error('socket type: {}'.format(socket_type))
_sys.exit(1)
def message_proxy(self, work_dir):
"""
drone_data_inboud is for data comming from drones
drone_data_outbound is for commands to the drone, topic must either be a drone ID or all for sending
a broadcast message to all drones
"""
ctx = zmq.Context()
public_keys_dir = os.path.join(work_dir, 'certificates', 'public_keys')
secret_keys_dir = os.path.join(work_dir, 'certificates', 'private_keys')
# start and configure auth worker
auth = IOLoopAuthenticator()
auth.start()
auth.allow('127.0.0.1')
auth.configure_curve(domain='*', location=public_keys_dir)
# external interfaces for communicating with drones
server_secret_file = os.path.join(secret_keys_dir, 'beeswarm_server.pri')
server_public, server_secret = load_certificate(server_secret_file)
drone_data_inbound = ctx.socket(zmq.PULL)
drone_data_inbound.curve_secretkey = server_secret
drone_data_inbound.curve_publickey = server_public
drone_data_inbound.curve_server = True
drone_data_inbound.bind('tcp://*:{0}'.format(self.config['network']['zmq_port']))
drone_data_outbound = ctx.socket(zmq.PUB)
drone_data_outbound.curve_secretkey = server_secret
drone_data_outbound.curve_publickey = server_public
drone_data_outbound.curve_server = True
drone_data_outbound.bind('tcp://*:{0}'.format(self.config['network']['zmq_command_port']))
# internal interfaces
# all inbound session data from drones will be replayed in this socket
sessionPublisher = ctx.socket(zmq.PUB)
sessionPublisher.bind('ipc://sessionPublisher')
# all commands received on this will be published on the external interface
drone_command_receiver = ctx.socket(zmq.PULL)
drone_command_receiver.bind('ipc://droneCommandReceiver')
poller = zmq.Poller()
poller.register(drone_data_inbound, zmq.POLLIN)
poller.register(drone_command_receiver, zmq.POLLIN)
while True:
# .recv() gives no context switch - why not? using poller with timeout instead
socks = dict(poller.poll(1))
gevent.sleep()
if drone_command_receiver in socks and socks[drone_command_receiver] == zmq.POLLIN:
data = drone_command_receiver.recv()
topic, _ = data.split(' ', 1)
logger.debug("Sending drone command to: {0}".format(topic))
# pub socket takes care of filtering
drone_data_outbound.send(data)
elif drone_data_inbound in socks and socks[drone_data_inbound] == zmq.POLLIN:
topic, data = drone_data_inbound.recv().split(' ', 1)
logger.debug("Received {0} data.".format(topic))
if topic == Messages.SESSION_HONEYPOT or topic == Messages.SESSION_CLIENT:
sessionPublisher.send('{0} {1}'.format(topic, data))
elif topic == Messages.KEY or topic == Messages.CERT:
# for now we just store the fingerprint
# in the future it might be relevant to store the entire public key and private key
# for forensic purposes
if topic == Messages.CERT:
drone_id, cert = data.split(' ', 1)
digest = generate_cert_digest(cert)
logging.debug('Storing public key digest: {0} for drone {1}.'.format(digest, drone_id))
db_session = database_setup.get_session()
drone = db_session.query(Drone).filter(Drone.id == drone_id).one()
drone.cert_digest = digest
db_session.add(drone)
db_session.commit()
elif topic == Messages.PING:
drone_id = data
db_session = database_setup.get_session()
drone = db_session.query(Drone).filter(Drone.id == drone_id).one()
drone.last_activity = datetime.now()
db_session.add(drone)
db_session.commit()
else:
logger.warn('Message with unknown topic received: {0}'.format(topic))
def message_proxy(self, work_dir):
"""
drone_data_inboud is for data comming from drones
drone_data_outbound is for commands to the drones, topic must either be a drone ID or all for sending
a broadcast message to all drones
"""
public_keys_dir = os.path.join(work_dir, 'certificates', 'public_keys')
secret_keys_dir = os.path.join(work_dir, 'certificates',
'private_keys')
# start and configure auth worker
auth = IOLoopAuthenticator()
auth.start()
auth.allow('127.0.0.1')
auth.configure_curve(domain='*', location=public_keys_dir)
# external interfaces for communicating with drones
server_secret_file = os.path.join(secret_keys_dir,
'beeswarm_server.pri')
server_public, server_secret = load_certificate(server_secret_file)
drone_data_inbound = beeswarm.shared.zmq_context.socket(zmq.PULL)
drone_data_inbound.curve_secretkey = server_secret
drone_data_inbound.curve_publickey = server_public
drone_data_inbound.curve_server = True
drone_data_inbound.bind('tcp://*:{0}'.format(
self.config['network']['zmq_port']))
drone_data_outbound = beeswarm.shared.zmq_context.socket(zmq.PUB)
drone_data_outbound.curve_secretkey = server_secret
drone_data_outbound.curve_publickey = server_public
drone_data_outbound.curve_server = True
drone_data_outbound.bind('tcp://*:{0}'.format(
self.config['network']['zmq_command_port']))
# internal interfaces
# all inbound session data from drones will be replayed in this socket
sessionPublisher = beeswarm.shared.zmq_context.socket(zmq.PUB)
sessionPublisher.bind('inproc://sessionPublisher')
# all commands received on this will be published on the external interface
drone_command_receiver = beeswarm.shared.zmq_context.socket(zmq.PULL)
drone_command_receiver.bind('inproc://droneCommandReceiver')
poller = zmq.Poller()
poller.register(drone_data_inbound, zmq.POLLIN)
poller.register(drone_command_receiver, zmq.POLLIN)
while True:
# .recv() gives no context switch - why not? using poller with timeout instead
socks = dict(poller.poll(100))
gevent.sleep()
if drone_command_receiver in socks and socks[
drone_command_receiver] == zmq.POLLIN:
data = drone_command_receiver.recv()
drone_id, _ = data.split(' ', 1)
logger.debug("Sending drone command to: {0}".format(drone_id))
# pub socket takes care of filtering
drone_data_outbound.send(data)
elif drone_data_inbound in socks and socks[
drone_data_inbound] == zmq.POLLIN:
split_data = drone_data_inbound.recv().split(' ', 2)
if len(split_data) == 3:
topic, drone_id, data = split_data
else:
data = None
topic, drone_id, = split_data
logger.debug("Received {0} message from {1}.".format(
topic, drone_id))
db_session = database_setup.get_session()
drone = db_session.query(Drone).filter(
Drone.id == drone_id).one()
drone.last_activity = datetime.now()
db_session.add(drone)
db_session.commit()
if topic == Messages.SESSION_HONEYPOT or topic == Messages.SESSION_CLIENT:
sessionPublisher.send('{0} {1}'.format(topic, data))
elif topic == Messages.KEY or topic == Messages.CERT:
# for now we just store the fingerprint
# in the future it might be relevant to store the entire public key and private key
# for forensic purposes
if topic == Messages.CERT:
cert = data.split(' ', 1)[1]
digest = generate_cert_digest(cert)
logging.debug(
'Storing public key digest: {0} for drone {1}.'.
format(digest, drone_id))
db_session = database_setup.get_session()
drone = db_session.query(Drone).filter(
Drone.id == drone_id).one()
drone.cert_digest = digest
db_session.add(drone)
db_session.commit()
elif topic == Messages.PING:
pass
elif topic == Messages.IP:
ip_address = data
logging.debug('Drone {0} reported ip: {1}'.format(
drone_id, ip_address))
db_session = database_setup.get_session()
drone = db_session.query(Drone).filter(
Drone.id == drone_id).one()
if drone.ip_address != ip_address:
drone.ip_address = ip_address
db_session.add(drone)
db_session.commit()
send_zmq_request(
'inproc://configCommands',
'{0} {1}'.format(Messages.DRONE_CONFIG_CHANGED,
drone_id))
# drone want it's config transmitted
elif topic == Messages.DRONE_CONFIG:
config_dict = send_zmq_request(
'inproc://configCommands',
'{0} {1}'.format(Messages.DRONE_CONFIG, drone_id))
drone_data_outbound.send('{0} {1} {2}'.format(
drone_id, Messages.CONFIG, json.dumps(config_dict)))
else:
logger.warn(
'Message with unknown topic received: {0}'.format(
topic))
def message_proxy(self, work_dir):
"""
drone_data_inboud is for data comming from drones
drone_data_outbound is for commands to the drones, topic must either be a drone ID or all for sending
a broadcast message to all drones
"""
public_keys_dir = os.path.join(work_dir, 'certificates', 'public_keys')
secret_keys_dir = os.path.join(work_dir, 'certificates',
'private_keys')
# start and configure auth worker
auth = IOLoopAuthenticator()
auth.start()
auth.allow('127.0.0.1')
auth.configure_curve(domain='*', location=public_keys_dir)
# external interfaces for communicating with drones
server_secret_file = os.path.join(secret_keys_dir,
'beeswarm_server.pri')
server_public, server_secret = load_certificate(server_secret_file)
drone_data_inbound = beeswarm.shared.zmq_context.socket(zmq.PULL)
drone_data_inbound.curve_secretkey = server_secret
drone_data_inbound.curve_publickey = server_public
drone_data_inbound.curve_server = True
drone_data_inbound.bind('tcp://*:{0}'.format(
self.config['network']['zmq_port']))
drone_data_outbound = beeswarm.shared.zmq_context.socket(zmq.PUB)
drone_data_outbound.curve_secretkey = server_secret
drone_data_outbound.curve_publickey = server_public
drone_data_outbound.curve_server = True
drone_data_outbound.bind('tcp://*:{0}'.format(
self.config['network']['zmq_command_port']))
# internal interfaces
# all inbound session data from drones will be replayed on this socket
drone_data_socket = beeswarm.shared.zmq_context.socket(zmq.PUB)
drone_data_socket.bind(SocketNames.DRONE_DATA.value)
# all commands received on this will be published on the external interface
drone_command_socket = beeswarm.shared.zmq_context.socket(zmq.PULL)
drone_command_socket.bind(SocketNames.DRONE_COMMANDS.value)
poller = zmq.Poller()
poller.register(drone_data_inbound, zmq.POLLIN)
poller.register(drone_command_socket, zmq.POLLIN)
while True:
# .recv() gives no context switch - why not? using poller with timeout instead
socks = dict(poller.poll(100))
gevent.sleep()
if drone_command_socket in socks and socks[
drone_command_socket] == zmq.POLLIN:
data = drone_command_socket.recv()
drone_id, _ = data.split(' ', 1)
logger.debug("Sending drone command to: {0}".format(drone_id))
# pub socket takes care of filtering
drone_data_outbound.send(data)
elif drone_data_inbound in socks and socks[
drone_data_inbound] == zmq.POLLIN:
raw_msg = drone_data_inbound.recv()
split_data = raw_msg.split(' ', 2)
if len(split_data) == 3:
topic, drone_id, data = split_data
else:
data = None
topic, drone_id, = split_data
logger.debug("Received {0} message from {1}.".format(
topic, drone_id))
# relay message on internal socket
drone_data_socket.send(raw_msg)
def message_proxy(self, work_dir):
"""
drone_data_inboud is for data comming from drones
drone_data_outbound is for commands to the drones, topic must either be a drone ID or all for sending
a broadcast message to all drones
"""
public_keys_dir = os.path.join(work_dir, 'certificates', 'public_keys')
secret_keys_dir = os.path.join(work_dir, 'certificates', 'private_keys')
# start and configure auth worker
auth = IOLoopAuthenticator()
auth.start()
auth.allow('127.0.0.1')
auth.configure_curve(domain='*', location=public_keys_dir)
# external interfaces for communicating with drones
server_secret_file = os.path.join(secret_keys_dir, 'beeswarm_server.pri')
server_public, server_secret = load_certificate(server_secret_file)
drone_data_inbound = beeswarm.shared.zmq_context.socket(zmq.PULL)
drone_data_inbound.curve_secretkey = server_secret
drone_data_inbound.curve_publickey = server_public
drone_data_inbound.curve_server = True
drone_data_inbound.bind('tcp://*:{0}'.format(self.config['network']['zmq_port']))
drone_data_outbound = beeswarm.shared.zmq_context.socket(zmq.PUB)
drone_data_outbound.curve_secretkey = server_secret
drone_data_outbound.curve_publickey = server_public
drone_data_outbound.curve_server = True
drone_data_outbound.bind('tcp://*:{0}'.format(self.config['network']['zmq_command_port']))
# internal interfaces
# all inbound session data from drones will be replayed on this socket
drone_data_socket = beeswarm.shared.zmq_context.socket(zmq.PUB)
drone_data_socket.bind(SocketNames.DRONE_DATA.value)
# all commands received on this will be published on the external interface
drone_command_socket = beeswarm.shared.zmq_context.socket(zmq.PULL)
drone_command_socket.bind(SocketNames.DRONE_COMMANDS.value)
poller = zmq.Poller()
poller.register(drone_data_inbound, zmq.POLLIN)
poller.register(drone_command_socket, zmq.POLLIN)
while True:
# .recv() gives no context switch - why not? using poller with timeout instead
socks = dict(poller.poll(100))
gevent.sleep()
if drone_command_socket in socks and socks[drone_command_socket] == zmq.POLLIN:
data = drone_command_socket.recv()
drone_id, _ = data.split(' ', 1)
logger.debug("Sending drone command to: {0}".format(drone_id))
# pub socket takes care of filtering
drone_data_outbound.send(data)
elif drone_data_inbound in socks and socks[drone_data_inbound] == zmq.POLLIN:
raw_msg = drone_data_inbound.recv()
split_data = raw_msg.split(' ', 2)
if len(split_data) == 3:
topic, drone_id, data = split_data
else:
data = None
topic, drone_id, = split_data
logger.debug("Received {0} message from {1}.".format(topic, drone_id))
# relay message on internal socket
drone_data_socket.send(raw_msg)
