beeswarm/server/db/entities.py
class Honeypot(Drone):
    __tablename__ = 'honeypot'
    id = Column(String, ForeignKey('drone.id'), primary_key=True)
    __mapper_args__ = {'polymorphic_identity': 'honeypot'}

    sessions = relationship('Session', cascade='all, delete-orphan', backref='honeypot')
    clients = relationship('Client', secondary=honeypot_client_mtm)
    # current capabilities
    capabilities = relationship('Capability', cascade='all, delete-orphan', backref='honeypot')
    # fingerprint of the public key used to interact with attackers and clients
    cert_digest = Column(String)

    # The following certificate attribute are temporarily
    # generation of certificate should be done on the server, hence no need for
    # this information here
    cert_common_name = Column(String)
    cert_country = Column(String)
    cert_state = Column(String)
    cert_locality = Column(String)
    cert_organization = Column(String)
    cert_organization_unit = Column(String)

    def add_capability(self, protocol, port, protocol_specific_data):
        capability = Capability(protocol=protocol, port=port, protocol_specific_data=json.dumps(protocol_specific_data))
        self.capabilities.append(capability)
beeswarm/server/db/database_actor.py
class DatabaseActor(gevent.Greenlet):

    def __init__(self, max_sessions, clear_sessions=False, delay_seconds=30):
        assert delay_seconds > 1
        Greenlet.__init__(self)
        db_session = database_setup.get_session()
        self.enabled = True
        # pending session will be converted to attacks if we cannot match with bait traffic
        # with this period
        self.delay_seconds = delay_seconds
        # clear all pending sessions on startup, pending sessions on startup
        pending_classification = db_session.query(Classification).filter(
            Classification.type == 'pending').one()
        pending_deleted = db_session.query(Session).filter(
            Session.classification == pending_classification).delete()
        db_session.commit()
        logging.info(
            'Cleaned {0} pending sessions on startup'.format(pending_deleted))
        self.do_classify = False
        self.do_maintenance = False
        if clear_sessions or max_sessions == 0:
            db_session = database_setup.get_session()
            count = db_session.query(Session).delete()
            logging.info('Deleting {0} sessions on startup.'.format(count))
            db_session.commit()

        self.max_session_count = max_sessions
        if max_sessions:
            logger.info(
                'Database has been limited to contain {0} sessions.'.format(max_sessions))

        context = beeswarm.shared.zmq_context

        # prepare sockets
        self.drone_data_socket = context.socket(zmq.SUB)
        self.processedSessionsPublisher = context.socket(zmq.PUB)
        self.databaseRequests = context.socket(zmq.REP)
        self.config_actor_socket = context.socket(zmq.REQ)
        self.drone_command_receiver = context.socket(zmq.PUSH)

    def _run(self):
        # connect and bind to all relevant sockets
        # raw data from drones
        self.drone_data_socket.connect(SocketNames.DRONE_DATA.value)
        self.drone_data_socket.setsockopt(zmq.SUBSCRIBE, '')
        # requests that this actor needs to respond to
        self.databaseRequests.bind(SocketNames.DATABASE_REQUESTS.value)
        # will publish session after they have been processed on this socket
        self.processedSessionsPublisher.bind(
            SocketNames.PROCESSED_SESSIONS.value)
        # needed to be able to probe for options and get zmq keys
        self.config_actor_socket.connect(SocketNames.CONFIG_COMMANDS.value)
        # needed to send data directly to drones
        self.drone_command_receiver.connect(SocketNames.DRONE_COMMANDS.value)

        poller = zmq.Poller()
        poller.register(self.drone_data_socket, zmq.POLLIN)
        poller.register(self.databaseRequests, zmq.POLLIN)
        gevent.spawn(self._start_recurring_classify_set)
        gevent.spawn(self._start_recurring_maintenance_set)
        while self.enabled:
            socks = dict(poller.poll(100))
            if self.do_classify:
                self._classify_malicious_sessions()
                self.do_classify = False
            elif self.do_maintenance:
                self._db_maintenance()
                self.do_maintenance = False
            elif self.databaseRequests in socks and socks[self.databaseRequests] == zmq.POLLIN:
                data = self.databaseRequests.recv()
                if ' ' in data:
                    cmd, data = data.split(' ', 1)
                else:
                    cmd = data

                if cmd == Messages.DRONE_CONFIG.value:
                    self._handle_command_get_droneconfig(data)
                elif cmd == Messages.BAIT_USER_ADD.value:
                    self._handle_command_bait_user_add(data)
                    self.databaseRequests.send(
                        '{0} {1}'.format(Messages.OK.value, '{}'))
                elif cmd == Messages.BAIT_USER_DELETE.value:
                    self._handle_command_bait_user_delete(data)
                    self.databaseRequests.send(
                        '{0} {1}'.format(Messages.OK.value, '{}'))
                elif cmd == Messages.DRONE_DELETE.value:
                    self._handle_command_delete_drone(data)
                    self.databaseRequests.send(
                        '{0} {1}'.format(Messages.OK.value, '{}'))
                elif cmd == Messages.DRONE_ADD.value:
                    self._handle_command_add_drone()
                elif cmd == Messages.GET_DB_STATS.value:
                    result = self._handle_command_get_db_stats()
                    self.databaseRequests.send('{0} {1}'.format(
                        Messages.OK.value, json.dumps(result)))
                elif cmd == Messages.GET_SESSIONS_ALL.value or cmd == Messages.GET_SESSIONS_ATTACKS.value \
                        or cmd == Messages.GET_SESSIONS_BAIT.value:
                    # TODO: Accept to/from param to facilitate pagination
                    result = self._handle_command_get_sessions(cmd)
                    self.databaseRequests.send('{0} {1}'.format(
                        Messages.OK.value, json.dumps(result)))
                elif cmd == Messages.GET_SESSION_CREDENTIALS.value:
                    result = self._handle_command_get_credentials(data)
                    self.databaseRequests.send('{0} {1}'.format(
                        Messages.OK.value, json.dumps(result)))
                elif cmd == Messages.GET_SESSION_TRANSCRIPT.value:
                    result = self._handle_command_get_transcript(data)
                    self.databaseRequests.send('{0} {1}'.format(
                        Messages.OK.value, json.dumps(result)))
                elif cmd == Messages.GET_BAIT_USERS.value:
                    result = self._handle_command_get_bait_users()
                    self.databaseRequests.send('{0} {1}'.format(
                        Messages.OK.value, json.dumps(result)))
                elif cmd == Messages.CONFIG_DRONE.value:
                    # .send on socket is handled internally since it can send errors back
                    self._handle_command_config_drone(data)
                elif cmd == Messages.GET_DRONE_LIST.value:
                    result = self._handle_command_get_drone_list(data)
                    self.databaseRequests.send('{0} {1}'.format(
                        Messages.OK.value, json.dumps(result)))
                elif cmd == Messages.PING_ALL_DRONES.value:
                    self._handle_ping_all_drones()
                    self.databaseRequests.send(
                        '{0} {1}'.format(Messages.OK.value, ' '))
                else:
                    logger.error('Unknown message received: {0}'.format(data))
                    assert False
            elif self.drone_data_socket in socks and socks[self.drone_data_socket] == zmq.POLLIN:
                split_data = self.drone_data_socket.recv().split(' ', 2)
                if len(split_data) == 3:
                    topic, drone_id, data = split_data
                else:
                    data = None
                    topic, drone_id, = split_data
                self._update_drone_last_activity(drone_id)
                if topic == Messages.SESSION_HONEYPOT.value or topic == Messages.SESSION_CLIENT.value:
                    self.persist_session(topic, data)
                elif topic == Messages.CERT.value:
                    self._handle_cert_message(topic, drone_id, data)
                elif topic == Messages.KEY.value:
                    pass
                elif topic == Messages.IP.value:
                    self._handle_message_ip(topic, drone_id, data)
                elif topic == Messages.DRONE_WANT_CONFIG.value:
                    config_dict = self._get_drone_config(drone_id)
                    self.drone_command_receiver.send('{0} {1} {2}'.format(drone_id, Messages.CONFIG.value,
                                                                          json.dumps(config_dict)))
                elif topic == Messages.PING.value:
                    self._update_drone_last_activity(drone_id)
                    logger.debug('Received ping from {0}'.format(drone_id))
                else:
                    logger.debug(
                        'This actor cannot process this message: {0}'.format(topic))

    def stop(self):
        self.enabled = False
        self.drone_data_socket.close()
        self.processedSessionsPublisher.close()
        self.databaseRequests.close()
        self.config_actor_socket.close()
        self.drone_command_receiver.close()

    def _update_drone_last_activity(self, drone_id):
        db_session = database_setup.get_session()
        drone = db_session.query(Drone).filter(Drone.id == drone_id).one_or_none()
        if drone:
            drone.last_activity = datetime.now()
            db_session.add(drone)
            db_session.commit()
        else:
            logger.warning(
                'Trying to update last activity non-exting drone with id {0}'.format(drone_id))

    def _start_recurring_maintenance_set(self):
        while self.enabled:
            sleep_time = 60 * 60
            gevent.sleep(sleep_time)
            self.do_maintenance = True

    def _start_recurring_classify_set(self):
        while self.enabled:
            sleep_time = self.delay_seconds / 2
            gevent.sleep(sleep_time)
            self.do_classify = True

    def _handle_message_ip(self, topic, drone_id, data):
        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_or_none()
        if drone:
            if drone.ip_address != ip_address:
                drone.ip_address = ip_address
                db_session.add(drone)
                db_session.commit()
        else:
            logger.warning(
                'Trying to update IP on non-exting drone with id {0}'.format(drone_id))

    def _handle_cert_message(self, topic, drone_id, data):
        # 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
        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_or_none()
        if drone:
            drone.cert_digest = digest
            db_session.add(drone)
            db_session.commit()
        else:
            logger.warning(
                'Trying to update cert on non-exting drone with id {0}'.format(drone_id))

    def persist_session(self, session_type, session_json):
        db_session = database_setup.get_session()

        if self.max_session_count == 0:
            return
        elif db_session.query(Session).count() == self.max_session_count:
            session_to_delete = db_session.query(
                Session, func.min(Session.timestamp)).first()[0]
            db_session.delete(session_to_delete)
        try:
            data = json.loads(session_json)
        except UnicodeDecodeError:
            data = json.loads(unicode(session_json, "ISO-8859-1"))
        logger.debug('Persisting {0} session: {1}'.format(session_type, data))

        classification = db_session.query(Classification).filter(
            Classification.type == 'pending').one()

        assert data['honeypot_id'] is not None
        _honeypot = db_session.query(Honeypot).filter(
            Honeypot.id == data['honeypot_id']).one_or_none()
        if not _honeypot:
            logger.warning('Trying to persist session for non-existing honeypot drone with id {0}'.format(data['honeypot_id']))
            return
        if session_type == Messages.SESSION_HONEYPOT.value:
            session = Session()
            for entry in data['transcript']:
                transcript_timestamp = isoformatToDatetime(entry['timestamp'])
                transcript = Transcript(timestamp=transcript_timestamp, direction=entry['direction'],
                                        data=entry['data'])
                session.transcript.append(transcript)

            for auth in data['login_attempts']:
                authentication = self.extract_auth_entity(auth)
                session.authentication.append(authentication)
        elif session_type == Messages.SESSION_CLIENT.value:
            ignore_failed_bait_sessions = self.send_config_request('{0} {1}'.format(Messages.GET_CONFIG_ITEM.value,
                                                                                    'ignore_failed_bait_session'))
            if not data['did_complete'] and ignore_failed_bait_sessions:
                logger.debug('Ignore failed bait session.')
                return
            session = BaitSession()
            client = db_session.query(Client).filter(
                Client.id == data['client_id']).one_or_none()
            if not client:
                logger.warning('Trying to persist session for non-existing client drone with id {0}'.format(data['client_id']))
                return
            client.last_activity = datetime.now()
            session.did_connect = data['did_connect']
            session.did_login = data['did_login']
            session.did_complete = data['did_complete']
            session.client = client
            for auth in data['login_attempts']:
                authentication = self.extract_auth_entity(auth)
                session.authentication.append(authentication)
        else:
            logger.warn('Unknown message type: {0}'.format(session_type))
            return
        session.id = data['id']
        session.classification = classification
        session.timestamp = isoformatToDatetime(data['timestamp'])
        session.received = datetime.utcnow()
        session.protocol = data['protocol']
        session.destination_ip = data['destination_ip']
        session.destination_port = data['destination_port']
        session.source_ip = data['source_ip']
        session.source_port = data['source_port']
        session.honeypot = _honeypot
        _dronename = db_session.query(Drone).filter(
            Drone.id == _honeypot.id).first().name
        db_session.add(session)
        db_session.commit()
        matching_session = self.get_matching_session(session, db_session)
        if session_type == Messages.SESSION_HONEYPOT.value:
            if matching_session:
                matching_session.name = _dronename
                self.merge_bait_and_session(
                    session, matching_session, db_session)
        elif session_type == Messages.SESSION_CLIENT.value:
            if matching_session:
                session.name = _dronename
                self.merge_bait_and_session(
                    matching_session, session, db_session)
        else:
            assert False

    def extract_auth_entity(self, auth_data):
        username = auth_data.get('username', '')
        password = auth_data.get('password', '')
        authentication = Authentication(id=auth_data['id'], username=username, password=password,
                                        successful=auth_data['successful'],
                                        timestamp=isoformatToDatetime(auth_data['timestamp']))
        return authentication

    def get_matching_session(self, session, db_session, timediff=5):
        """
        Tries to match a session with it's counterpart. For bait session it will try to match it with honeypot sessions
        and the other way around.

        :param session: session object which will be used as base for query.
        :param timediff: +/- allowed time difference between a session and a potential matching session.
        """
        db_session = db_session
        min_datetime = session.timestamp - timedelta(seconds=timediff)
        max_datetime = session.timestamp + timedelta(seconds=timediff)
        # default return value
        match = None
        classification = db_session.query(Classification).filter(
            Classification.type == 'pending').one()
        # get all sessions that match basic properties.
        sessions = db_session.query(Session).options(joinedload(Session.authentication)) \
            .filter(Session.protocol == session.protocol) \
            .filter(Session.honeypot == session.honeypot) \
            .filter(Session.timestamp >= min_datetime) \
            .filter(Session.timestamp <= max_datetime) \
            .filter(Session.id != session.id) \
            .filter(Session.classification == classification)

        # identify the correct session by comparing authentication.
        # this could properly also be done using some fancy ORM/SQL construct.
        for potential_match in sessions:
            if potential_match.discriminator == session.discriminator:
                continue
            assert potential_match.id != session.id
            for honey_auth in session.authentication:
                for session_auth in potential_match.authentication:
                    if session_auth.username == honey_auth.username and \
                            session_auth.password == honey_auth.password and \
                            session_auth.successful == honey_auth.successful:
                        assert potential_match.id != session.id
                        match = potential_match
                        break

        return match

    def merge_bait_and_session(self, honeypot_session, bait_session, db_session):
        logger.debug('Classifying bait session with id {0} as legit bait and deleting '
                     'matching honeypot_session with id {1}'.format(bait_session.id, honeypot_session.id))
        bait_session.classification = db_session.query(Classification).filter(
            Classification.type == 'bait_session').one()
        bait_session.transcript = honeypot_session.transcript
        bait_session.session_data = honeypot_session.session_data
        # the client ip own detection can be flawed, but we are sure the the honeypot tcp source ip
        # is the same as the clients source ip.
        bait_session.source_ip = honeypot_session.source_ip
        db_session.add(bait_session)
        db_session.delete(honeypot_session)
        db_session.commit()
        self.processedSessionsPublisher.send(
            '{0} {1} {2}'.format(Messages.DELETED_DUE_TO_MERGE.value, honeypot_session.id,
                                 bait_session.id))
        self.processedSessionsPublisher.send('{0} {1}'.format(Messages.SESSION.value,
                                                              json.dumps(bait_session.to_dict())))

    def _classify_malicious_sessions(self):
        """
        Will classify all unclassified sessions as malicious activity.

        :param delay_seconds: no sessions newer than (now - delay_seconds) will be processed.
        """
        min_datetime = datetime.utcnow() - timedelta(seconds=self.delay_seconds)

        db_session = database_setup.get_session()

        # find and process bait sessions that did not get classified during
        # persistence.
        bait_sessions = db_session.query(BaitSession).options(joinedload(BaitSession.authentication)) \
            .filter(BaitSession.classification_id == 'pending') \
            .filter(BaitSession.did_complete == True) \
            .filter(BaitSession.received < min_datetime).all()

        for bait_session in bait_sessions:
            logger.debug(
                'Classifying bait session with id {0} as MITM'.format(bait_session.id))
            bait_session.classification = db_session.query(
                Classification).filter(Classification.type == 'mitm').one()
            db_session.commit()

        # find and process honeypot sessions that did not get classified during
        # persistence.
        sessions = db_session.query(Session, Drone.name).filter(Session.discriminator == None) \
            .filter(Session.timestamp <= min_datetime) \
            .filter(Session.classification_id == 'pending') \
            .all()

        for entry in sessions:
            # Check if the attack used credentials leaked by beeswarm drones
            session = entry[0]
            bait_match = None
            for a in session.authentication:
                bait_match = db_session.query(BaitSession) \
                    .filter(BaitSession.authentication.any(username=a.username, password=a.password)).first()
                if bait_match:
                    break

            if bait_match:
                logger.debug('Classifying session with id {0} as attack which involved the reuse '
                             'of previously transmitted credentials.'.format(session.id))
                session.classification = db_session.query(Classification).filter(
                    Classification.type == 'credentials_reuse').one()
            elif len(session.authentication) == 0:
                logger.debug(
                    'Classifying session with id {0} as probe.'.format(session.id))
                session.classification = db_session.query(
                    Classification).filter(Classification.type == 'probe').one()
            else:
                # we have never transmitted this username/password combo
                logger.debug(
                    'Classifying session with id {0} as bruteforce attempt.'.format(session.id))
                session.classification = db_session.query(Classification).filter(
                    Classification.type == 'bruteforce').one()
            db_session.commit()
            session.name = entry[1]
            self.processedSessionsPublisher.send(
                '{0} {1}'.format(Messages.SESSION.value, json.dumps(session.to_dict())))

    def _handle_command_delete_drone(self, data):
        drone_id = data
        logger.debug('Deleting drone: {0}'.format(drone_id))
        db_session = database_setup.get_session()
        drone_to_delete = db_session.query(
            Drone).filter(Drone.id == drone_id).first()
        if drone_to_delete:
            db_session.delete(drone_to_delete)
            db_session.commit()
            # tell the drone to kill itself
            self.drone_command_receiver.send(
                '{0} {1} '.format(drone_id, Messages.DRONE_DELETE.value))
            self.send_config_request('{0} {1}'.format(
                Messages.DELETE_ZMQ_KEYS.value, drone_id))
            self._reconfigure_all_clients()

    def _handle_command_add_drone(self):
        db_session = database_setup.get_session()
        drone = Drone()
        db_session.add(drone)
        db_session.commit()
        logger.debug('New drone has been added with id: {0}'.format(drone.id))

        drone_config = self._get_drone_config(drone.id)
        self.databaseRequests.send('{0} {1}'.format(
            Messages.OK.value, json.dumps(drone_config)))

    def _reconfigure_all_clients(self):
        db_session = database_setup.get_session()
        db_session.query(DroneEdge).delete()
        db_session.commit()
        honeypots = db_session.query(Honeypot).all()
        clients = db_session.query(Client).all()
        # delete old architecture
        credentials = db_session.query(BaitUser).all()
        for honeypot in honeypots:
            for capability in honeypot.capabilities:
                for client in clients:
                    # following three variables should be make somewhat user
                    # configurable again
                    client_timings = json.loads(client.bait_timings)
                    if capability.protocol in client_timings:
                        # the time range in which to activate the bait sessions
                        activation_range = client_timings[
                            capability.protocol]['active_range']
                        # period to sleep before using activation_probability
                        sleep_interval = client_timings[
                            capability.protocol]['sleep_interval']
                        # the probability that a bait session will be
                        # activated, 1 is always activate
                        activation_probability = client_timings[
                            capability.protocol]['activation_probability']
                    else:
                        logger.warning('Bait timings for {0} not found on client drone {1}({2}), using defaults instead'
                                       .format(capability.protocol, client.name, client.id))
                        activation_range = '00:00 - 23:59'
                        sleep_interval = '60'
                        activation_probability = 1
                    bait_credentials = random.choice(credentials)
                    client.add_bait(capability, activation_range, sleep_interval,
                                    activation_probability, bait_credentials.username, bait_credentials.password)
        db_session.commit()

        drones = db_session.query(Drone).all()
        for drone in drones:
            self._send_config_to_drone(drone.id)

    def _handle_command_get_droneconfig(self, drone_id):
        result = self._get_drone_config(drone_id)
        if len(result) == 0:
            self.databaseRequests.send('{0} {1}'.format(
                Messages.FAIL.value, 'Drone could not be found'))
        else:
            self.databaseRequests.send('{0} {1}'.format(
                Messages.OK.value, json.dumps(result)))

    def _send_config_to_drone(self, drone_id):
        config = self._get_drone_config(drone_id)
        logger.debug('Sending config to {0}: {1}'.format(drone_id, config))
        self.drone_command_receiver.send('{0} {1} {2}'.format(
            drone_id, Messages.CONFIG.value, json.dumps(config)))

    def _handle_ping_all_drones(self):
        db_session = database_setup.get_session()
        drones = db_session.query(Drone)

        for drone in drones:
            logger.debug('Sending ping to {0}'.format(drone.id))
            self.drone_command_receiver.send(
                '{0} {1} {2}'.format(drone.id, Messages.PING.value, ''))

    def send_config_request(self, request):
        return send_zmq_request_socket(self.config_actor_socket, request)

    def _handle_command_drone_config_changed(self, drone_id):
        self._send_config_to_drone(drone_id)
        # TODO: Only Clients that communicate with this drone_id needs to get
        # reconfigured.
        self._reconfigure_all_clients()

    def _handle_command_bait_user_delete(self, data):
        bait_user_id = int(data)
        db_session = database_setup.get_session()
        bait_user = db_session.query(BaitUser).filter(
            BaitUser.id == bait_user_id).first()
        if bait_user:
            db_session.delete(bait_user)
            db_session.commit()
            drone_edge = db_session.query(DroneEdge).filter(DroneEdge.username == bait_user.username,
                                                            DroneEdge.password == bait_user.password).first()
            # A drone is using the bait users, reconfigure all
            # TODO: This is lazy, we should only reconfigure the drone(s) who are actually
            # using the credentials
            if drone_edge:
                self._reconfigure_all_clients()
        else:
            logger.warning(
                'Tried to delete non-existing bait user with id {0}.'.format(bait_user_id))

    def _handle_command_bait_user_add(self, data):
        username, password = data.split(' ')
        db_session = database_setup.get_session()
        existing_bait_user = db_session.query(BaitUser).filter(BaitUser.username == username,
                                                               BaitUser.password == password).first()
        if not existing_bait_user:
            new_bait_user = BaitUser(username=username, password=password)
            db_session.add(new_bait_user)
            db_session.commit()

    def _handle_command_get_bait_users(self):
        db_session = database_setup.get_session()
        bait_users = db_session.query(BaitUser)
        return_rows = []
        for bait_user in bait_users:
            row = {'id': bait_user.id, 'username': bait_user.username,
                   'password': bait_user.password}
            return_rows.append(row)
        return return_rows

    def _get_drone_config(self, drone_id):
        db_session = database_setup.get_session()
        drone = db_session.query(Honeypot).filter(Drone.id == drone_id).first()
        # lame! what is the correct way?
        if not drone:
            drone = db_session.query(Client).filter(
                Drone.id == drone_id).first()
        if not drone:
            drone = db_session.query(Drone).filter(
                Drone.id == drone_id).first()
        if not drone:
            # drone not found
            return {}

        host = self.send_config_request('{0} {1}'.format(
            Messages.GET_CONFIG_ITEM.value, 'network,server_host'))
        zmq_port = self.send_config_request('{0} {1}'.format(
            Messages.GET_CONFIG_ITEM.value, 'network,zmq_port'))
        zmq_command_port = self.send_config_request(
            '{0} {1}'.format(Messages.GET_CONFIG_ITEM.value, 'network,zmq_command_port'))

        server_zmq_url = 'tcp://{0}:{1}'.format(host, zmq_port)
        server_zmq_command_url = 'tcp://{0}:{1}'.format(host, zmq_command_port)

        zmq_keys = self.send_config_request(
            '{0} {1}'.format(Messages.GET_ZMQ_KEYS.value, drone_id))
        zmq_server_key = self.send_config_request(
            '{0} {1}'.format(Messages.GET_ZMQ_KEYS.value, 'beeswarm_server'))
        zmq_server_key = zmq_server_key['public_key']
        # common section that goes for all types of drones
        drone_config = {
            'general': {
                'mode': drone.discriminator,
                'id': int(drone.id),
                'fetch_ip': False,
                'name': drone.name
            },
            'beeswarm_server': {
                'zmq_url': server_zmq_url,
                'zmq_command_url': server_zmq_command_url,
                'zmq_server_public': zmq_server_key,
                'zmq_own_public': zmq_keys['public_key'],
                'zmq_own_private': zmq_keys['private_key'],
            },
            'timecheck': {
                'enabled': True,
                'poll': 5,
                'ntp_pool': 'pool.ntp.org'
            }
        }

        if drone.discriminator == 'honeypot':
            drone_config['certificate_info'] = {
                'common_name': drone.cert_common_name,
                'country': drone.cert_country,
                'state': drone.cert_state,
                'locality': drone.cert_locality,
                'organization': drone.cert_organization,
                'organization_unit': drone.cert_organization_unit
            }
            drone_config['capabilities'] = {}
            for capability in drone.capabilities:
                users = {}
                for bait in capability.baits:
                    users[bait.username] = bait.password
                drone_config['capabilities'][capability.protocol] = {'port': capability.port,
                                                                     'enabled': True,
                                                                     'protocol_specific_data': json.loads(
                                                                         capability.protocol_specific_data),
                                                                     'users': users}
        elif drone.discriminator == 'client':
            drone_config['baits'] = {}
            for bait in drone.baits:
                _bait = {'server': bait.capability.honeypot.ip_address,
                         'port': bait.capability.port,
                         'honeypot_id': bait.capability.honeypot_id,
                         'username': bait.username,
                         'password': bait.password,
                         'active_range': bait.activation_range,
                         'sleep_interval': bait.sleep_interval,
                         'activation_probability': bait.activation_probability}
                if bait.capability.honeypot_id not in drone_config['baits']:
                    drone_config['baits'][bait.capability.honeypot_id] = {}
                assert bait.capability.protocol not in drone_config[
                    'baits'][bait.capability.honeypot_id]
                drone_config['baits'][bait.capability.honeypot_id][
                    bait.capability.protocol] = _bait
            if drone.bait_timings:
                drone_config['bait_timings'] = json.loads(drone.bait_timings)
            else:
                drone_config['bait_timings'] = {}

        return drone_config

    def _handle_command_get_db_stats(self):
        db_session = database_setup.get_session()
        database_stats = {
            'count_honeypots': db_session.query(Honeypot).count(),
            'count_clients': db_session.query(Client).count(),
            'count_sessions': db_session.query(Session).count(),
            'count_all_baits': db_session.query(BaitSession).count(),
            'count_all_attacks': db_session.query(Session).filter(Session.classification_id != 'bait_session')
            .filter(Session.classification_id != 'pending')
            .filter(Session.classification_id is not None).count(),
            'count_attack_type': {
                'http': self._get_num_attacks('http', db_session),
                'vnc': self._get_num_attacks('vnc', db_session),
                'ssh': self._get_num_attacks('ssh', db_session),
                'ftp': self._get_num_attacks('ftp', db_session),
                'https': self._get_num_attacks('https', db_session),
                'pop3': self._get_num_attacks('pop3', db_session),
                'pop3s': self._get_num_attacks('pop3s', db_session),
                'smtp': self._get_num_attacks('smtp', db_session),
                'telnet': self._get_num_attacks('telnet', db_session),
            },
            'baits': {
                'successful': db_session.query(BaitSession).filter(BaitSession.did_login).count(),
                'failed': db_session.query(BaitSession).filter(not BaitSession.did_login).count(),

            }
        }
        return database_stats

    def _get_num_attacks(self, protocol, db_session):
        return db_session.query(Session).filter(Session.classification_id != 'bait_session') \
            .filter(Session.classification_id is not None) \
            .filter(Session.protocol == protocol).count()

    def _handle_command_get_sessions(self, _type):
        db_session = database_setup.get_session()
        # the database_setup will not get hit until we start iterating the
        # query object
        query_iterators = {
            Messages.GET_SESSIONS_ALL.value: db_session.query(Session),
            Messages.GET_SESSIONS_BAIT.value: db_session.query(BaitSession),
            Messages.GET_SESSIONS_ATTACKS.value: db_session.query(Session).filter(
                Session.classification_id != 'bait_session')
        }

        if _type not in query_iterators:
            logger.warning(
                'Query for sessions with unknown type: {0}'.format(_type))
            return []

        # select which iterator to use
        entries = query_iterators[_type].order_by(desc(Session.timestamp))
        rows = []

        for session in entries:
            rows.append(session.to_dict())

        return rows

    def _handle_command_get_credentials(self, session_id):
        db_session = database_setup.get_session()

        credentials = db_session.query(Authentication).filter(
            Authentication.session_id == session_id)
        return_rows = []
        for credential in credentials:
            return_rows.append(credential.to_dict())
        return return_rows

    def _handle_command_get_transcript(self, session_id):
        db_session = database_setup.get_session()

        transcripts = db_session.query(Transcript).filter(
            Transcript.session_id == session_id)
        return_rows = []
        for transcript in transcripts:
            row = {'time': transcript.timestamp.strftime('%Y-%m-%d %H:%M:%S'), 'direction': transcript.direction,
                   'data': transcript.data}
            return_rows.append(row)
        return return_rows

    def _handle_command_get_drone_list(self, drone_type):
        db_session = database_setup.get_session()
        if drone_type == 'all':
            drones = db_session.query(Drone).all()
        elif drone_type == 'unassigned':
            drones = db_session.query(Drone).filter(
                Drone.discriminator == None)
        else:
            drones = db_session.query(Drone).filter(
                Drone.discriminator == drone_type)

        drone_list = []
        for drone in drones:
            drone_list.append(drone.to_dict())
        return drone_list

    def _handle_command_config_drone(self, data):
        drone_id, config = data.split(' ', 1)
        config = json.loads(config)

        db_session = database_setup.get_session()
        drone = db_session.query(Drone).filter(Drone.id == drone_id).one_or_none()
        if not drone:
            self.databaseRequests.send('{0} {1}'.format(Messages.FAIL.value, 'Drone with id {0} could not '
                                                                             'found'.format(drone_id)))
        elif config['mode'] == 'honeypot':
            # it is a honeypot
            self._config_honeypot(drone, db_session, config)
            self.databaseRequests.send('{0} {1}'.format(Messages.OK.value, {}))
        elif config['mode'] == 'client':
            # it is a client
            self._config_client(drone, db_session, config)
            self.databaseRequests.send('{0} {1}'.format(Messages.OK.value, {}))

        else:
            logger.error(
                'Could not detect mode for drone config, drone id: {0}'.format(drone_id))
            self.databaseRequests.send('{0} {1}'.format(Messages.FAIL.value, 'Malformed data in drone config data.'
                                                                             'Drone id: {0}'.format(drone_id)))

    def _config_honeypot(self, drone, db_session, config):
        if drone.discriminator != 'honeypot':
            # meh, better way do do this?
            drone_id = drone.id
            ip_address = drone.ip_address
            db_session.delete(drone)
            db_session.commit()
            drone = Honeypot(id=drone_id)
            drone.ip_address = ip_address
            db_session.add(drone)
            db_session.commit()

        # common properties
        drone.name = config['name']

        # certificate information
        drone.cert_common_name = config['certificate']['common_name']
        drone.cert_country = config['certificate']['country']
        drone.cert_state = config['certificate']['state']
        drone.cert_locality = config['certificate']['locality']
        drone.cert_organization = config['certificate']['organization']
        drone.cert_organization_unit = config[
            'certificate']['organization_unit']

        # add capabilities
        drone.capabilities = []
        for protocol_name, protocol_config in config['capabilities'].items():
            if 'protocol_specific_data' in protocol_config:
                protocol_specific_data = protocol_config[
                    'protocol_specific_data']
            else:
                protocol_specific_data = {}
            drone.add_capability(protocol_name, protocol_config[
                                 'port'], protocol_specific_data)

        db_session.add(drone)
        db_session.commit()
        self._handle_command_drone_config_changed(drone.id)

    def _config_client(self, drone, db_session, config):
        if drone.discriminator != 'client':
            # meh, better way do do this?
            # TODO: this cascade delete sessions, find a way to maintain
            # sessions for deleted drones.
            ip_address = drone.ip_address
            drone_id = drone.id
            db_session.delete(drone)
            db_session.commit()
            drone = Client(id=drone_id)
            drone.ip_address = ip_address
            db_session.add(drone)
            db_session.commit()
        drone.bait_timings = json.dumps(config['bait_timings'])
        drone.name = config['name']
        db_session.add(drone)
        db_session.commit()
        self._handle_command_drone_config_changed(drone.id)

    def _db_maintenance(self):
        logger.debug('Doing database maintenance')
        bait_session_retain_days = int(self.send_config_request('{0} {1}'.format(Messages.GET_CONFIG_ITEM.value,
                                                                                 'bait_session_retain')))
        bait_retain = datetime.utcnow() - timedelta(days=bait_session_retain_days)
        malicious_session_retain_days = int(self.send_config_request('{0} {1}'.format(Messages.GET_CONFIG_ITEM.value,
                                                                                      'malicious_session_retain')))
        malicious_retain = datetime.utcnow() - timedelta(days=malicious_session_retain_days)

        db_session = database_setup.get_session()

        malicious_deleted_count = db_session.query(Session).filter(Session.classification_id != 'bait_session') \
            .filter(Session.timestamp < malicious_retain).delete()

        bait_deleted_count = db_session.query(Session).filter(Session.classification_id == 'bait_session') \
            .filter(Session.timestamp < bait_retain).delete()
        db_session.commit()

        logger.info('Database maintenance finished. Deleted {0} bait_sessions and {1} malicious sessions)'
                    .format(bait_deleted_count, malicious_deleted_count))
