Improve PoC editor

- Change internal tree representation for experiments
- Store python backend logic as files without folder
- Existing url_queue.txt files can be editor through web UI

Author: GJFR

bci/evaluations/custom/custom_evaluation.py

@@ -20,19 +20,29 @@ def __init__(self):
 
     @staticmethod
     def initialize_dir_tree() -> dict:
+        """
+        Initializes directory tree of experiments.
+        """
         path = Global.custom_page_folder
-
-        def path_to_dict(path):
-            if os.path.isdir(path):
-                return {
-                    sub_folder: path_to_dict(os.path.join(path, sub_folder))
-                    for sub_folder in os.listdir(path)
-                    if sub_folder != 'url_queue.txt'
-                }
+        dir_tree = {}
+
+        def set_nested_value(d: dict, keys: list[str], value: dict):
+            nested_dict = d
+            for key in keys[:-1]:
+                nested_dict = nested_dict[key]
+            nested_dict[keys[-1]] = value
+
+        for root, dirs, files in os.walk(path):
+            # Remove base path from root
+            root = root[len(path):]
+            keys = root.split('/')[1:]
+            subdir_tree = {dir: {} for dir in dirs} | {file: None for file in files}
+            if root:
+                set_nested_value(dir_tree, keys, subdir_tree)
             else:
-                return os.path.basename(path)
+                dir_tree = subdir_tree
 
-        return path_to_dict(path)
+        return dir_tree
 
     @staticmethod
     def initialize_tests_and_url_queues(dir_tree: dict) -> dict:
@@ -70,11 +80,14 @@ def __get_url_queue(project: str, project_path: str, experiment: str) -> list[st
         raise AttributeError(f"Could not infer url queue for experiment '{experiment}' in project '{project}'")
 
     @staticmethod
-    def is_runnable_experiment(project: str, poc: str, dir_tree: dict) -> bool:
+    def is_runnable_experiment(project: str, poc: str, dir_tree: dict[str,dict]) -> bool:
         domains = dir_tree[project][poc]
-        if not (poc_main_path := [paths for domain, paths in domains.items() if 'main' in paths]):
+        # Should have exactly one main folder
+        main_paths = [paths for paths in domains.values() if paths is not None and 'main' in paths.keys()]
+        if len(main_paths) != 1:
             return False
-        if 'index.html' not in poc_main_path[0]['main'].keys():
+        # Main should have index.html
+        if 'index.html' not in main_paths[0]['main'].keys():
             return False
         return True
 
@@ -194,7 +207,7 @@ def get_default_file_content(file_type: str) -> str:
         if file_type != 'py':
             return ''
 
-        with open('experiments/pages/Support/PythonServer/a.test/py-server/index.py', 'r') as file:
+        with open('experiments/pages/Support/PythonServer/a.test/server.py', 'r') as file:
             template_content = file.read()
 
         return template_content

bci/web/blueprints/api.py

@@ -16,8 +16,10 @@
 THREAD = None
 
 
-def start_thread(func, args=None) -> bool:
+def __start_thread(func, args=None) -> bool:
     global THREAD
+    if args is None:
+        args = []
     if THREAD and THREAD.is_alive():
         return False
     else:
@@ -26,7 +28,6 @@ def start_thread(func, args=None) -> bool:
         return True
 
 
-
 @api.before_request
 def check_readiness():
     if not bci_api.is_ready():
@@ -55,9 +56,15 @@ def add_headers(response):
 
 @api.route('/evaluation/start/', methods=['POST'])
 def start_evaluation():
+    if request.json is None:
+        return {
+            'status': 'NOK',
+            'msg': "No evaluation parameters found"
+        }
+
     data = request.json.copy()
     params = evaluation_factory(data)
-    if start_thread(bci_api.run, args=[params]):
+    if __start_thread(bci_api.run, args=[params]):
         return {
             'status': 'OK'
         }
@@ -69,6 +76,12 @@ def start_evaluation():
 
 @api.route('/evaluation/stop/', methods=['POST'])
 def stop_evaluation():
+    if request.json is None:
+        return {
+            'status': 'NOK',
+            'msg': "No stop parameters found"
+        }
+
     data = request.json.copy()
     forcefully = data.get('forcefully', False)
     if forcefully:
@@ -140,6 +153,12 @@ def log():
 
 @api.route('/data/', methods=['PUT'])
 def data_source():
+    if request.json is None:
+        return {
+            'status': 'NOK',
+            'msg': "No data parameters found"
+        }
+
     params = request.json.copy()
     revision_data, version_data = bci_api.get_data_sources(params)
     if revision_data or version_data:
@@ -172,30 +191,41 @@ def poc(project: str, poc: str):
     }
 
 
