OAuth Token Management

# Manual OAuth Token Management (for non-SDK users)

import requests
import time
import json
from typing import Optional, Dict
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import ed25519

class ManualAIMClient:
    """Manual implementation of AIM OAuth token management"""

    def __init__(self, agent_name: str, private_key_path: str):
        self.agent_name = agent_name
        self.base_url = "https://api.aim.example.com"
        self.access_token = None
        self.refresh_token = None
        self.expires_at = None

        # Load Ed25519 private key for signing
        with open(private_key_path, 'rb') as f:
            self.private_key = serialization.load_pem_private_key(
                f.read(), password=None
            )

    def register_agent(self) -> Dict:
        """Initial agent registration with signed request"""
        # Create registration payload
        payload = {
            "name": self.agent_name,
            "public_key": self._get_public_key(),
            "capabilities": ["read", "write", "execute"],
            "timestamp": int(time.time())
        }

        # Sign the payload
        signature = self._sign_payload(payload)

        # Register agent
        response = requests.post(
            f"{self.base_url}/v1/agents/register",
            json=payload,
            headers={
                "X-Agent-Signature": signature
            }
        )

        if response.status_code == 201:
            data = response.json()
            self.access_token = data["access_token"]
            self.refresh_token = data["refresh_token"]
            self.expires_at = time.time() + data["expires_in"]
            return data
        else:
            raise Exception(f"Registration failed: {response.text}")

    def refresh_tokens(self) -> Dict:
        """Refresh expired access token using refresh token"""
        if not self.refresh_token:
            raise Exception("No refresh token available")

        response = requests.post(
            f"{self.base_url}/v1/auth/refresh",
            json={
                "grant_type": "refresh_token",
                "refresh_token": self.refresh_token,
                "client_id": self.agent_name
            }
        )

        if response.status_code == 200:
            data = response.json()
            # Update tokens (rotation)
            self.access_token = data["access_token"]
            self.refresh_token = data["refresh_token"]  # New refresh token!
            self.expires_at = time.time() + data["expires_in"]
            return data
        else:
            # Refresh token expired or revoked
            # Need to re-register the agent
            return self.register_agent()

    def make_api_call(self, endpoint: str, method: str = "GET", **kwargs) -> Dict:
        """Make an API call with automatic token refresh"""
        # Check if token needs refresh
        if self._needs_refresh():
            self.refresh_tokens()

        # Make the API call
        response = requests.request(
            method,
            f"{self.base_url}{endpoint}",
            headers={
                "Authorization": f"Bearer {self.access_token}"
            },
            **kwargs
        )

        # Handle 401 Unauthorized (token expired)
        if response.status_code == 401:
            self.refresh_tokens()
            # Retry the request
            response = requests.request(
                method,
                f"{self.base_url}{endpoint}",
                headers={
                    "Authorization": f"Bearer {self.access_token}"
                },
                **kwargs
            )

        return response.json()

    def _needs_refresh(self) -> bool:
        """Check if access token needs refresh (5 min buffer)"""
        if not self.expires_at:
            return True
        return time.time() > (self.expires_at - 300)

    def _get_public_key(self) -> str:
        """Get public key in PEM format"""
        public_key = self.private_key.public_key()
        return public_key.public_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PublicFormat.SubjectPublicKeyInfo
        ).decode('utf-8')

    def _sign_payload(self, payload: Dict) -> str:
        """Sign payload with Ed25519 private key"""
        message = json.dumps(payload, sort_keys=True).encode()
        signature = self.private_key.sign(message)
        return signature.hex()

# Usage example
client = ManualAIMClient("my-agent", "/path/to/private_key.pem")
client.register_agent()

# Make API calls - tokens refresh automatically
agents = client.make_api_call("/v1/agents")
trust_score = client.make_api_call("/v1/agents/my-agent/trust")