beeswarm/server/db/entities.py
class Client(Drone):
    __tablename__ = 'client'
    __mapper_args__ = {'polymorphic_identity': 'client'}

    id = Column(String, ForeignKey('drone.id'), primary_key=True)
    bait_sessions = relationship("BaitSession", cascade="all, delete-orphan", backref='client')

    baits = relationship('DroneEdge', cascade='all, delete-orphan', backref='client')
    bait_timings = Column(String)

    def add_bait(self, capability, activation_range, sleep_interval, activation_probability, username, password):
        bait = DroneEdge(capability=capability, activation_range=activation_range, sleep_interval=sleep_interval,
                         activation_probability=activation_probability, username=username, password=password)
        bait.client_id = self.id
        self.baits.append(bait)
beeswarm/drones/client/models/session.py
class BaitSession(BaseSession):
    client_id = ''

    def __init__(self, protocol, destination_ip, destination_port, honeypot_id):
        super(BaitSession, self).__init__(protocol, destination_ip=destination_ip,
                                          destination_port=destination_port)

        assert BaitSession.client_id

        self.client_id = BaitSession.client_id
        self.honeypot_id = honeypot_id

        self.did_connect = False
        self.did_login = False
        self.alldone = False
        self.did_complete = False
        self.protocol_data = {}

    def to_dict(self):
        return vars(self)

    def end_session(self):
        super(BaitSession, self).end_session(Messages.SESSION_CLIENT.value)
