# Insecure Deserialization via Pickle

Language: Python
Severity: Critical
CWE: CWE-502

## Source
6

## Flow
6-9

## Sink
9

## Vulnerable Code
```python
import pickle
import redis
from multiprocessing import Pool
def fetch_and_run_task(worker_id):
    cache = redis.StrictRedis(host='task-broker.internal', port=6379, db=2)
    raw_payload = cache.lpop(f'iot_worker_queue:{worker_id}')
    if not raw_payload:
        return None
    task_obj = pickle.loads(raw_payload)
    result = task_obj.execute()
    cache.rpush(f'iot_results:{worker_id}', str(result))
    return result
```

## Explanation

The raw bytes fetched from Redis via `cache.lpop()` on line 6 are untrusted external data. On line 9, `pickle.loads(raw_payload)` deserializes these bytes without any integrity check, signature verification, or allowlist validation. An attacker with write access to Redis or the ability to perform a MITM attack on the Redis connection can inject a crafted pickle payload that executes arbitrary OS commands during deserialization.

## Remediation

The fix eliminates `pickle.loads()` entirely and replaces it with `json.loads()`, which cannot execute arbitrary code during deserialization. An HMAC signature (SHA-256) is verified before the payload is parsed, ensuring that only messages produced by trusted parties (who share the secret) are accepted. Finally, an explicit allowlist of task types is enforced and dispatched to hard-coded handler functions, preventing any form of arbitrary code execution.

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

# This secret must be shared securely between task producers and workers.
# In production, load from a secrets manager or environment variable.
import os
HMAC_SECRET = os.environ["TASK_HMAC_SECRET"].encode("utf-8")

ALLOWED_TASK_TYPES = {
    "sensor_read",
    "firmware_update",
    "health_check",
    "config_sync",
}


def _compute_hmac(payload_bytes: bytes) -> str:
    return hmac.new(HMAC_SECRET, payload_bytes, hashlib.sha256).hexdigest()


def fetch_and_run_task(worker_id):
    cache = redis.StrictRedis(host='task-broker.internal', port=6379, db=2)
    raw_message = cache.lpop(f'iot_worker_queue:{worker_id}')
    if not raw_message:
        return None

    # ---------- integrity check ----------
    # The producer must push messages in the format:
    #   <hex-hmac>:<json-payload>
    # e.g.  cache.rpush(queue, f"{hmac_hex}:{json_payload}")
    try:
        raw_str = raw_message.decode("utf-8")
        sep_index = raw_str.index(":")
        received_sig = raw_str[:sep_index]
        json_payload = raw_str[sep_index + 1:]
    except (ValueError, UnicodeDecodeError):
        raise ValueError("Malformed task message: missing HMAC signature")

    expected_sig = _compute_hmac(json_payload.encode("utf-8"))
    if not hmac.compare_digest(received_sig, expected_sig):
        raise ValueError("Task message HMAC verification failed – payload rejected")

    # ---------- safe deserialization ----------
    task_data = json.loads(json_payload)

    if not isinstance(task_data, dict):
        raise ValueError("Task payload must be a JSON object")

    task_type = task_data.get("task_type")
    if task_type not in ALLOWED_TASK_TYPES:
        raise ValueError(f"Disallowed or unknown task type: {task_type!r}")

    task_params = task_data.get("params", {})
    if not isinstance(task_params, dict):
        raise ValueError("Task params must be a JSON object")

    # ---------- dispatch to a known, safe handler ----------
    result = _execute_task(task_type, task_params)

    cache.rpush(f'iot_results:{worker_id}', json.dumps({"result": result}))
    return result


def _execute_task(task_type: str, params: dict):
    """Route to a hard-coded handler based on the validated task type."""
    handler = _TASK_HANDLERS.get(task_type)
    if handler is None:
        raise ValueError(f"No handler registered for task type: {task_type!r}")
    return handler(params)


# ---- Register concrete, safe task handlers here ----
def _handle_sensor_read(params):
    return {"status": "ok", "sensor": params.get("sensor_id")}

def _handle_firmware_update(params):
    return {"status": "ok", "version": params.get("version")}

def _handle_health_check(params):
    return {"status": "healthy"}

def _handle_config_sync(params):
    return {"status": "synced"}


_TASK_HANDLERS = {
    "sensor_read": _handle_sensor_read,
    "firmware_update": _handle_firmware_update,
    "health_check": _handle_health_check,
    "config_sync": _handle_config_sync,
}
```