412 lines
13 KiB
Plaintext
412 lines
13 KiB
Plaintext
{
|
|
"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
|
|
}
|