beeswarm/drones/honeypot/models/session.py
class Session(BaseSession):
    authenticator = None
    default_timeout = 25
    honeypot_id = None

    def __init__(self, source_ip, source_port, protocol, users, destination_port=None, destination_ip=None):

        super(Session, self).__init__(protocol, source_ip, source_port, destination_ip, destination_port)

        self.connected = True
        self.authenticated = False
        self.honeypot_id = Session.honeypot_id
        self.users = users

        # for session specific volatile data (will not get logged)
        self.vdata = {}
        self.last_activity = datetime.utcnow()

        self.send_log(Messages.SESSION_PART_HONEYPOT_SESSION_START.value, self.to_dict())

    def activity(self):
        self.last_activity = datetime.utcnow()

    def is_connected(self):
        return self.connected

    def try_auth(self, _type, **kwargs):
        authenticated = False
        if _type == 'plaintext':
            if kwargs.get('username') in self.users:
                if self.users[kwargs.get('username')] == kwargs.get('password'):
                    authenticated = True

        elif _type == 'cram_md5':
            def encode_cram_md5(challenge, user, password):
                response = user + ' ' + hmac.HMAC(password, challenge).hexdigest()
                return response

            if kwargs.get('username') in self.users:
                uname = kwargs.get('username')
                digest = kwargs.get('digest')
                s_pass = self.users[uname]
                challenge = kwargs.get('challenge')
                ideal_response = encode_cram_md5(challenge, uname, s_pass)
                _, ideal_digest = ideal_response.split()
                if ideal_digest == digest:
                    authenticated = True
        elif _type == 'des_challenge':
            challenge = kwargs.get('challenge')
            response = kwargs.get('response')
            for valid_password in self.users.values():
                aligned_password = (valid_password + '\0' * 8)[:8]
                des = RFBDes(aligned_password)
                expected_response = des.encrypt(challenge)
                if response == expected_response:
                    authenticated = True
                    kwargs['password'] = aligned_password
                    break
        else:
            assert False

        if authenticated:
            self.authenticated = True
            self.add_auth_attempt(_type, True, **kwargs)
        else:
            self.add_auth_attempt(_type, False, **kwargs)

        if _type == 'des_challenge':
            kwargs['challenge'] = kwargs.get('challenge').encode('hex')
            kwargs['response'] = kwargs.get('response').encode('hex')

        self.send_log(Messages.SESSION_PART_HONEYPOT_AUTH.value, self.login_attempts[-1])
        logger.debug('{0} authentication attempt from {1}:{2}. Credentials: {3}'.format(self.protocol, self.source_ip,
                                                                                        self.source_port,
                                                                                        json.dumps(kwargs)))
        return authenticated

    def end_session(self):
        super(Session, self).end_session(Messages.SESSION_HONEYPOT.value)
