wiliot-certificate 4.5.0a2__py3-none-any.whl → 4.5.0a4__py3-none-any.whl

gui_certificate/server.py

@@ -1,7 +1,8 @@
 import os
 import json
 import tempfile
-from flask import Flask, render_template, request, session, redirect, url_for, flash, send_file, jsonify
+import datetime
+from flask import Flask, render_template, request, session, redirect, url_for, flash, send_file
 from jinja2 import ChoiceLoader, PackageLoader
 from werkzeug.utils import secure_filename
 from certificate.certificate_cli import CertificateCLI
@@ -31,9 +32,21 @@ TEMP_BASE = tempfile.gettempdir()
 UPLOAD_FOLDER = os.path.join(TEMP_BASE, "wiliot_certificate_uploaded_schemas")
 CONFIG_FOLDER = os.path.join(TEMP_BASE, "wiliot_certificate_saved_configs")
 TEST_LIST_FOLDER = os.path.join(TEMP_BASE, "wiliot_certificate_test_lists")
+CUSTOM_BROKER_FOLDER = os.path.join(TEMP_BASE, "wiliot_certificate_custom_brokers")
 ALLOWED_EXTENSIONS = {'json'}
 MAX_UPLOAD_SIZE = 16 * 1024 * 1024  # 16MB
 
+# Default custom broker configuration (from hivemq.json)
+DEFAULT_CUSTOM_BROKER = {
+    "port": 8883,
+    "brokerUrl": "mqtts://broker.hivemq.com",
+    "username": "",
+    "password": "",
+    "updateTopic": "update/wiliot/<gatewayId>",
+    "statusTopic": "status/wiliot/<gatewayId>",
+    "dataTopic": "data/wiliot/<gatewayId>"
+}
+
 app = Flask(__name__)
 app.secret_key = os.urandom(24)  # Required for sessions
 app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
@@ -43,6 +56,7 @@ app.config['MAX_CONTENT_LENGTH'] = MAX_UPLOAD_SIZE
 os.makedirs(UPLOAD_FOLDER, exist_ok=True)
 os.makedirs(CONFIG_FOLDER, exist_ok=True)
 os.makedirs(TEST_LIST_FOLDER, exist_ok=True)