-@api.route('/poc/<string:project>/<string:poc>/<string:domain>/<string:path>/<string:file>/', methods=['GET'])
-def poc_file(project: str, poc: str, domain: str, path: str, file: str):
-    return {
-        'status': 'OK',
-        'content': bci_api.get_poc_file(project, poc, domain, path, file)
-    }
-
-
-@api.route('/poc/<string:project>/<string:poc>/<string:domain>/<string:path>/<string:file>/', methods=['POST'])
-def update_poc_file(project: str, poc: str, domain: str, path: str, file: str):
-    data = request.json.copy()
-    success = bci_api.update_poc_file(project, poc, domain, path, file, data['content'])
-    if success:
+@api.route('/poc/<string:project>/<string:poc>/<string:file>/', methods=['GET', 'POST'])
+def get_poc_file_content(project: str, poc: str, file: str):
+    domain = request.args.get('domain', '')
+    path = request.args.get('path', '')
+    if request.method == 'GET':
         return {
-            'status': 'OK'
-        }
-    else :
-        return {
-            'status': 'NOK'
+            'status': 'OK',
+            'content': bci_api.get_poc_file(project, poc, domain, path, file)
         }
+    else:
+        if not request.json:
+            return {
+                'status': 'NOK',
+                'msg': 'No content to update file with'
+            }
+        data = request.json.copy()
+        success = bci_api.update_poc_file(project, poc, domain, path, file, data['content'])
+        if success:
+            return {
+                'status': 'OK'
+            }
+        else :
+            return {
+                'status': 'NOK'
+            }
 
 
 @api.route('/poc/<string:project>/<string:poc>/', methods=['POST'])
 def add_page(project: str, poc: str):
+    if request.json is None:
+        return {
+            'status': 'NOK',
+            'msg': "No page parameters found"
+        }
+
     data = request.json.copy()
     success = bci_api.add_page(project, poc, data['domain'], data['page'], data['file_type'])
     if success:
@@ -218,6 +248,12 @@ def get_available_domains():
 
 @api.route('/poc/<string:project>/', methods=['POST'])
 def create_experiment(project: str):