beeswarm/shared/helpers.py
def send_zmq_request_socket(socket, request):
    socket.send(request)
    result = socket.recv()
    status, data = result.split(' ', 1)
    if status != Messages.OK.value:
        logger.error('Receiving actor on socket {0} failed to respond properly. The request was: {1}'.format(socket,
                                                                                                             request))
        assert False
    else:
        if data.startswith('{') or data.startswith('['):
            return json.loads(result.split(' ', 1)[1])
        else:
            return data
beeswarm/server/db/database_setup.py
def setup_db(connection_string):
def clear_db():
def get_session():
beeswarm/server/db/entities.py
class Session(Base):
    __tablename__ = 'session'
    discriminator = Column('type', String(50))
    __mapper_args__ = {'polymorphic_on': discriminator}

    id = Column(String, primary_key=True)
    received = Column(DateTime)
    timestamp = Column(DateTime, index=True)
    protocol = Column(String)
    authentication = relationship('Authentication', cascade="all, delete-orphan")
    source_ip = Column(String)
    source_port = Column(Integer)
    session_data = relationship('SessionData', cascade="all, delete-orphan", backref='Session')
    transcript = relationship('Transcript', cascade="all, delete-orphan", backref='Session')
    destination_ip = Column(String)
    destination_port = Column(Integer)
    honeypot_id = Column(String, ForeignKey('honeypot.id'))
    drone = relationship(Drone, primaryjoin=honeypot_id==Drone.id, foreign_keys=honeypot_id, backref='session')
    classification_id = Column(String, ForeignKey('classification.type'), nullable=False, default='pending')
    classification = relationship('Classification')

    # for display purposes
    def to_dict(self):

        auth_attempts = []
        for attempt in self.authentication:
            auth_attempts.append(attempt.to_dict())
        classification = self.classification_id.replace('_', ' ').capitalize()
        result = {'time': self.timestamp.strftime('%Y-%m-%d %H:%M:%S'),
                  'protocol': self.protocol,
                  'ip_address': self.source_ip,
                  'honeypot_id': self.honeypot_id,
                  'classification': classification,
                  'id': self.id,
                  'drone_name': '',
                  }
                  # drone can be None when unit testing
        if self.drone:
            result['drone_name'] = self.drone.name

        return result
