Privacy & Security

# pip install opacus torch presidio-analyzer presidio-anonymizer

# ─── 1. DP-SGD with Opacus (PyTorch) ─────────────────────────
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
from opacus import PrivacyEngine
from opacus.validators import ModuleValidator
import numpy as np

# Simple model
class LinearClassifier(nn.Module):
    def __init__(self, input_dim=20, num_classes=2):
        super().__init__()
        self.fc = nn.Sequential(
            nn.Linear(input_dim, 64), nn.ReLU(),
            nn.Linear(64, num_classes)
        )
    def forward(self, x):
        return self.fc(x)

# Synthetic dataset
X = torch.randn(1000, 20)
y = (X[:, 0] + X[:, 1] > 0).long()
dataset = TensorDataset(X, y)
loader = DataLoader(dataset, batch_size=64, shuffle=True)

model = LinearClassifier()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()

# Make model compatible with Opacus (fixes BatchNorm → GroupNorm)
model = ModuleValidator.fix(model)
errors = ModuleValidator.validate(model, strict=False)
print(f"Model validation errors: {errors}")

# Attach DP engine
privacy_engine = PrivacyEngine()
model, optimizer, loader = privacy_engine.make_private_with_epsilon(
    module=model,
    optimizer=optimizer,
    data_loader=loader,
    epochs=5,
    target_epsilon=3.0,    # Privacy budget
    target_delta=1e-5,     # Failure probability
    max_grad_norm=1.0,     # Gradient clipping threshold
)

print(f"\nDP-SGD Configuration:")
print(f"  Target ε = 3.0, δ = 1e-5")
print(f"  Noise multiplier σ = {optimizer.noise_multiplier:.4f}")
print(f"  Max gradient norm C = 1.0")

# Training loop
for epoch in range(3):
    model.train()
    total_loss = 0
    for batch_x, batch_y in loader:
        optimizer.zero_grad()
        output = model(batch_x)
        loss = criterion(output, batch_y)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()

    epsilon = privacy_engine.get_epsilon(delta=1e-5)
    print(f"  Epoch {epoch+1}: loss={total_loss/len(loader):.4f}, ε spent={epsilon:.4f}")

print(f"\nFinal privacy guarantee: ε={privacy_engine.get_epsilon(1e-5):.3f}, δ=1e-5")

# ─── 2. Membership Inference Attack + Detection ───────────────
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import roc_auc_score

# Train a target model
X_data = np.random.randn(2000, 20)
y_data = (X_data[:, 0] > 0).astype(int)
X_train_mi, X_shadow, y_train_mi, y_shadow = train_test_split(X_data, y_data, test_size=0.5)
target_model = RandomForestClassifier(n_estimators=50, random_state=42)
target_model.fit(X_train_mi, y_train_mi)

# Simple membership inference: training samples have higher confidence
def mi_attack_confidence(model, X_members, X_non_members):
    """Naive MI attack: members → high confidence, non-members → lower confidence."""
    member_conf = model.predict_proba(X_members).max(axis=1)
    nonmember_conf = model.predict_proba(X_non_members).max(axis=1)

    all_conf = np.concatenate([member_conf, nonmember_conf])
    all_labels = np.concatenate([np.ones(len(X_members)), np.zeros(len(X_non_members))])
    auc = roc_auc_score(all_labels, all_conf)
    print(f"  MI Attack AUC: {auc:.4f} (0.5=random=perfectly private, 1.0=full membership leaked)")
    print(f"  Member avg conf:     {member_conf.mean():.4f}")
    print(f"  Non-member avg conf: {nonmember_conf.mean():.4f}")
    return auc

print("\nMembership Inference Attack:")
mi_attack_confidence(target_model, X_train_mi[:100], X_shadow[:100])

# ─── 3. Simulated Federated Learning ─────────────────────────
class FederatedServer:
    def __init__(self, global_model):
        self.global_weights = {k: v.clone() for k, v in global_model.state_dict().items()}

    def aggregate(self, client_updates, client_sizes):
        """FedAvg: weighted average of client model updates."""
        total = sum(client_sizes)
        new_weights = {}
        for key in self.global_weights:
            new_weights[key] = sum(
                (n / total) * updates[key]
                for updates, n in zip(client_updates, client_sizes)
            )
        self.global_weights = new_weights
        return self.global_weights

