{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# AI Security Lab\n",
    "\n",
    "This notebook provides hands-on exercises for understanding AI security vulnerabilities and defenses.\n",
    "\n",
    "## Learning Objectives\n",
    "\n",
    "1. Understand common AI security threats\n",
    "2. Implement basic security checks\n",
    "3. Test adversarial robustness\n",
    "4. Apply secure coding practices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setup and imports\n",
    "import numpy as np\n",
    "import hashlib\n",
    "import pickle\n",
    "import json\n",
    "from typing import Any, Dict, List, Tuple\n",
    "\n",
    "# For demonstrations\n",
    "from sklearn.datasets import load_iris\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.ensemble import RandomForestClassifier\n",
    "from sklearn.metrics import accuracy_score"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise 1: Input Validation\n",
    "\n",
    "Implement input validation to prevent malicious inputs from reaching your ML model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class InputValidator:\n",
    "    \"\"\"Validate inputs before sending to ML model.\"\"\"\n",
    "    \n",
    "    def __init__(self, feature_count: int, feature_ranges: Dict[int, Tuple[float, float]]):\n",
    "        self.feature_count = feature_count\n",
    "        self.feature_ranges = feature_ranges\n",
    "    \n",
    "    def validate(self, features: List[float]) -> Tuple[bool, str]:\n",
    "        \"\"\"\n",
    "        Validate input features.\n",
    "        \n",
    "        Returns:\n",
    "            Tuple of (is_valid, error_message)\n",
    "        \"\"\"\n",
    "        # TODO: Implement validation checks\n",
    "        # 1. Check feature count\n",
    "        # 2. Check for NaN/Inf values\n",
    "        # 3. Check value ranges\n",
    "        # 4. Check for injection patterns\n",
    "        \n",
    "        return True, \"\"\n",
    "\n",
    "# Test your implementation\n",
    "validator = InputValidator(\n",
    "    feature_count=4,\n",
    "    feature_ranges={\n",
    "        0: (0, 10),  # Sepal length\n",
    "        1: (0, 5),   # Sepal width\n",
    "        2: (0, 10),  # Petal length\n",
    "        3: (0, 5)    # Petal width\n",
    "    }\n",
    ")\n",
    "\n",
    "# Test cases\n",
    "test_inputs = [\n",
    "    [5.1, 3.5, 1.4, 0.2],     # Valid\n",
    "    [5.1, 3.5, 1.4],           # Wrong count\n",
    "    [5.1, float('nan'), 1.4, 0.2],  # NaN\n",
    "    [5.1, 3.5, 1.4, 100],     # Out of range\n",
    "]\n",
    "\n",
    "for i, features in enumerate(test_inputs):\n",
    "    is_valid, msg = validator.validate(features)\n",
    "    print(f\"Input {i+1}: {is_valid} - {msg}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise 2: Model Integrity Verification\n",
    "\n",
    "Implement hash verification to ensure model integrity before loading."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ModelIntegrityChecker:\n",
    "    \"\"\"Verify model integrity using cryptographic hashes.\"\"\"\n",
    "    \n",
    "    def __init__(self):\n",
    "        self.model_hashes = {}  # model_name -> expected_hash\n",
    "    \n",
    "    def compute_hash(self, model_data: bytes) -> str:\n",
    "        \"\"\"Compute SHA-256 hash of model data.\"\"\"\n",
    "        # TODO: Implement hash computation\n",
    "        return \"\"\n",
    "    \n",
    "    def register_model(self, name: str, model_data: bytes) -> str:\n",
    "        \"\"\"Register model and store its hash.\"\"\"\n",
    "        # TODO: Compute and store hash\n",
    "        return \"\"\n",
    "    \n",
    "    def verify_model(self, name: str, model_data: bytes) -> bool:\n",
    "        \"\"\"Verify model integrity against stored hash.\"\"\"\n",
    "        # TODO: Implement verification\n",
    "        return False\n",
    "\n",
    "# Test your implementation\n",
    "checker = ModelIntegrityChecker()\n",
    "\n",
    "# Simulate model data\n",
    "model_data = pickle.dumps(RandomForestClassifier())\n",
    "\n",
    "# Register model\n",
    "expected_hash = checker.register_model(\"iris_model\", model_data)\n",
    "print(f\"Expected hash: {expected_hash}\")\n",
    "\n",
    "# Verify model\n",
    "is_valid = checker.verify_model(\"iris_model\", model_data)\n",
    "print(f\"Model valid: {is_valid}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise 3: Adversarial Example Generation\n",
    "\n",
    "Create adversarial examples to test model robustness."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# First, train a simple model\n",
    "X, y = load_iris(return_X_y=True)\n",
    "X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)\n",
    "\n",
    "model = RandomForestClassifier(n_estimators=100, random_state=42)\n",
    "model.fit(X_train, y_train)\n",
    "\n",
    "print(f\"Model accuracy: {accuracy_score(y_test, model.predict(X_test)):.2%}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AdversarialAttacker:\n",
    "    \"\"\"Generate adversarial examples for testing.\"\"\"\n",
    "    \n",
    "    def __init__(self, model, epsilon: float = 0.1):\n",
    "        self.model = model\n",
    "        self.epsilon = epsilon\n",
    "    \n",
    "    def fgsm_attack(self, x: np.ndarray, true_label: int) -> np.ndarray:\n",
    "        \"\"\"\n",
    "        Fast Gradient Sign Method (FGSM) attack.\n",
    "        \n",
    "        Args:\n",
    "            x: Original input\n",
    "            true_label: True label of the input\n",
    "        \n",
    "        Returns:\n",
    "            Adversarial example\n",
    "        \"\"\"\n",
    "        # TODO: Implement FGSM\n",
    "        # 1. Compute gradient of loss w.r.t. input\n",
    "        # 2. Apply sign of gradient * epsilon\n",
    "        # 3. Clip to valid range\n",
    "        \n",
    "        return x  # Placeholder\n",
    "    \n",
    "    def test_robustness(self, X: np.ndarray, y: np.ndarray) -> dict:\n",
    "        \"\"\"Test model robustness against adversarial examples.\"\"\"\n",
    "        results = {\n",
    "            'original_accuracy': 0,\n",
    "            'adversarial_accuracy': 0,\n",
    "            'attack_success_rate': 0\n",
    "        }\n",
    "        \n",
    "        # TODO: Implement robustness testing\n",
    "        # 1. Test original accuracy\n",
    "        # 2. Generate adversarial examples\n",
    "        # 3. Test adversarial accuracy\n",
    "        # 4. Calculate attack success rate\n",
    "        \n",
    "        return results\n",
    "\n",
    "# Test your implementation\n",
    "attacker = AdversarialAttacker(model, epsilon=0.1)\n",
    "results = attacker.test_robustness(X_test, y_test)\n",
    "print(f\"Original accuracy: {results['original_accuracy']:.2%}\")\n",
    "print(f\"Adversarial accuracy: {results['adversarial_accuracy']:.2%}\")\n",
    "print(f\"Attack success rate: {results['attack_success_rate']:.2%}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise 4: Secure Model Serialization\n",
    "\n",
    "Implement secure serialization to prevent deserialization attacks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SecureSerializer:\n",
    "    \"\"\"Secure serialization/deserialization of ML models.\"\"\"\n",
    "    \n",
    "    ALLOWED_TYPES = {\"sklearn.ensemble._forest.RandomForestClassifier\"}\n",
    "    \n",
    "    def secure_serialize(self, model: Any) -> bytes:\n",
    "        \"\"\"\n",
    "        Serialize model securely.\n",
    "        \n",
    "        Steps:\n",
    "        1. Validate model type\n",
    "        2. Serialize with restricted pickle\n",
    "        3. Add integrity check\n",
    "        \"\"\"\n",
    "        # TODO: Implement secure serialization\n",
    "        return b\"\"\n",
    "    \n",
    "    def secure_deserialize(self, data: bytes, expected_hash: str) -> Any:\n",
    "        \"\"\"\n",
    "        Deserialize model securely.\n",
    "        \n",
    "        Steps:\n",
    "        1. Verify integrity\n",
    "        2. Validate types during deserialization\n",
    "        3. Return model\n",
    "        \"\"\"\n",
    "        # TODO: Implement secure deserialization\n",
    "        return None\n",
    "\n",
    "# Test your implementation\n",
    "serializer = SecureSerializer()\n",
    "\n",
    "# Serialize model\n",
    "model_data = serializer.secure_serialize(model)\n",
    "print(f\"Serialized size: {len(model_data)} bytes\")\n",
    "\n",
    "# Compute expected hash\n",
    "expected_hash = hashlib.sha256(model_data).hexdigest()\n",
    "print(f\"Expected hash: {expected_hash}\")\n",
    "\n",
    "# Deserialize and verify\n",
    "loaded_model = serializer.secure_deserialize(model_data, expected_hash)\n",
    "print(f\"Model loaded: {loaded_model is not None}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise 5: Security Audit Checklist\n",
    "\n",
    "Create a comprehensive security audit checklist for ML deployments."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MLSecurityAuditor:\n",
    "    \"\"\"Audit ML systems for security vulnerabilities.\"\"\"\n",
    "    \n",
    "    def __init__(self):\n",
    "        self.checks = []\n",
    "    \n",
    "    def add_check(self, name: str, category: str, severity: str):\n",
    "        \"\"\"Add a security check.\"\"\"\n",
    "        self.checks.append({\n",
    "            'name': name,\n",
    "            'category': category,\n",
    "            'severity': severity,\n",
    "            'status': 'pending'\n",
    "        })\n",
    "    \n",
    "    def run_audit(self, ml_system: Dict[str, Any]) -> Dict[str, Any]:\n",
    "        \"\"\"\n",
    "        Run security audit on ML system.\n",
    "        \n",
    "        Args:\n",
    "            ml_system: Dictionary describing the ML system\n",
    "        \n",
    "        Returns:\n",
    "            Audit report\n",
    "        \"\"\"\n",
    "        report = {\n",
    "            'total_checks': len(self.checks),\n",
    "            'passed': 0,\n",
    "            'failed': 0,\n",
    "            'findings': []\n",
    "        }\n",
    "        \n",
    "        # TODO: Implement audit logic\n",
    "        # For each check, evaluate the ml_system\n",
    "        # Record findings\n",
    "        \n",
    "        return report\n",
    "\n",
    "# Create auditor and add checks\n",
    "auditor = MLSecurityAuditor()\n",
    "\n",
    "# Add security checks\n",
    "auditor.add_check(\"Input Validation\", \"Data\", \"HIGH\")\n",
    "auditor.add_check(\"Model Integrity\", \"Model\", \"CRITICAL\")\n",
    "auditor.add_check(\"API Rate Limiting\", \"Infrastructure\", \"MEDIUM\")\n",
    "auditor.add_check(\"Dependency Scanning\", \"Supply Chain\", \"HIGH\")\n",
    "auditor.add_check(\"Logging & Monitoring\", \"Operations\", \"MEDIUM\")\n",
    "\n",
    "# Define sample ML system\n",
    "sample_system = {\n",
    "    'model_type': 'sklearn.ensemble.RandomForestClassifier',\n",
    "    'input_validation': True,\n",
    "    'rate_limiting': False,\n",
    "    'logging_enabled': True,\n",
    "    'dependencies_scanned': True\n",
    "}\n",
    "\n",
    "# Run audit\n",
    "report = auditor.run_audit(sample_system)\n",
    "print(f\"Audit complete: {report['passed']}/{report['total_checks']} checks passed\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "In this lab, you practiced:\n",
    "\n",
    "1. **Input Validation** - Preventing malicious inputs from reaching ML models\n",
    "2. **Model Integrity** - Verifying models haven't been tampered with\n",
    "3. **Adversarial Testing** - Generating adversarial examples to test robustness\n",
    "4. **Secure Serialization** - Preventing deserialization attacks\n",
    "5. **Security Auditing** - Systematic assessment of ML security\n",
    "\n",
    "## Next Steps\n",
    "\n",
    "1. Complete the exercises above\n",
    "2. Run the `ml-security-scanner` tool on your own ML pipelines\n",
    "3. Review the [OWASP ML Top 10](https://owasp.org/www-project-machine-learning-security-top-10/) for more threats\n",
    "4. Implement security testing in your CI/CD pipeline"
   ]
  }
 ],
 "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",
   "nbformat_minor": 4,
   "version": "3.8.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}