beeswarm/shared/message_enum.py
class Messages(Enum):
    STOP = 'STOP'
    START = 'START'
    CONFIG = 'CONFIG'
    # dump of all configuration elements known to the sender
    CONFIG_FULL = 'CONFIG_FULL'
    # mapping between clients, honeypots, capabilities and bait users
    CONFIG_ARCHITECTURE = 'CONFIG_ARCHITECTURE'
    BROADCAST = 'BROADCAST'

    OK = 'OK'
    FAIL = 'FAIL'
    # KEY DRONE_ID DRONE_PRIVATE_KEY
    KEY = 'KEY'
    # CERT DRONE_ID DRONE_CERT
    CERT = 'CERT'

    # All data on an entire session, this also indicates session end.
    SESSION_HONEYPOT = 'SESSION_HONEYPOT'
    # Network connect to a honeypot, data format is as in SESSION_HONEYPOT, but not all fields has data.
    SESSION_PART_HONEYPOT_SESSION_START = 'SESSION_PART_HONEYPOT_SESSION_START'
    # Authentication attempt, format is as in self.login_attempts list SESSION_HONEYPOT
    SESSION_PART_HONEYPOT_AUTH = 'SESSION_PART_HONEYPOT_AUTH'
    SESSION_CLIENT = 'SESSION_CLIENT'

    SET_CONFIG_ITEM = 'SET'
    GET_CONFIG_ITEM = 'GET'
    GET_ZMQ_KEYS = 'GET_ZMQ_KEYS'
    DELETE_ZMQ_KEYS = 'DELETE_ZMQ_KEYS'
    PING = 'PING'
    PONG = 'PING'
    IP = 'IP'

    DRONE_WANT_CONFIG = 'DRONE_WANT_CONFIG'
    DRONE_CONFIG = 'DRONE_CONFIG'
    # ID USERNAME PASSWORD
    BAIT_USER_ADD = 'BAIT_USER_ADD'
    BAIT_USER_DELETE = 'BAIT_USER_DELETE'
    GET_BAIT_USERS = 'GET_BAIT_USERS'
    DRONE_DELETE = 'DRONE_DELETE'
    DRONE_ADD = 'DRONE_ADD'

    # Session was deleted because it was matched with an existing session
    # This happens when a bait session is matched with a honeypot session.
    # ID_OF_DELETED_SESSION ID_OF_THE_MERGED_SESSION
    DELETED_DUE_TO_MERGE = "DELETED_DUE_TO_MERGE"

    # A classified session is transmitted
    # Parameters is data in json.
    SESSION = "SESSION"

    # Database requests
    GET_DB_STATS = 'GET_DB_STATS'
    # TODO: Following three should have to/from params to facilitate pagination
    GET_SESSIONS_ALL = 'GET_SESSIONS_ALL'
    GET_SESSIONS_BAIT = 'GET_SESSIONS_BAIT'
    GET_SESSIONS_ATTACKS = 'GET_SESSIONS_ATTACKS'

    # Param: session id
    GET_SESSION_TRANSCRIPT = 'GET_SESSION_TRANSCRIPT'
    # Param: session id
    GET_SESSION_CREDENTIALS = 'GET_SESSION_CREDENTIALS'
    # Param dronetype
    GET_DRONE_LIST = 'GET_DRONE_LIST'
    # Param: clientId config
    CONFIG_DRONE = 'CONFIG_DRONE'

    # Param: clientId config
    PING_ALL_DRONES = 'PING_ALL_DRONES'