def client_train(global_weights, local_X, local_y, local_epochs=5, lr=0.01):
    """Train a local model on client data and return updated weights."""
    model = LinearClassifier()
    model.load_state_dict(global_weights)
    optimizer = optim.SGD(model.parameters(), lr=lr)
    criterion = nn.CrossEntropyLoss()
    dataset = TensorDataset(
        torch.FloatTensor(local_X),
        torch.LongTensor(local_y)
    )
    loader = DataLoader(dataset, batch_size=32, shuffle=True)
    for _ in range(local_epochs):
        for bx, by in loader:
            optimizer.zero_grad()
            loss = criterion(model(bx), by)
            loss.backward()
            optimizer.step()
    return {k: v.detach().clone() for k, v in model.state_dict().items()}, len(local_X)

# Simulate 3 clients (e.g., 3 hospitals with private patient data)
global_model = LinearClassifier()
server = FederatedServer(global_model)

client_data = [
    (np.random.randn(300, 20), (np.random.randn(300) > 0).astype(int)),
    (np.random.randn(200, 20), (np.random.randn(200) > 0).astype(int)),
    (np.random.randn(400, 20), (np.random.randn(400) > 0).astype(int)),
]

print("\nFederated Learning Simulation (3 clients, 3 rounds):")
for round_num in range(3):
    updates, sizes = [], []
    for i, (cx, cy) in enumerate(client_data):
        client_weights, n = client_train(server.global_weights, cx, cy)
        updates.append(client_weights)
        sizes.append(n)
    server.global_weights = server.aggregate(updates, sizes)
    print(f"  Round {round_num+1}: aggregated {sum(sizes)} samples from {len(client_data)} clients")

# ─── 4. PII Detection and Anonymization ──────────────────────
from presidio_analyzer import AnalyzerEngine
from presidio_anonymizer import AnonymizerEngine
from presidio_anonymizer.entities import OperatorConfig

analyzer = AnalyzerEngine()
anonymizer = AnonymizerEngine()

texts_with_pii = [
    "Patient John Smith (DOB: 15/03/1985) was diagnosed at john.smith@email.com. His SSN is 123-45-6789.",
    "Call Dr. Sarah Johnson at +1-555-0198 regarding claim #AB12345 for Mary Williams.",
    "Credit card 4532-1234-5678-9012 was used at 123 Main St, New York, NY 10001.",
]

print("\nPII Detection and Anonymization:")
for text in texts_with_pii:
    results = analyzer.analyze(
        text=text,
        entities=["PERSON", "EMAIL_ADDRESS", "PHONE_NUMBER", "US_SSN",
                  "CREDIT_CARD", "DATE_TIME", "LOCATION", "US_DRIVER_LICENSE"],
        language="en"
    )
    anonymized = anonymizer.anonymize(
        text=text,
        analyzer_results=results,
        operators={
            "PERSON": OperatorConfig("replace", {"new_value": "[PERSON]"}),
            "EMAIL_ADDRESS": OperatorConfig("replace", {"new_value": "[EMAIL]"}),
            "PHONE_NUMBER": OperatorConfig("replace", {"new_value": "[PHONE]"}),
            "US_SSN": OperatorConfig("replace", {"new_value": "[SSN]"}),
            "CREDIT_CARD": OperatorConfig("replace", {"new_value": "[CC]"}),
            "DATE_TIME": OperatorConfig("replace", {"new_value": "[DATE]"}),
            "LOCATION": OperatorConfig("replace", {"new_value": "[LOCATION]"}),
        }
    )
    print(f"\n  Original:    {text[:80]}")
    print(f"  Anonymized:  {anonymized.text[:80]}")
    detected = [(r.entity_type, text[r.start:r.end]) for r in results]
    print(f"  Detected PII: {detected}")