# Insecure Deserialization via pickle.loads()

Language: Python
Severity: Critical
CWE: CWE-502

## Source
6

## Flow
6-8-9

## Sink
9

## Vulnerable Code
```python
import pickle
import redis
from multiprocessing import Pool
def fetch_and_execute_sensor_task(device_id):
    r = redis.StrictRedis(host='iot-broker', port=6379, db=0)
    raw_payload = r.lpop(f'sensor_queue:{device_id}')
    if not raw_payload:
        return None
    task_obj = pickle.loads(raw_payload)
    result = task_obj['handler'](task_obj['params'])
    return result
def dispatch_sensor_workers(dev_ids):
    with Pool(processes=4) as wp:
        outcomes = wp.map(fetch_and_execute_sensor_task, dev_ids)
    return outcomes
```

## Explanation

The raw_payload retrieved from Redis at line 6 is untrusted external data that flows directly into pickle.loads() at line 9 without any validation, signature verification, or integrity check. An attacker who can write to the Redis queue (e.g., via a compromised IoT device, network interception, or direct Redis access) can inject a crafted pickle payload that executes arbitrary OS commands during deserialization, before any application logic runs.

## Remediation

The fix eliminates pickle.loads() entirely, replacing it with JSON deserialization which cannot execute arbitrary code during parsing. An HMAC-SHA256 signature is required on every payload to ensure integrity and authenticity before any deserialization occurs, preventing tampered or injected messages from being processed. Additionally, a handler whitelist ensures only explicitly registered function names can be invoked, blocking any attempt to call arbitrary callables even if the JSON content were somehow manipulated.

## Secure Code
```python
import json
import hmac
import hashlib
import redis
from multiprocessing import Pool

# HMAC secret key — must be stored securely (e.g., env var, secrets manager)
# and shared only between trusted producer and consumer components.
import os
HMAC_SECRET_KEY = os.environ.get('SENSOR_QUEUE_HMAC_KEY', '').encode('utf-8')

# Whitelist of allowed handler function names mapped to actual callables.
# Only explicitly registered handlers can be invoked by task payloads.
ALLOWED_HANDLERS = {}

def register_handler(name, func):
    """Register a callable handler by name for task dispatch."""
    if not callable(func):
        raise ValueError(f"Handler '{name}' must be callable")
    ALLOWED_HANDLERS[name] = func

def _verify_hmac(message_bytes, received_mac_hex):
    """Verify HMAC-SHA256 signature of the message."""
    if not HMAC_SECRET_KEY:
        raise RuntimeError("SENSOR_QUEUE_HMAC_KEY environment variable is not set")
    expected_mac = hmac.new(HMAC_SECRET_KEY, message_bytes, hashlib.sha256).hexdigest()
    return hmac.compare_digest(expected_mac, received_mac_hex)

def _validate_task(task_obj):
    """Validate the structure and contents of a deserialized task object."""
    if not isinstance(task_obj, dict):
        raise ValueError("Task object must be a dictionary")
    if 'handler' not in task_obj or 'params' not in task_obj:
        raise ValueError("Task object must contain 'handler' and 'params' keys")
    handler_name = task_obj['handler']
    if not isinstance(handler_name, str):
        raise ValueError("Handler must be a string name")
    if handler_name not in ALLOWED_HANDLERS:
        raise ValueError(f"Unknown or disallowed handler: '{handler_name}'")
    if not isinstance(task_obj['params'], (dict, list, str, int, float, bool, type(None))):
        raise ValueError("Params must be a JSON-safe primitive or container type")
    return handler_name, task_obj['params']

def fetch_and_execute_sensor_task(device_id):
    r = redis.StrictRedis(host='iot-broker', port=6379, db=0)
    raw_payload = r.lpop(f'sensor_queue:{device_id}')
    if not raw_payload:
        return None

    # Expect payload format: "<hex_hmac>:<json_body>"
    try:
        raw_str = raw_payload.decode('utf-8')
    except (UnicodeDecodeError, AttributeError):
        raise ValueError("Payload is not valid UTF-8")

    if ':' not in raw_str:
        raise ValueError("Payload missing HMAC signature prefix")

    received_mac, json_body = raw_str.split(':', 1)

    # Step 1: Verify HMAC integrity/authenticity before parsing
    if not _verify_hmac(json_body.encode('utf-8'), received_mac):
        raise ValueError("HMAC verification failed — payload integrity compromised")

    # Step 2: Deserialize using safe JSON (no arbitrary code execution)
    task_obj = json.loads(json_body)

    # Step 3: Validate structure and whitelist the handler
    handler_name, params = _validate_task(task_obj)

    # Step 4: Invoke only the whitelisted handler
    result = ALLOWED_HANDLERS[handler_name](params)
    return result

def dispatch_sensor_workers(dev_ids):
    with Pool(processes=4) as wp:
        outcomes = wp.map(fetch_and_execute_sensor_task, dev_ids)
    return outcomes

# --- Example: Producer-side helper (for trusted task publishers) ---
# def publish_task(device_id, handler_name, params):
#     import json, hmac, hashlib, os, redis
#     secret = os.environ['SENSOR_QUEUE_HMAC_KEY'].encode('utf-8')
#     body = json.dumps({'handler': handler_name, 'params': params})
#     mac = hmac.new(secret, body.encode('utf-8'), hashlib.sha256).hexdigest()
#     r = redis.StrictRedis(host='iot-broker', port=6379, db=0)
#     r.lpush(f'sensor_queue:{device_id}', f'{mac}:{body}')
```