beeswarm/server/misc/config_actor.py
class ConfigActor(Greenlet):
    def __init__(self, config_file, work_dir):
        Greenlet.__init__(self)
        self.config_file = os.path.join(work_dir, config_file)
        if not os.path.exists(self.config_file):
            self.config = {}
            self._save_config_file()
        self.config = json.load(open(self.config_file, 'r'))
        self.work_dir = work_dir

        context = beeswarm.shared.zmq_context
        self.config_commands = context.socket(zmq.REP)
        self.enabled = True

    def stop(self):
        self.enabled = False

        if self.config_commands:
            self.config_commands.close()

    def _run(self):
        self.config_commands.bind(SocketNames.CONFIG_COMMANDS.value)

        poller = zmq.Poller()
        poller.register(self.config_commands, zmq.POLLIN)
        while self.enabled:
            socks = dict(poller.poll(500))
            if self.config_commands in socks and socks[self.config_commands] == zmq.POLLIN:
                self._handle_commands()

    def _handle_commands(self):
        msg = self.config_commands.recv()

        if ' ' in msg:
            cmd, data = msg.split(' ', 1)
        else:
            cmd = msg
        logger.debug('Received command: {0}'.format(cmd))

        if cmd == Messages.SET_CONFIG_ITEM.value:
            self._handle_command_set(data)
            self.config_commands.send('{0} {1}'.format(Messages.OK.value, '{}'))
        elif cmd == Messages.GET_CONFIG_ITEM.value:
            value = self._handle_command_get(data)
            self.config_commands.send('{0} {1}'.format(Messages.OK.value, value))
        elif cmd == Messages.GET_ZMQ_KEYS.value:
            self._handle_command_getkeys(data)
        elif cmd == Messages.DELETE_ZMQ_KEYS.value:
            self._remove_zmq_keys(data)
            self.config_commands.send('{0} {1}'.format(Messages.OK.value, '{}'))
        else:
            logger.error('Unknown command received: {0}'.format(cmd))
            self.config_commands.send(Messages.FAIL.value)

    def _handle_command_set(self, data):
        new_config = json.loads(data)
        self.config.update(new_config)
        self._save_config_file()

    def _handle_command_get(self, data):
        # example: 'network,host' will lookup self.config['network']['host']
        keys = data.split(',')
        value = self._retrieve_nested_config(keys, self.config)
        return value

    def _retrieve_nested_config(self, keys, dict):
        if keys[0] in dict:
            if len(keys) == 1:
                return dict[keys[0]]
            else:
                return self._retrieve_nested_config(keys[1:], dict[keys[0]])

    def _handle_command_getkeys(self, name):
        private_key, publickey = self._get_zmq_keys(name)
        self.config_commands.send(Messages.OK.value + ' ' + json.dumps({'public_key': publickey,
                                                                        'private_key': private_key}))

    def _save_config_file(self):
        with open(self.config_file, 'w+') as config_file:
            config_file.write(json.dumps(self.config, indent=4))

    def _get_zmq_keys(self, id):
        cert_path = os.path.join(self.work_dir, 'certificates')
        public_keys = os.path.join(cert_path, 'public_keys')
        private_keys = os.path.join(cert_path, 'private_keys')
        public_key_path = os.path.join(public_keys, '{0}.pub'.format(id))
        private_key_path = os.path.join(private_keys, '{0}.pri'.format(id))

        if not os.path.isfile(public_key_path) or not os.path.isfile(private_key_path):
            logging.debug('Generating ZMQ keys for: {0}.'.format(id))
            for _path in [cert_path, public_keys, private_keys]:
                if not os.path.isdir(_path):
                    os.mkdir(_path)

            tmp_key_dir = tempfile.mkdtemp()
            try:
                public_key, private_key = create_certificates(tmp_key_dir, id)
                # the final location for keys
                shutil.move(public_key, public_key_path)
                shutil.move(private_key, private_key_path)
            finally:
                shutil.rmtree(tmp_key_dir)

        # return copy of keys
        return open(private_key_path, "r").readlines(), open(public_key_path, "r").readlines()

    def _remove_zmq_keys(self, id):
        cert_path = os.path.join(self.work_dir, 'certificates')
        public_keys = os.path.join(cert_path, 'public_keys')
        private_keys = os.path.join(cert_path, 'private_keys')
        public_key_path = os.path.join(public_keys, '{0}.pub'.format(id))
        private_key_path = os.path.join(private_keys, '{0}.pri'.format(id))

        for _file in [public_key_path, private_key_path]:
            if os.path.isfile(_file):
                os.remove(_file)
