Secret Rotation

from agentsecretstore import AgentVault

async def rotate_openai_key(new_key: str):
    """Manual rotation: update a secret with a new value."""
    async with AgentVault() as vault:
        # Update creates a new version automatically
        updated = await vault.update_secret(
            path="production/openai/api-key",
            value=new_key,
            # Optional: keep old version valid for 15 minutes (grace period)
            grace_period_seconds=900,
        )
        print(f"Rotated to version: {updated.version}")
        print(f"Old version valid until: {updated.grace_expires_at}")

# Rotate via REST API
curl -X PATCH https://api.agentsecretstore.com/v1/secrets/production%2Fopenai%2Fapi-key \
  -H "Authorization: Bearer $ASS_AGENT_KEY" \
  -H "Content-Type: application/json" \
  -d '{
    "value": "sk-proj-new-key-here",
    "grace_period_seconds": 900
  }'

# Common cron expressions for rotation schedules

# Every 24 hours at midnight UTC
0 0 * * *

# Every Sunday at 2 AM UTC (weekly)
0 2 * * 0

# 1st of every month at midnight (monthly)
0 0 1 * *

# Every 90 days (quarterly — use a cron job that fires daily, checks 90-day mark)
0 3 * * *   # fires daily; rotation logic checks elapsed time

# Every 6 hours (for high-rotation environments)
0 */6 * * *

from agentsecretstore import AgentVault
from agentsecretstore.rotation import RotationSchedule, RotationStrategy

async def configure_scheduled_rotation():
    async with AgentVault() as vault:
        # Rotate every 30 days, generate a new key via the provider
        await vault.rotation.create_schedule(
            path="production/openai/api-key",
            schedule=RotationSchedule(
                cron="0 2 * * 0",          # Every Sunday at 2 AM UTC
                strategy=RotationStrategy.PROVIDER_WEBHOOK,
                provider="openai",
                grace_period_seconds=1800, # Old key valid 30 min after rotation
            ),
        )

        # Monthly rotation (1st of each month at midnight)
        await vault.rotation.create_schedule(
            path="production/stripe/restricted-key",
            schedule=RotationSchedule(
                cron="0 0 1 * *",
                strategy=RotationStrategy.EXTERNAL_WEBHOOK,
                webhook_url="https://your-rotation-server.com/rotate/stripe",
                grace_period_seconds=3600,
            ),
        )

        # List all active schedules
        schedules = await vault.rotation.list_schedules()
        for s in schedules:
            print(f"{s.path}: {s.cron} (next: {s.next_run})")

# rotation_server.py — Receives rotation webhooks from Agent Secret Store
# and handles the provider-side key rotation

import hmac
import hashlib
import os
import asyncio
from fastapi import FastAPI, Request, HTTPException
import httpx

app = FastAPI()

WEBHOOK_SECRET = os.environ["ASS_WEBHOOK_SECRET"]

def verify_signature(payload: bytes, signature: str) -> bool:
    """Verify Agent Secret Store webhook signature."""
    expected = hmac.new(
        WEBHOOK_SECRET.encode(),
        payload,
        hashlib.sha256,
    ).hexdigest()
    return hmac.compare_digest(f"sha256={expected}", signature)

@app.post("/rotate/openai")
async def rotate_openai_key(request: Request):
    """Handle OpenAI key rotation webhook."""
    payload = await request.body()
    sig = request.headers.get("X-ASS-Signature", "")

    if not verify_signature(payload, sig):
        raise HTTPException(status_code=401, detail="Invalid signature")

    event = await request.json()

    if event["event"] != "rotation.requested":
        return {"status": "ignored"}

    # Step 1: Create a new key at the provider
    async with httpx.AsyncClient() as client:
        resp = await client.post(
            "https://api.openai.com/v1/api_keys",
            headers={"Authorization": f"Bearer {os.environ['OPENAI_ADMIN_KEY']}"},
            json={"name": f"rotated-{event['rotation_id']}"},
        )
        new_key = resp.json()["key"]

    # Step 2: Confirm rotation — vault updates the secret with the new value
    # and starts the grace period
    from agentsecretstore import AgentVault
    async with AgentVault() as vault:
        await vault.rotation.confirm(
            rotation_id=event["rotation_id"],
            new_value=new_key,
        )

    # Step 3: Delete the old key at the provider AFTER grace period
    # (vault fires a rotation.grace_period_ended webhook)
    return {"status": "rotated", "rotation_id": event["rotation_id"]}

@app.post("/rotation/grace-ended")
async def grace_period_ended(request: Request):
    """Called when the grace period expires — safe to delete old key."""
    payload = await request.body()
    sig = request.headers.get("X-ASS-Signature", "")

    if not verify_signature(payload, sig):
        raise HTTPException(status_code=401, detail="Invalid signature")

    event = await request.json()
    old_key_id = event["metadata"]["old_key_id"]

    # Revoke the old key at the provider
    async with httpx.AsyncClient() as client:
        await client.delete(
            f"https://api.openai.com/v1/api_keys/{old_key_id}",
            headers={"Authorization": f"Bearer {os.environ['OPENAI_ADMIN_KEY']}"},
        )

    return {"status": "old_key_deleted"}

from agentsecretstore import AgentVault
from agentsecretstore.webhooks import WebhookServer
import os

WEBHOOK_SECRET = os.environ["ASS_WEBHOOK_SECRET"]

async def handle_rotation(event: dict):
    """
    Called when a secret is rotated.
    Agent should invalidate its cached copy and re-fetch.
    """
    if event["event"] in ("secret.rotated", "rotation.grace_period_ended"):
        path = event["resource_path"]
        print(f"Secret rotated: {path}")

        # Invalidate local cache entry
        secret_cache.delete(path)

        # Optionally: pre-warm cache with new value
        async with AgentVault() as vault:
            new_secret = await vault.get_secret(path)
            secret_cache.set(path, new_secret.value, ttl=300)

        print(f"Cache refreshed for {path}")

server = WebhookServer(secret=WEBHOOK_SECRET)
server.on("secret.rotated", handle_rotation)
server.on("rotation.grace_period_ended", handle_rotation)
await server.listen(port=8080)

from agentsecretstore import AgentVault

async def rollback_bad_rotation():
    """
    Rollback to previous version if the new credential is broken.
    The grace period must still be active OR the old version must still exist.
    """
    async with AgentVault() as vault:
        # List versions to find the target
        versions = await vault.list_versions("production/openai/api-key")
        for v in versions:
            print(f"  {v.version}: active={v.is_active}, created={v.created_at}")

        # Rollback to the previous version
        await vault.rollback_secret(
            path="production/openai/api-key",
            version="v2",
        )
        print("Rolled back to v2 — investigate why v3 was broken")