# Pickle Deserialization via __reduce__

Language: Python
Severity: Critical
CWE: CWE-502

## Source
10

## Flow
10-13-14-15

## Sink
15

## Vulnerable Code
```python
import pickle
import base64
from flask import Flask, request, jsonify

app = Flask(__name__)

@app.route('/iot/device/restore', methods=['POST'])
def restore_device_state():
    device_id = request.json.get('device_id')
    encoded_state = request.json.get('state_snapshot')
    
    if not encoded_state:
        return jsonify({'error': 'Missing state snapshot'}), 400
    
    try:
        state_data = base64.b64decode(encoded_state)
        device_config = pickle.loads(state_data)
        
        return jsonify({
            'status': 'restored',
            'device_id': device_id,
            'config': str(device_config)
        }), 200
    except Exception as e:
        return jsonify({'error': 'Failed to restore state'}), 500
```

## Explanation

The application accepts user-controlled base64-encoded data via the 'state_snapshot' parameter, decodes it, and directly deserializes it using pickle.loads(). Python's pickle module can execute arbitrary code during deserialization when specially crafted payloads exploit the __reduce__ method, allowing Remote Code Execution (RCE).

## Remediation

The fix replaces pickle.loads() with json.loads() for deserialization, which is inherently safe as JSON cannot execute arbitrary code during parsing. Additionally, a validation function checks that the deserialized configuration only contains expected keys with appropriate data types, providing defense-in-depth against malformed input.

## Secure Code
```python
import json
import base64
from flask import Flask, request, jsonify

app = Flask(__name__)

# Define allowed keys and types for device configuration
ALLOWED_CONFIG_KEYS = {
    'device_name': str,
    'firmware_version': str,
    'network_settings': dict,
    'sensors': list,
    'thresholds': dict,
    'enabled': bool,
    'update_interval': (int, float),
    'mode': str
}

def validate_device_config(config):
    """Validate that the deserialized config only contains expected keys and types."""
    if not isinstance(config, dict):
        raise ValueError("Configuration must be a dictionary")
    
    for key, value in config.items():
        if key not in ALLOWED_CONFIG_KEYS:
            raise ValueError(f"Unexpected configuration key: {key}")
        expected_type = ALLOWED_CONFIG_KEYS[key]
        if not isinstance(value, expected_type):
            raise ValueError(f"Invalid type for key '{key}'")
    
    return config

@app.route('/iot/device/restore', methods=['POST'])
def restore_device_state():
    device_id = request.json.get('device_id')
    encoded_state = request.json.get('state_snapshot')
    
    if not encoded_state:
        return jsonify({'error': 'Missing state snapshot'}), 400
    
    try:
        state_data = base64.b64decode(encoded_state)
        # Use JSON for safe deserialization instead of pickle
        device_config = json.loads(state_data)
        
        # Validate the structure of the deserialized config
        validated_config = validate_device_config(device_config)
        
        return jsonify({
            'status': 'restored',
            'device_id': device_id,
            'config': validated_config
        }), 200
    except (json.JSONDecodeError, ValueError) as e:
        return jsonify({'error': f'Invalid state snapshot: {str(e)}'}), 400
    except Exception as e:
        return jsonify({'error': 'Failed to restore state'}), 500
```