beeswarm/shared/socket_enum.py
class SocketNames(Enum):
    #### Sockets used on server ####
    # All data received from drones will be published on this socket
    DRONE_DATA = 'inproc://droneData'
    # After sessions has been classified they will get retransmitted on this socket.
    # TODO: Does not actually happen yet
    PROCESSED_SESSIONS = 'inproc://processedSessionPublisher'
    # Request / Reply to config actor
    CONFIG_COMMANDS = 'inproc://configCommands'
    # Data sent on this socket will be retransmitted to the correct drone, the data must be prefixed with
    # the id of the drone.
    DRONE_COMMANDS = 'inproc://droneCommands'


    # Requests to and from the databsae
    DATABASE_REQUESTS = 'inproc://databaseRequests'

    #### Sockets used on drones ####
    # Drone commands received from the server will be retransmitted  on this socket.
    SERVER_COMMANDS = 'inproc://serverCommands'
    # All messages transmitted on this socket will get retransmitted to the server
    SERVER_RELAY = 'inproc://serverRelay'


Question: import unittest
import uuid
import json
import tempfile
import os
import zmq.green as zmq
import gevent
import beeswarm.shared
from datetime import datetime
from beeswarm.server.db import database_setup
from beeswarm.server.db.entities import Honeypot, Client
from beeswarm.server.db.entities import Session
from beeswarm.drones.honeypot.models.session import Session as HoneypotSession
from beeswarm.shared.socket_enum import SocketNames
from beeswarm.shared.message_enum import Messages
from beeswarm.server.db.database_actor import DatabaseActor
from beeswarm.drones.client.models.session import BaitSession
from beeswarm.server.misc.config_actor import ConfigActor
from beeswarm.shared.helpers import send_zmq_request_socket
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.






