Extracted process telem to hooks.py

This commit is contained in:
Shay Nehmad 2019-08-19 14:43:23 +03:00
parent 757af1c6af
commit a6789a53b2
2 changed files with 18 additions and 10 deletions

View File

@ -9,7 +9,7 @@ from flask import request
from monkey_island.cc.auth import jwt_required from monkey_island.cc.auth import jwt_required
from monkey_island.cc.database import mongo from monkey_island.cc.database import mongo
from monkey_island.cc.services.node import NodeService from monkey_island.cc.services.node import NodeService
from monkey_island.cc.services.telemetry.processing.hooks import TELEMETRY_CATEGORY_TO_PROCESSING_FUNC from monkey_island.cc.services.telemetry.processing.hooks import process_telemetry
from monkey_island.cc.models.monkey import Monkey from monkey_island.cc.models.monkey import Monkey
__author__ = 'Barak' __author__ = 'Barak'
@ -48,16 +48,9 @@ class Telemetry(flask_restful.Resource):
Monkey.get_single_monkey_by_guid(telemetry_json['monkey_guid']).renew_ttl() Monkey.get_single_monkey_by_guid(telemetry_json['monkey_guid']).renew_ttl()
monkey = NodeService.get_monkey_by_guid(telemetry_json['monkey_guid']) monkey = NodeService.get_monkey_by_guid(telemetry_json['monkey_guid'])
NodeService.update_monkey_modify_time(monkey["_id"])
try: process_telemetry(telemetry_json)
NodeService.update_monkey_modify_time(monkey["_id"])
telem_category = telemetry_json.get('telem_category')
if telem_category in TELEMETRY_CATEGORY_TO_PROCESSING_FUNC:
TELEMETRY_CATEGORY_TO_PROCESSING_FUNC[telem_category](telemetry_json)
else:
logger.info('Got unknown type of telemetry: %s' % telem_category)
except Exception as ex:
logger.error("Exception caught while processing telemetry. Info: {}".format(ex.message), exc_info=True)
telem_id = mongo.db.telemetry.insert(telemetry_json) telem_id = mongo.db.telemetry.insert(telemetry_json)
return mongo.db.telemetry.find_one_or_404({"_id": telem_id}) return mongo.db.telemetry.find_one_or_404({"_id": telem_id})

View File

@ -1,5 +1,9 @@
import logging
from monkey_island.cc.services.telemetry.processing import * from monkey_island.cc.services.telemetry.processing import *
logger = logging.getLogger(__name__)
TELEMETRY_CATEGORY_TO_PROCESSING_FUNC = \ TELEMETRY_CATEGORY_TO_PROCESSING_FUNC = \
{ {
'tunnel': process_tunnel_telemetry, 'tunnel': process_tunnel_telemetry,
@ -12,3 +16,14 @@ TELEMETRY_CATEGORY_TO_PROCESSING_FUNC = \
'trace': lambda *args, **kwargs: None, 'trace': lambda *args, **kwargs: None,
'attack': lambda *args, **kwargs: None, 'attack': lambda *args, **kwargs: None,
} }
def process_telemetry(telemetry_json):
try:
telem_category = telemetry_json.get('telem_category')
if telem_category in TELEMETRY_CATEGORY_TO_PROCESSING_FUNC:
TELEMETRY_CATEGORY_TO_PROCESSING_FUNC[telem_category](telemetry_json)
else:
logger.info('Got unknown type of telemetry: %s' % telem_category)
except Exception as ex:
logger.error("Exception caught while processing telemetry. Info: {}".format(ex.message), exc_info=True)