+os.makedirs(CUSTOM_BROKER_FOLDER, exist_ok=True)
 
 # extend Jinja search path to include shared dir
 app.jinja_loader = ChoiceLoader([
@@ -70,9 +84,32 @@ def get_mandatory_modules_for_flavor(flavor):
         return [MODULE_CLOUD_CONNECTIVITY, MODULE_EDGE_MGMT]
     return []
 
+def get_mandatory_modules_by_schema(schema_data, flavor):
+    """Get modules that are mandatory by schema (modules declared in schema)."""
+    if not schema_data or "modules" not in schema_data:
+        return []
+    
+    # Only apply to bridge flavors (bridge_only and combo)
+    if flavor not in (FLAVOR_BRIDGE_ONLY, FLAVOR_COMBO):
+        return []
+    
+    schema_modules = schema_data.get("modules", {})
+    if not isinstance(schema_modules, dict):
+        return []
+    
+    # Map schema module names to test module names
+    mandatory_by_schema = []
+    for schema_mod_name in schema_modules.keys():
+        # Map schema module name to test module name
+        test_mod_name = test_module_and_schema_module_to_other(schema_mod_name)
+        if test_mod_name not in mandatory_by_schema and test_mod_name != 'custom':
+            mandatory_by_schema.append(test_mod_name)
+    
+    return mandatory_by_schema
+
 # TODO: This is a temporary mapping to map test module names to schema module names - remove this
-def test_module_to_schema_module(test_module_name):
-    """Map test module names to validation schema module names."""
+def test_module_and_schema_module_to_other(module_name):
+    """Map test module names and schema module names to one another."""
     mapping = {
         "calibration": "calibration",
         "datapath": "datapath",
@@ -82,7 +119,12 @@ def test_module_to_schema_module(test_module_name):
         "sensors": "externalSensor",  # different name
         "custom": "custom"
     }
-    return mapping.get(test_module_name, test_module_name)  # Return as-is if not in mapping
+    for k, v in mapping.items():
+        if module_name == k:
+            return v
+        if module_name == v:
+            return k
+    return module_name  # Return as-is if not in mapping
 
 def filter_modules_by_flavor(tests_schema, flavor, schema_data=None):
     """Filter modules and tests based on selected flavor."""
@@ -115,12 +157,17 @@ def filter_modules_by_flavor(tests_schema, flavor, schema_data=None):
                 # Include all tests for combo
                 filtered_tests.append(test)
         
-        # Only include module if it has tests or is mandatory
+        # Get modules mandatory by schema
+        mandatory_by_schema = get_mandatory_modules_by_schema(schema_data, flavor) if schema_data else []
+        is_mandatory_by_schema = mod_name in mandatory_by_schema
+        
+        # Only include module if it has tests or is mandatory (by flavor or schema)
         # For mandatory modules, include even if no tests match (they'll be shown but empty)
-        if filtered_tests or mod_name in mandatory_modules:
+        if filtered_tests or mod_name in mandatory_modules or is_mandatory_by_schema:
             mod_copy = mod.copy()
             mod_copy["tests"] = filtered_tests
             mod_copy["is_mandatory"] = mod_name in mandatory_modules
+            mod_copy["is_mandatory_by_schema"] = is_mandatory_by_schema
             filtered_modules.append(mod_copy)
     
     # Sort modules: cloud_connectivity first, then edge_mgmt, then by schema order
@@ -136,7 +183,7 @@ def filter_modules_by_flavor(tests_schema, flavor, schema_data=None):
         if schema_module_order:
             try:
                 # Map test module name to schema module name
-                schema_mod_name = test_module_to_schema_module(mod_name)
+                schema_mod_name = test_module_and_schema_module_to_other(mod_name)
                 idx = schema_module_order.index(schema_mod_name)
                 return (2, idx)
             except ValueError:
@@ -285,30 +332,37 @@ def verify_schema_matches_selection(schema_data, flavor, selected_modules, selec
             bridge_modules_to_check = [m for m in selected_modules if m not in (MODULE_CLOUD_CONNECTIVITY, MODULE_EDGE_MGMT)]
             for test_mod_name in bridge_modules_to_check:
                 # Map test module name to schema module name
-                schema_mod_name = test_module_to_schema_module(test_mod_name)
+                schema_mod_name = test_module_and_schema_module_to_other(test_mod_name)
                 if schema_mod_name not in schema_module_names:
                     warnings.append(f"Module '{test_mod_name}' may not be supported according to the uploaded schema")
     
     return len(errors) == 0, errors, warnings
 
-def check_certification_status(flavor, selected_modules, selected_tests, tests_schema):
+def check_certification_status(flavor, selected_modules, selected_tests, tests_schema, schema_data=None, unsterile_run=False):
     """Check if the selection qualifies for certification or is test-only."""
     mandatory_modules = get_mandatory_modules_for_flavor(flavor)
+    mandatory_by_schema = get_mandatory_modules_by_schema(schema_data, flavor) if schema_data else []
     all_mandatory_tests = []
     
-    # Get all mandatory tests from mandatory modules
+    # Get selected module names (handle both dict and string formats)
+    selected_module_names = [m.get("name") if isinstance(m, dict) else m for m in selected_modules]
+    
+    # Get all mandatory tests from ALL selected modules (not just mandatory modules)
+    # When a module is selected, its mandatory tests become mandatory
     for mod in tests_schema.get("modules", []):
         mod_name = mod.get("name", "")
-        if mod_name in mandatory_modules:
+        if mod_name in selected_module_names:
             for test in mod.get("tests", []):
                 if test.get("meta", {}).get("mandatory", 0) == 1:
                     all_mandatory_tests.append(test.get("id"))
     
     # Check if all mandatory modules are selected
-    selected_module_names = [m.get("name") if isinstance(m, dict) else m for m in selected_modules]
     missing_modules = [m for m in mandatory_modules if m not in selected_module_names]
     
-    # Check if all mandatory tests are selected
+    # Check if all schema-mandatory modules are selected
+    missing_modules_by_schema = [m for m in mandatory_by_schema if m not in selected_module_names]
+    
+    # Check if all mandatory tests (from selected modules) are selected
     missing_tests = [t for t in all_mandatory_tests if t not in selected_tests]
     
     # Check for "at least one additional module" requirement for Bridge Only and Combo
@@ -319,9 +373,11 @@ def check_certification_status(flavor, selected_modules, selected_tests, tests_s
         if len(additional_modules) == 0:
             missing_additional_module = True
     
-    is_certified = len(missing_modules) == 0 and len(missing_tests) == 0 and not missing_additional_module
+    # A certifying run must be sterile - if unsterile_run is set, force non-certifying
+    is_certified = (len(missing_modules) == 0 and len(missing_modules_by_schema) == 0 and 
+                   len(missing_tests) == 0 and not missing_additional_module and not unsterile_run)
     
-    return is_certified, missing_modules, missing_tests, missing_additional_module
+    return is_certified, missing_modules, missing_tests, missing_additional_module, missing_modules_by_schema
 
 def _prepare_form_data_for_template(session, cert_schema):
     """Prepare form_data for template, ensuring custom_broker_path is included if available."""
@@ -424,6 +480,43 @@ def _write_testlist(selected_ids: list[str], form_data=None, is_certified=False)
     print(f"Custom test list saved in {path}")
     return path
 
+def _write_custom_broker_config(form_data, cert_schema_title="cert_run") -> str:
+    """
+    Generate a custom broker JSON configuration file from form data.
+    Extracts broker field values and creates a JSON file matching hivemq.json format.
+    """
+    # Extract broker configuration from form_data
+    broker_config = {}
+    
+    # Field names follow pattern: {schema_title}_custom_broker_{field_name}
+    field_prefix = f"{cert_schema_title}_custom_broker_"
+    
+    # Extract values from form_data, using defaults if not provided
+    for field in DEFAULT_CUSTOM_BROKER.keys():
+        field_name = f"{field_prefix}{field}"
+        value = form_data.get(field_name, "")
+        
+        # Convert port to int if it's a number
+        if field == "port":
+            try:
+                broker_config[field] = int(value) if value else DEFAULT_CUSTOM_BROKER[field]
+            except (ValueError, TypeError):
+                broker_config[field] = DEFAULT_CUSTOM_BROKER[field]
+        else:
+            broker_config[field] = value
+    
+    # Generate filename with timestamp
+    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+    filename = f"custom_broker_{timestamp}.json"
+    filepath = os.path.join(CUSTOM_BROKER_FOLDER, filename)
+    
+    # Write JSON file
+    with open(filepath, 'w', encoding='utf-8') as f:
+        json.dump(broker_config, f, indent=4)
+    
+    print(f"Custom broker config saved in {filepath}")
+    return filepath
+
 @app.route("/")
 def index():
     """Redirect to step 1 or show current step."""
@@ -451,6 +544,7 @@ def step(step_num):
         session['selected_modules'] = []
         session['selected_tests'] = []
         session['form_data'] = {}
+        session['schema_mandatory_initialized'] = False
     
     title = "Run Certificate"
     tests_schema = web_utils.scan_tests_dir(CERT_TESTS_ROOT)
@@ -505,7 +599,22 @@ def step(step_num):
                             flavor = session.get('flavor')
                             if 'selected_modules' not in session or not session.get('selected_modules'):
                                 mandatory_modules = get_mandatory_modules_for_flavor(flavor)
-                                session['selected_modules'] = mandatory_modules
+                                # Add schema-mandatory modules
+                                schema_data = result  # result is the validated schema data
+                                mandatory_by_schema = get_mandatory_modules_by_schema(schema_data, flavor)
+                                all_mandatory = list(set(mandatory_modules + mandatory_by_schema))
+                                session['selected_modules'] = all_mandatory
+                                # Mark schema-mandatory modules as initialized
+                                session['schema_mandatory_initialized'] = True
+                            else:
+                                # If modules already exist, add schema-mandatory modules and mark as initialized
+                                schema_data = result
+                                mandatory_by_schema = get_mandatory_modules_by_schema(schema_data, flavor)
+                                current_modules = session.get('selected_modules', [])
+                                updated_modules = list(set(current_modules + mandatory_by_schema))
+                                if len(updated_modules) > len(current_modules):
+                                    session['selected_modules'] = updated_modules
+                                session['schema_mandatory_initialized'] = True
                             # For GW only, skip step 3 (modules)
                             if session.get('flavor') == FLAVOR_GW_ONLY:
                                 session['current_step'] = 4
@@ -575,43 +684,19 @@ def step(step_num):
                 else:
                     form_data_dict[key] = request.form.get(key)
             
-            # Handle custom_broker file upload
+            # Handle custom_broker configuration from form fields
             cert_cli = CertificateCLI()
             cert_schema = web_utils.parser_to_schema(cert_cli.parser)
             custom_broker_field_name = f"{cert_schema['title']}_custom_broker"
-            custom_broker_file_key = f"{custom_broker_field_name}_file"
-            
-            # Preserve existing custom_broker_path if no new file is uploaded
-            existing_custom_broker_path = session.get('custom_broker_path')
-            if existing_custom_broker_path:
-                form_data_dict[custom_broker_field_name] = existing_custom_broker_path
             
-            if custom_broker_file_key in request.files:
-                file = request.files[custom_broker_file_key]
-                if file and file.filename and allowed_file(file.filename):
-                    filename = secure_filename(file.filename)
-                    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
-                    filename = f"{timestamp}_{filename}"
-                    filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
-                    file.save(filepath)
-                    
-                    # Validate file is valid JSON
-                    try:
-                        with open(filepath, "r", encoding="utf-8") as f:
-                            json.load(f)
-                        # Store normalized path
-                        normalized_path = os.path.normpath(filepath)
-                        session['custom_broker_path'] = normalized_path
-                        form_data_dict[custom_broker_field_name] = normalized_path
-                    except json.JSONDecodeError as e:
-                        error = f"Invalid JSON in custom broker file: {e}"
-                        os.remove(filepath)
-                    except Exception as e:
-                        error = f"Error reading custom broker file: {e}"
-                        if os.path.exists(filepath):
-                            os.remove(filepath)
-                elif file and file.filename:
-                    error = "Custom broker file must be a JSON file"
+            # Generate custom broker JSON file from form fields
+            try:
+                broker_config_path = _write_custom_broker_config(form_data_dict, cert_schema['title'])
+                normalized_path = os.path.normpath(broker_config_path)
+                session['custom_broker_path'] = normalized_path
+                form_data_dict[custom_broker_field_name] = normalized_path
+            except Exception as e:
+                error = f"Error generating custom broker configuration: {e}"
             
             # Validate all required fields (required for moving forward, but allow going back)
             if action == 'next':
@@ -631,13 +716,22 @@ def step(step_num):
                     field_name = f"{cert_schema['title']}_{field['name']}"
                     field_value = form_data_dict.get(field_name, '')
                     
-                    # Special handling for custom_broker file upload
+                    # Special handling for custom_broker - check if all required broker fields are filled
                     if field['name'] == 'custom_broker':
-                        # Check if file was uploaded or path exists in session
-                        custom_broker_file_key = f"{field_name}_file"
-                        if custom_broker_file_key not in request.files or not request.files[custom_broker_file_key].filename:
-                            # Check if path exists in form_data or session
-                            if not field_value and not session.get('custom_broker_path'):
+                        # Check all broker fields except username and password are filled
+                        broker_field_prefix = f"{cert_schema['title']}_custom_broker_"
+                        optional_fields = {'username', 'password'}
+                        for broker_field_key in DEFAULT_CUSTOM_BROKER.keys():
+                            if broker_field_key not in optional_fields:
+                                broker_field_name = f"{broker_field_prefix}{broker_field_key}"
+                                broker_field_value = form_data_dict.get(broker_field_name, '')
+                                if not broker_field_value or (isinstance(broker_field_value, str) and not broker_field_value.strip()):
+                                    missing_required.append(f"custom_broker.{broker_field_key}")
+                                    break
+                        # Also check if path was generated (should be generated above if fields are valid)
+                        if not field_value and not session.get('custom_broker_path'):
+                            # Only add if we haven't already added a broker field error
+                            if not any('custom_broker.' in req for req in missing_required):
                                 missing_required.append(field.get('name', field['label']))
                     else:
                         # For other fields, check if value is provided
@@ -681,9 +775,30 @@ def step(step_num):
     # This ensures mandatory items are only pre-checked at the beginning
     if flavor and not selected_modules:
         mandatory_modules = get_mandatory_modules_for_flavor(flavor)
-        selected_modules = mandatory_modules
+        # Add schema-mandatory modules if schema is available
+        schema_data = load_schema_data(schema_path) if schema_path else None
+        mandatory_by_schema = get_mandatory_modules_by_schema(schema_data, flavor) if schema_data else []
+        all_mandatory = list(set(mandatory_modules + mandatory_by_schema))
+        selected_modules = all_mandatory
         session['selected_modules'] = selected_modules
     
+    # Ensure schema-mandatory modules are added when first reaching step 3 (only if not already initialized)
+    # Use a session flag to track if schema-mandatory modules have been initialized
+    if step_num == 3 and flavor and schema_path:
+        schema_mandatory_initialized = session.get('schema_mandatory_initialized', False)
+        if not schema_mandatory_initialized:
+            schema_data = load_schema_data(schema_path) if schema_path else None
+            if schema_data:
+                mandatory_by_schema = get_mandatory_modules_by_schema(schema_data, flavor)
+                # Add schema-mandatory modules that aren't already selected
+                current_modules = session.get('selected_modules', [])
+                updated_modules = list(set(current_modules + mandatory_by_schema))
+                if len(updated_modules) > len(current_modules):
+                    selected_modules = updated_modules
+                    session['selected_modules'] = selected_modules
+                # Mark as initialized so we don't re-add them when going back
+                session['schema_mandatory_initialized'] = True
+    
     # Initialize mandatory tests only if not already in session and we have selected modules
     # Only initialize when first reaching step 4 (not on every GET request)
     if flavor and selected_modules and not selected_tests and step_num == 4:
@@ -714,17 +829,28 @@ def step(step_num):
         schema_data = load_schema_data(schema_path) if schema_path else None
         filtered_modules = filter_modules_by_flavor(tests_schema, flavor, schema_data)
     
+    # Get CLI schemas for form fields - use certificate_cli for all flavors (needed early for unsterile_run check)
+    cert_cli = CertificateCLI()
+    cert_schema = web_utils.parser_to_schema(cert_cli.parser)
+    
+    # Check if unsterile_run is set in form_data
+    form_data = _prepare_form_data_for_template(session, cert_schema)
+    unsterile_run_field = f"{cert_schema['title']}_unsterile_run"
+    unsterile_run = bool(form_data.get(unsterile_run_field))
+    
     # Get certification status
     is_certified = True
     missing_modules = []
     missing_tests = []
     missing_tests_details = []  # List of dicts with test id and label
     missing_additional_module = False
+    missing_modules_by_schema = []
     if step_num >= 4 and flavor and selected_modules and selected_tests:
         # Convert selected_modules to module dicts for check_certification_status
         module_dicts = [m for m in filtered_modules if m.get('name') in selected_modules]
-        is_certified, missing_modules, missing_tests, missing_additional_module = check_certification_status(
-            flavor, module_dicts, selected_tests, tests_schema
+        schema_data = load_schema_data(schema_path) if schema_path else None
+        is_certified, missing_modules, missing_tests, missing_additional_module, missing_modules_by_schema = check_certification_status(
+            flavor, module_dicts, selected_tests, tests_schema, schema_data, unsterile_run
         )
         
         # Get test details (labels) for missing tests
@@ -751,10 +877,12 @@ def step(step_num):
             )
             schema_errors = errors
             schema_warnings = warnings
+            # Get missing schema-mandatory modules for display (recalculate if not already set)
+            if not missing_modules_by_schema and step_num >= 4 and flavor and selected_modules:
+                _, _, _, _, missing_modules_by_schema = check_certification_status(
+                    flavor, module_dicts, selected_tests, tests_schema, schema_data, unsterile_run
+                )
     
-    # Get CLI schemas for form fields - use certificate_cli for all flavors
-    cert_cli = CertificateCLI()
-    cert_schema = web_utils.parser_to_schema(cert_cli.parser)
     
     # Calculate total test count including mandatory tests from selected modules
     total_test_count = len(selected_tests)
@@ -770,6 +898,8 @@ def step(step_num):
     # Prepare data for JavaScript validation
     # Always initialize these to avoid Undefined errors in template
     mandatory_modules_for_js = get_mandatory_modules_for_flavor(flavor) if flavor else []
+    schema_data_for_js = load_schema_data(schema_path) if schema_path else None
+    mandatory_modules_by_schema_for_js = get_mandatory_modules_by_schema(schema_data_for_js, flavor) if (schema_data_for_js and flavor) else []
     mandatory_tests_for_js = []
     
     if step_num in [3, 4] and flavor:
@@ -839,16 +969,21 @@ def step(step_num):
                            missing_additional_module=missing_additional_module,
                            schema_errors=schema_errors,
                            schema_warnings=schema_warnings,
+                           missing_modules_by_schema=missing_modules_by_schema,
                            mandatory_modules_for_js=mandatory_modules_for_js,
+                           mandatory_modules_by_schema_for_js=mandatory_modules_by_schema_for_js,
                            mandatory_tests_for_js=mandatory_tests_for_js,
                            error=error,
                            warning=warning,
                            form_data=_prepare_form_data_for_template(session, cert_schema),
+                           custom_broker_defaults=DEFAULT_CUSTOM_BROKER,
+                           custom_broker_field_keys=list(DEFAULT_CUSTOM_BROKER.keys()),
                            run_completed=run_completed,
                            run_terminal=run_terminal,
                            run_pid=run_pid,
                            run_is_certified=run_is_certified,
-                           available_ports=available_ports)
+                           available_ports=available_ports,
+                           unsterile_run=unsterile_run)
 
 def execute_certificate():
     """Execute the certificate run based on session data."""
@@ -874,7 +1009,13 @@ def execute_certificate():
     tests_schema = web_utils.scan_tests_dir(CERT_TESTS_ROOT)
     filtered_modules = filter_modules_by_flavor(tests_schema, flavor)
     module_dicts = [m for m in filtered_modules if m.get('name') in session.get('selected_modules', [])]
-    is_certified, _, _, _ = check_certification_status(flavor, module_dicts, selected_tests, tests_schema)
+    schema_data = load_schema_data(schema_path) if schema_path else None
+    # Check if unsterile_run is set in form_data
+    cert_cli = CertificateCLI()
+    cert_schema = web_utils.parser_to_schema(cert_cli.parser)
+    unsterile_run_field = f"{cert_schema['title']}_unsterile_run"
+    unsterile_run = bool(form_data.get(unsterile_run_field))
+    is_certified, _, _, _, _ = check_certification_status(flavor, module_dicts, selected_tests, tests_schema, schema_data, unsterile_run)
     
     full_cmd = []
     
@@ -937,6 +1078,17 @@ def export_config():
         flash("No configuration to export. Please complete at least step 1 and 2.", "error")
         return redirect(url_for('step', step_num=session.get('current_step', 1)))
     
+    # Extract custom broker form field values from form_data
+    form_data = session.get('form_data', {})
+    cert_cli = CertificateCLI()
+    cert_schema = web_utils.parser_to_schema(cert_cli.parser)
+    custom_broker_fields = {}
+    field_prefix = f"{cert_schema['title']}_custom_broker_"
+    for field_name in DEFAULT_CUSTOM_BROKER.keys():
+        full_field_name = f"{field_prefix}{field_name}"
+        if full_field_name in form_data:
+            custom_broker_fields[field_name] = form_data[full_field_name]
+    
     # Prepare configuration data
     config = {
         'version': '1.0',
@@ -945,9 +1097,10 @@ def export_config():
         'schema_path': session.get('schema_path'),
         'selected_modules': session.get('selected_modules', []),
         'selected_tests': session.get('selected_tests', []),
-        'form_data': session.get('form_data', {}),
+        'form_data': form_data,
         'schema_data': load_schema_data(session.get('schema_path')),  # Load schema data from file for export
-        'custom_broker_path': session.get('custom_broker_path')  # Include custom broker file path
+        'custom_broker_path': session.get('custom_broker_path'),  # Include custom broker file path for backward compatibility
+        'custom_broker_fields': custom_broker_fields  # Include custom broker form field values
     }
     
     # Create filename with timestamp
@@ -1008,17 +1161,47 @@ def import_config():
                     json.dump(schema_data, f, indent=2)
             except Exception:
                 pass  # If we can't save, user will need to re-upload
-        session['custom_broker_path'] = config_data.get('custom_broker_path')
+        # Handle custom broker configuration
+        custom_broker_fields = config_data.get('custom_broker_fields', {})
+        if custom_broker_fields:
+            # Restore broker form field values to form_data
+            cert_cli = CertificateCLI()
+            cert_schema = web_utils.parser_to_schema(cert_cli.parser)
+            field_prefix = f"{cert_schema['title']}_custom_broker_"
+            for field_name, field_value in custom_broker_fields.items():
+                full_field_name = f"{field_prefix}{field_name}"
+                session['form_data'][full_field_name] = field_value
+            # Regenerate broker config file from form fields
+            try:
+                broker_config_path = _write_custom_broker_config(session['form_data'], cert_schema['title'])
+                session['custom_broker_path'] = os.path.normpath(broker_config_path)
+            except Exception:
+                pass  # If generation fails, user can reconfigure in step 4
+        else:
+            # Backward compatibility: try to use existing path
+            session['custom_broker_path'] = config_data.get('custom_broker_path')
         
         # Validate schema file still exists
         if session['schema_path'] and not os.path.exists(session['schema_path']):
             flash(f"Warning: Schema file not found at {session['schema_path']}. Please re-upload it in step 2.", "warning")
             session['current_step'] = 2
-        # Validate custom_broker file still exists if present
+        # Validate custom_broker file still exists if present (or regenerate if we have fields)
         elif session.get('custom_broker_path') and not os.path.exists(session['custom_broker_path']):
-            flash(f"Warning: Custom broker file not found at {session['custom_broker_path']}. Please re-upload it in step 4.", "warning")
-            if session.get('current_step', 0) < 4:
-                session['current_step'] = 4
+            if custom_broker_fields:
+                # Regenerate from fields if file doesn't exist
+                try:
+                    cert_cli = CertificateCLI()
+                    cert_schema = web_utils.parser_to_schema(cert_cli.parser)
+                    broker_config_path = _write_custom_broker_config(session['form_data'], cert_schema['title'])
+                    session['custom_broker_path'] = os.path.normpath(broker_config_path)
+                except Exception:
+                    flash(f"Warning: Custom broker configuration will be regenerated in step 4.", "warning")
+                    if session.get('current_step', 0) < 4:
+                        session['current_step'] = 4
+            else:
+                flash(f"Warning: Custom broker file not found at {session['custom_broker_path']}. Please reconfigure it in step 4.", "warning")
+                if session.get('current_step', 0) < 4:
+                    session['current_step'] = 4
         else:
             # Determine appropriate step based on what's configured
             if session.get('selected_tests'):
@@ -1047,6 +1230,8 @@ def import_config():
 def clear_session():
     """Clear session and start fresh."""
     session.clear()
+    # Initialize session flags
+    session['schema_mandatory_initialized'] = False
     return redirect(url_for('step', step_num=1))
 
 @app.route("/parser")