diff --git a/monkey/monkey_island/cc/app.py b/monkey/monkey_island/cc/app.py index be2430dda..4c5c64ec8 100644 --- a/monkey/monkey_island/cc/app.py +++ b/monkey/monkey_island/cc/app.py @@ -1,8 +1,11 @@ +import json +import logging import os +import time import uuid import flask_restful -from flask import Flask, send_from_directory, Response +from flask import Flask, send_from_directory, Response, g from werkzeug.exceptions import NotFound from monkey_island.cc.auth import init_jwt @@ -41,6 +44,8 @@ __author__ = 'Barak' HOME_FILE = 'index.html' +logger = logging.getLogger(__name__) + def serve_static_file(static_path): if static_path.startswith('api/'): @@ -119,6 +124,25 @@ def init_api_resources(api): api.add_resource(LogTest, '/api/test/log') +def init_app_execution_time_calc(app): + @app.before_request + def before_request(): + g.start = time.time() + logger.debug("inb4 next request?") + + @app.after_request + def after_request(response): + diff_in_ms = int((time.time() - g.start)*1000) + logger.debug(f"After request") + if response.response: + if response.content_type == "application/json": + response_data = json.loads(response.get_data()) + if isinstance(response_data, dict): + response_data["execution_time_ms"] = diff_in_ms + response.set_data(json.dumps(response_data)) + return response + + def init_app(mongo_url): app = Flask(__name__) @@ -128,6 +152,7 @@ def init_app(mongo_url): init_app_config(app, mongo_url) init_app_services(app) init_app_url_rules(app) + init_app_execution_time_calc(app) init_api_resources(api) return app