Patched code sample

import ast
import os
import shlex

# This dictionary simulates the arguments provided to a roslaunch file.
# An attacker could control the value of 'my_param'.
launch_args = {
    'my_param_benign': "'hello world'",
    'my_param_malicious': "__import__('os').system('echo VULNERABILITY EXPLOITED')"
}

def _is_string_literal(s):
    """
    Check if a string is a valid python string literal.
    This is a simplified check for demonstration.
    """
    return (s.startswith("'") and s.endswith("'")) or \
           (s.startswith('"') and s.endswith('"'))

def _safe_eval(s):
    """
    Safely evaluate a string to a Python literal (string, number, list, etc.).
    This prevents the execution of arbitrary code by restricting evaluation
    to simple, static values. This is the core of the fix.
    """
    try:
        return ast.literal_eval(s)
    except (ValueError, SyntaxError, TypeError):
        # If it's not a valid literal, it should be treated as a plain string.
        # The original vulnerability failed to make this distinction.
        return s

def substitute_args_fixed(arg_string):
    """
    This function represents the fixed substitution logic in roslaunch.
    It safely processes arguments, avoiding arbitrary code execution.
    """
    # In the actual roslaunch code, this logic is more complex, involving
    # parsing for $(arg my_param). We simulate the final step where the
    # parameter's value is resolved and evaluated.
    if arg_string in launch_args:
        resolved_arg = launch_args[arg_string]
        
        # --- THE FIX ---
        # VULNERABLE CODE (REPLACED):
        # return eval(resolved_arg)
        
        # FIXED CODE (NOW IN USE):
        # The fix is to use a safe evaluation method or, if the argument
        # is intended to be a command, to handle it as such without eval.
        # If the argument is meant to be a literal, ast.literal_eval is used.
        # If it's not a literal, it's treated as a raw string.
        if _is_string_literal(resolved_arg):
             # Safely evaluate things like "'hello'" to "hello"
             return _safe_eval(resolved_arg)
        else:
             # For potentially dangerous expressions, run a safe evaluation
             # which will fail for code and return the original string.
             # This prevents the exploit.
             return _safe_eval(resolved_arg)
    return arg_string


if __name__ == '__main__':
    print("--- Demonstrating the FIX for CVE-2024-39835 ---")

    print("\n1. Processing a benign parameter...")
    benign_param_name = 'my_param_benign'
    benign_value = launch_args[benign_param_name]
    print(f"   Input argument value: {benign_value}")
    result = substitute_args_fixed(benign_param_name)
    print(f"   Safely evaluated result: '{result}' (type: {type(result).__name__})")
    print("   Result: Correctly evaluated as a string literal.")

    print("\n2. Processing a malicious parameter...")
    malicious_param_name = 'my_param_malicious'
    malicious_value = launch_args[malicious_param_name]
    print(f"   Input argument value: {malicious_value}")
    result = substitute_args_fixed(malicious_param_name)
    print(f"   Safely evaluated result: '{result}' (type: {type(result).__name__})")
    print("   Result: Malicious code was NOT executed and was treated as a plain string.")