Commit 7f77062f authored by Ondřej Borýsek's avatar Ondřej Borýsek
Browse files

Switch to different single thread init mechanism

parent a81bc6e0

app.py

+14 −12
Original line number Diff line number Diff line
@@ -11,6 +11,7 @@ from werkzeug.middleware.proxy_fix import ProxyFix 


import config

import helpers.time_helper

import pwndoc_db_init

from helpers.file_utils import *



mimetypes.add_type('application/javascript', '.js')
@@ -50,9 +51,11 @@ def create_app(skip_pwndoc: bool = False): 
            migrate.init_app(app, db)



    # db.create_all() need to happen single threaded (or staggered)

    time.sleep(random.random())  # todo: do this properly

    with app.app_context():

        db.create_all()



    main_thread = pwndoc_db_init.determine_main_thread()

    if main_thread:

        with app.app_context():

            db.create_all()



    app.config['UPLOAD_FOLDER'] = config.FLASK_UPLOAD_FOLDER

    app.config['MAX_CONTENT_LENGTH'] = 100 * 1000 * 1000
@@ -61,7 +64,8 @@ def create_app(skip_pwndoc: bool = False): 


    Path(app.config['UPLOAD_FOLDER']).mkdir(parents=True, exist_ok=True)



    single_threaded_init(app, skip_pwndoc)

    if main_thread:

        single_threaded_init(app, skip_pwndoc)



    app.register_blueprint(process_findings_bp, url_prefix=f'{config.IMPORTER_URL_PREFIX}/findings')

    app.register_blueprint(templates_bp, url_prefix=f'{config.IMPORTER_URL_PREFIX}/templates')
@@ -106,24 +110,22 @@ def single_threaded_init(app, skip_pwndoc: bool): 
    from helpers.template_grouping import TemplateGrouping

    import pwndoc_db_init



    main_thread = True if skip_pwndoc else pwndoc_db_init.InitialData.setup_first_user()



    if not main_thread:

        return  # todo: All threads except the main one will be responsive before the full setup is completed. User actions could cause race condition.



    pwndoc_db_init.InitialData.upload_initial_data()

    logger.info("PwnDoc Init completed")



    pwndoc_db_init.PwnDocUpdate.setup_scan2report_plugin_folder()  # This will fail if run second time.

    TemplateGrouping.add_new_gids()



    if skip_pwndoc:

        return

    pwndoc_db_init.InitialData.setup_first_user()



    pwndoc_db_init.InitialData.upload_initial_data()

    logger.info("PwnDoc - initial configuration uploaded")



    pwndoc_db_init.PwnDocUpdate.add_new_fields_if_needed()

    logger.info("PwnDoc - update to new version completed")



    with app.app_context():

        pwndoc_db_init.PwnDocUpdate.upload_templates_from_scan2report_repository()

    logger.info("PwnDoc - public templates uploaded")





if __name__ == "__main__":

pwndoc_db_init.py

+19 −2
Original line number Diff line number Diff line
@@ -12,7 +12,8 @@ from typing import Dict, Optional, List 


import api_templates

import config

from helpers.file_utils import relative_path

from helpers.time_helper import current_timestamp

from helpers.file_utils import relative_path, json_safe_load, json_dump

from pathlib import Path



from pwndoc_api import session, test_connection
@@ -224,4 +225,20 @@ class PwnDocUpdate: 
                api_templates.upload_scan2report_templates(locale, template_filepaths)

            except Exception as e:

                logger.warning(f"Upload of initial templates failed for locale {locale}. Skipping. Reason: {e}")

                raise
 
 No newline at end of file
 
                raise





def determine_main_thread(new_init_after_x_seconds: int = 5, sleep_sec: float = 0.1) -> bool:

    CONCURENCY_FILE = relative_path("debug_tmp/init_main_thread.json")



    random_id = random.randrange(10 ** 9)

    current_time = current_timestamp()

    existing_file: Optional[dict] = json_safe_load(CONCURENCY_FILE)

    if existing_file and existing_file.get('timestamp', 0) > current_time - new_init_after_x_seconds:

        return False



    json_dump({'timestamp': current_time, 'worker_random_id': random_id}, CONCURENCY_FILE)

    time.sleep(sleep_sec)



    existing_file = json_safe_load(CONCURENCY_FILE)

    return existing_file['worker_random_id'] == random_id