{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Secure Storage\n",
    "\n",
    "`fleche` supports securing filesystem-based caches (`PickleFile` and `CloudpickleFile`) using HMAC-SHA256 signatures. This protects against untrusted code execution (RCE) via tampered pickle files.\n",
    "\n",
    "Security is controlled via the `FLECHE_SECRET_KEY` environment variable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-03-20T00:50:02.926538Z",
     "iopub.status.busy": "2026-03-20T00:50:02.926364Z",
     "iopub.status.idle": "2026-03-20T00:50:03.480592Z",
     "shell.execute_reply": "2026-03-20T00:50:03.479528Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "No config file found. Using default memory cache.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saved data with digest key: c38b5861...\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "from pathlib import Path\n",
    "import tempfile\n",
    "from fleche.storage.pickle_file import PickleFile\n",
    "\n",
    "# Set a strong secret key\n",
    "os.environ[\"FLECHE_SECRET_KEY\"] = \"my_super_secret_key_123\"\n",
    "\n",
    "# Initialize a storage backend. It will automatically load the key from the environment.\n",
    "temp_dir = tempfile.TemporaryDirectory()\n",
    "storage = PickleFile.with_pickle(root=temp_dir.name)\n",
    "\n",
    "# Save a value\n",
    "value_to_cache = {\"status\": \"secure\", \"data\": [1, 2, 3]}\n",
    "digest_key = storage.save(value_to_cache)\n",
    "print(f\"Saved data with digest key: {digest_key[:8]}...\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When we load the data, `fleche` automatically verifies the signature. If it matches, the data is deserialized."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-03-20T00:50:03.482624Z",
     "iopub.status.busy": "2026-03-20T00:50:03.482275Z",
     "iopub.status.idle": "2026-03-20T00:50:03.486344Z",
     "shell.execute_reply": "2026-03-20T00:50:03.485544Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded securely: {'status': 'secure', 'data': [1, 2, 3]}\n"
     ]
    }
   ],
   "source": [
    "loaded_value = storage.load(digest_key)\n",
    "print(\"Loaded securely:\", loaded_value)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Tampering Protection\n",
    "\n",
    "If an attacker tampers with the file, the signature verification will fail, preventing the potentially malicious payload from being deserialized."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-03-20T00:50:03.488228Z",
     "iopub.status.busy": "2026-03-20T00:50:03.488042Z",
     "iopub.status.idle": "2026-03-20T00:50:03.493494Z",
     "shell.execute_reply": "2026-03-20T00:50:03.492525Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Invalid signature for cache entry. Potential tampering or key mismatch.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Security Error Caught: ('c38b58619f23f9507daf70dc7ef009512035480102031e8345c7e95fa254dbeb', 'Value present but failed signature check.')\n"
     ]
    }
   ],
   "source": [
    "file_path = Path(temp_dir.name) / digest_key\n",
    "original_content = file_path.read_bytes()\n",
    "\n",
    "# Simulate an attacker appending malicious code or modifying the data\n",
    "tampered_content = b\"malicious_payload\" + original_content\n",
    "file_path.write_bytes(tampered_content)\n",
    "\n",
    "try:\n",
    "    storage.load(digest_key)\n",
    "except KeyError as e:\n",
    "    print(\"Security Error Caught:\", e)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Key Rotation\n",
    "\n",
    "To support key rotation without invalidating existing caches, you can provide a colon-separated list of keys. `fleche` will sign new entries with the first key, and verify using any of the keys."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-03-20T00:50:03.495354Z",
     "iopub.status.busy": "2026-03-20T00:50:03.495141Z",
     "iopub.status.idle": "2026-03-20T00:50:03.499482Z",
     "shell.execute_reply": "2026-03-20T00:50:03.498788Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded with rotated key: {'status': 'secure', 'data': [1, 2, 3]}\n"
     ]
    }
   ],
   "source": [
    "# We roll out a new key, but keep the old one for reading\n",
    "os.environ[\"FLECHE_SECRET_KEY\"] = \"new_secret_key_456:my_super_secret_key_123\"\n",
    "storage_rotated = PickleFile.with_pickle(root=temp_dir.name)\n",
    "\n",
    "# Restore the original (untampered) file to test reading with the old key\n",
    "file_path.write_bytes(original_content)\n",
    "\n",
    "# It successfully loads using the second key in the list!\n",
    "rotated_loaded_value = storage_rotated.load(digest_key)\n",
    "print(\"Loaded with rotated key:\", rotated_loaded_value)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}