fixed crash when expected entity could not be found

fix #233
honeynet · Dec 6, 2016 · 8f77dd5 · 8f77dd5
1 parent 4a85d03
commit 8f77dd5
Showing 1 changed file with 12 additions and 6 deletions.
diff --git a/beeswarm/server/db/database_actor.py b/beeswarm/server/db/database_actor.py
@@ -203,7 +203,7 @@ def stop(self):
 
     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()
+        drone = db_session.query(Drone).filter(Drone.id == drone_id).one_or_none()
         if drone:
             drone.last_activity = datetime.now()
             db_session.add(drone)
@@ -229,7 +229,7 @@ def _handle_message_ip(self, topic, drone_id, 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()
+        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
@@ -248,7 +248,7 @@ def _handle_cert_message(self, topic, drone_id, data):
         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 = db_session.query(Drone).filter(Drone.id == drone_id).one_or_none()
         if drone:
             drone.cert_digest = digest
             db_session.add(drone)
@@ -277,7 +277,10 @@ def persist_session(self, session_type, session_json):
 
         assert data['honeypot_id'] is not None
         _honeypot = db_session.query(Honeypot).filter(
-            Honeypot.id == data['honeypot_id']).one()
+            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']:
@@ -297,7 +300,10 @@ def persist_session(self, session_type, session_json):
                 return
             session = BaitSession()
             client = db_session.query(Client).filter(
-                Client.id == data['client_id']).one()
+                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']
@@ -802,7 +808,7 @@ def _handle_command_config_drone(self, data):
         config = json.loads(config)
 
         db_session = database_setup.get_session()
-        drone = db_session.query(Drone).filter(Drone.id == drone_id).one()
+        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)))