class ClassifierTests(unittest.TestCase):
    def setUp(self):
        beeswarm.shared.zmq_context = zmq.Context()
        fd, self.db_file = tempfile.mkstemp()
        os.close(fd)
        connection_string = 'sqlite:///{0}'.format(self.db_file)
        os.remove(self.db_file)
        database_setup.setup_db(connection_string)

    def tearDown(self):
        if os.path.isfile(self.db_file):
            os.remove(self.db_file)

    def test_matching_quick_succession(self):
        """
        Tests that attack sessions coming in quick succession are classified correctly.
        This test relates to issue #218
        """

        honeypot_id = 1
        honeypot = Honeypot(id=honeypot_id)

        db_session = database_setup.get_session()
        db_session.add(honeypot)
        db_session.commit()

        drone_data_socket = beeswarm.shared.zmq_context.socket(zmq.PUB)
        drone_data_socket.bind(SocketNames.DRONE_DATA.value)

        # startup session database
        database_actor = DatabaseActor(999, delay_seconds=2)
        database_actor.start()
        gevent.sleep(1)

        for x in xrange(0, 100):
            honeypot_session = HoneypotSession(source_ip='192.168.100.22', source_port=52311, protocol='pop3', users={},
                                               destination_port=110)
            honeypot_session.add_auth_attempt('plaintext', True, username='james', password='bond')
            honeypot_session.honeypot_id = honeypot_id
            drone_data_socket.send('{0} {1} {2}'.format(Messages.SESSION_HONEYPOT.value, honeypot_id,
                                                            json.dumps(honeypot_session.to_dict(), default=json_default,
                                                            ensure_ascii=False)))

        gevent.sleep(1)
        database_actor_request_socket = beeswarm.shared.zmq_context.socket(zmq.REQ)
        database_actor_request_socket.connect(SocketNames.DATABASE_REQUESTS.value)
        sessions = send_zmq_request_socket(database_actor_request_socket, '{0}'.format(Messages.GET_SESSIONS_ALL.value))

        for session in sessions:
            self.assertEqual(session['classification'], 'Bruteforce')

        self.assertEqual(len(sessions), 100)

    def test_bait_classification_honeypot_first(self):
        """
        Tests that bait sessions are paired correctly with their honeypot counter parts when honeypot message arrives
        first.
        """

        self.populate_bait(True)
        # TODO: Use message request - not orm access!
        db_session = database_setup.get_session()
        sessions = db_session.query(Session).all()
        for session in sessions:
            self.assertEqual(session.classification_id, 'bait_session')

        self.assertEqual(len(sessions), 1)

    def test_bait_classification_client_first(self):
        """
        Tests that bait sessions are paired correctly with their honeypot counter parts when client message arrives
        first.
        """

        self.populate_bait(False)
        # TODO: Use message request - not orm access!
        db_session = database_setup.get_session()
        sessions = db_session.query(Session).all()
        for session in sessions:
            self.assertEqual(session.classification_id, 'bait_session')

        self.assertEqual(len(sessions), 1)

    def populate_bait(self, honeypot_first):
        honeypot_id = 1
        client_id = 2
        honeypot = Honeypot(id=honeypot_id)
        client = Client(id=client_id)

        db_session = database_setup.get_session()
        db_session.add(honeypot)
        db_session.add(client)
        db_session.commit()

        drone_data_socket = beeswarm.shared.zmq_context.socket(zmq.PUB)
        drone_data_socket.bind(SocketNames.DRONE_DATA.value)

        fd, config_file = tempfile.mkstemp()
        os.close(fd)
        os.remove(config_file)
        # persistence actor needs to communicate with on config REQ/REP socket

Answer:         config_actor = ConfigActor(config_file, '')