+    if request.json is None:
+        return {
+            'status': 'NOK',
+            'msg': "No experiment parameters found"
+        }
+
     data = request.json.copy()
     if 'poc_name' not in data.keys():
         return {

bci/web/blueprints/experiments.py

@@ -5,7 +5,7 @@
 import sys
 
 import requests
-from flask import Blueprint, make_response, render_template, request, url_for
+from flask import Blueprint, Request, make_response, render_template, request, url_for
 
 from bci.web.page_parser import load_experiment_pages
 
@@ -26,6 +26,7 @@
 
 @exp.before_request
 def before_request():
+    __report(request)
     host = request.host.lower()
     if host not in ALLOWED_DOMAINS:
         logger.error(
@@ -34,29 +35,13 @@ def before_request():
         return f"Host '{host}' is not supported by this framework."
 
 
-@exp.route("/")
-def index():
-    return f"This page is visited over <b>{request.scheme}</b>."
-
-
-@exp.route("/report/", methods=["GET", "POST"])
-def report():
-    get_params = [item for item in get_all_GET_parameters(request).items()]
-    resp = make_response(
-        render_template("cookies.html", title="Report", get_params=get_params)
-    )
-
-    cookie_exp_date = datetime.datetime.now() + datetime.timedelta(weeks=4)
-    resp.set_cookie("generic", "1", expires=cookie_exp_date)
-    resp.set_cookie("secure", "1", expires=cookie_exp_date, secure=True)
-    resp.set_cookie("httpOnly", "1", expires=cookie_exp_date, httponly=True)
-    resp.set_cookie("lax", "1", expires=cookie_exp_date, samesite="lax")
-    resp.set_cookie("strict", "1", expires=cookie_exp_date, samesite="strict")
-
+def __report(request: Request) -> None:
+    """
+    Submit report to BugHog
+    """
     # Respond to collector on same IP
     # remote_ip = request.remote_addr
     remote_ip = request.headers.get("X-Real-IP")
-
     response_data = {
         "url": request.url,
         "method": request.method,
@@ -72,6 +57,29 @@ def send_report_to_collector():
 
     threading.Thread(target=send_report_to_collector).start()
 
+
+def __get_all_GET_parameters(request) -> dict[str,str]:
+    return {k: v for k, v in request.args.items()}
+
+
+@exp.route("/")
+def index():
+    return f"This page is visited over <b>{request.scheme}</b>."
+
+
+@exp.route("/report/", methods=["GET", "POST"])
+def report_endpoint():
+    get_params = [item for item in __get_all_GET_parameters(request).items()]
+    resp = make_response(
+        render_template("cookies.html", title="Report", get_params=get_params)
+    )
+
+    cookie_exp_date = datetime.datetime.now() + datetime.timedelta(weeks=4)
+    resp.set_cookie("generic", "1", expires=cookie_exp_date)
+    resp.set_cookie("secure", "1", expires=cookie_exp_date, secure=True)
+    resp.set_cookie("httpOnly", "1", expires=cookie_exp_date, httponly=True)
+    resp.set_cookie("lax", "1", expires=cookie_exp_date, samesite="lax")
+    resp.set_cookie("strict", "1", expires=cookie_exp_date, samesite="strict")
     return resp
 
 
@@ -81,7 +89,7 @@ def report_leak_if_using_http(target_scheme):
     Triggers request to /report/ if a request was received over the specified `scheme`.
     """
     used_scheme = request.headers.get("X-Forwarded-Proto")
-    params = get_all_GET_parameters(request)
+    params = __get_all_GET_parameters(request)
     if used_scheme == target_scheme:
         return "Redirect", 307, {"Location": url_for("experiments.report", **params)}
     else:
@@ -96,7 +104,7 @@ def report_leak_if_present(expected_header_name: str):
     if expected_header_name not in request.headers:
         return f"Header {expected_header_name} not found", 200, {"Allow-CSP-From": "*"}
 
-    params = get_all_GET_parameters(request)
+    params = __get_all_GET_parameters(request)
     return (
         "Redirect",
         307,
@@ -121,7 +129,7 @@ def report_leak_if_contains(expected_header_name: str, expected_header_value: st
             {"Allow-CSP-From": "*"},
         )
 
-    params = get_all_GET_parameters(request)
+    params = __get_all_GET_parameters(request)
     return (
         "Redirect",
         307,
@@ -132,15 +140,15 @@ def report_leak_if_contains(expected_header_name: str, expected_header_value: st
     )
 
 
-@exp.route("/<string:project>/<string:experiment>/<string:directory>/")
-def python_evaluation(project: str, experiment: str, directory: str):
+@exp.route("/<string:project>/<string:experiment>/<string:file_name>.py")
+def python_evaluation(project: str, experiment: str, file_name: str):
     """
     Evaluates the python script and returns its result.
     """
     host = request.host.lower()
 
-    module_name = f"{host}/{project}/{experiment}/{directory}"
-    path = f"experiments/pages/{project}/{experiment}/{host}/{directory}/index.py"
+    module_name = f"{host}/{project}/{experiment}"
+    path = f"experiments/pages/{project}/{experiment}/{host}/{file_name}.py"
 
     # Dynamically import the file
     sys.dont_write_bytecode = True
@@ -150,7 +158,3 @@ def python_evaluation(project: str, experiment: str, directory: str):
     spec.loader.exec_module(module)
 
     return module.main(request)
-
-  
-def get_all_GET_parameters(request):
-    return {k: v for k, v in request.args.items()}