# Unsafe Pickle Deserialization via pickle.loads

Language: Python
Severity: Critical
CWE: CWE-502

## Source
9

## Flow
9-10

## Sink
10

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

app = Flask(__name__)

@app.route('/iot/device/restore', methods=['POST'])
def restore_device_config():
    encoded_state = request.json.get('device_state')
    device_config = pickle.loads(base64.b64decode(encoded_state))
    device_config['last_restored'] = None
    apply_device_settings(device_config)
    return {'status': 'restored', 'device_id': device_config.get('id')}, 200
```

## Explanation

The application accepts user-controlled base64-encoded data through the 'device_state' parameter and directly deserializes it using pickle.loads() without any validation. An attacker can craft a malicious pickled object that executes arbitrary code when deserialized, leading to Remote Code Execution (RCE) on the server.

## Remediation

The fix replaces pickle deserialization with JSON deserialization, which cannot execute arbitrary code during parsing. Additionally, a strict validation function checks that only allowed keys and safe value types are present in the configuration dictionary, providing defense-in-depth against malformed input.

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

app = Flask(__name__)

ALLOWED_CONFIG_KEYS = {'id', 'name', 'firmware_version', 'network', 'sensors', 'thresholds', 'intervals', 'enabled', 'last_restored'}
ALLOWED_VALUE_TYPES = (str, int, float, bool, list, dict, type(None))


def validate_config(config):
    """Validate that the deserialized config only contains safe types and allowed keys."""
    if not isinstance(config, dict):
        raise ValueError("Device config must be a dictionary")
    for key, value in config.items():
        if key not in ALLOWED_CONFIG_KEYS:
            raise ValueError(f"Unexpected config key: {key}")
        if not isinstance(value, ALLOWED_VALUE_TYPES):
            raise ValueError(f"Invalid value type for key '{key}': {type(value).__name__}")
    if 'id' not in config:
        raise ValueError("Device config must contain an 'id' field")
    return config


@app.route('/iot/device/restore', methods=['POST'])
def restore_device_config():
    encoded_state = request.json.get('device_state')
    if not encoded_state or not isinstance(encoded_state, str):
        return {'status': 'error', 'message': 'Invalid device_state parameter'}, 400

    try:
        decoded_data = base64.b64decode(encoded_state)
        device_config = json.loads(decoded_data)
    except Exception:
        return {'status': 'error', 'message': 'Invalid device state format'}, 400

    try:
        device_config = validate_config(device_config)
    except ValueError as e:
        return {'status': 'error', 'message': str(e)}, 400

    device_config['last_restored'] = time.time()
    apply_device_settings(device_config)
    return {'status': 'restored', 'device_id': device_config.get('id')}, 200
```