import base64
import json
import time
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
from cryptography.hazmat.primitives import serialization

def sign_request(private_key_pem: str, request_body: dict) -> dict:
    """Sign a request with Ed25519 private key."""

    # Load private key
    private_key = serialization.load_pem_private_key(
        private_key_pem.encode(),
        password=None
    )

    # Create signing payload
    timestamp = int(time.time())
    payload = {
        "body": json.dumps(request_body, sort_keys=True),
        "timestamp": timestamp
    }
    payload_bytes = json.dumps(payload, sort_keys=True).encode()

    # Sign
    signature = private_key.sign(payload_bytes)
    signature_b64 = base64.b64encode(signature).decode()

    # Return headers to include
    return {
        "X-Bastio-Agent-Id": agent_id,
        "X-Bastio-Signature": signature_b64,
        "X-Bastio-Timestamp": str(timestamp)
    }

# Usage
headers = sign_request(AGENT_PRIVATE_KEY, tool_call_body)
response = requests.post(
    f"https://api.bastio.com/v1/guard/{proxy_id}/tool",
    headers={
        "Authorization": f"Bearer {API_KEY}",
        **headers
    },
    json=tool_call_body
)

import * as crypto from 'crypto';

interface SignedHeaders {
  'X-Bastio-Agent-Id': string;
  'X-Bastio-Signature': string;
  'X-Bastio-Timestamp': string;
}

function signRequest(
  privateKeyPem: string,
  agentId: string,
  requestBody: object
): SignedHeaders {
  // Create signing payload
  const timestamp = Math.floor(Date.now() / 1000);
  const payload = JSON.stringify({
    body: JSON.stringify(requestBody),
    timestamp,
  });

  // Sign with Ed25519
  const privateKey = crypto.createPrivateKey(privateKeyPem);
  const signature = crypto.sign(null, Buffer.from(payload), privateKey);

  return {
    'X-Bastio-Agent-Id': agentId,
    'X-Bastio-Signature': signature.toString('base64'),
    'X-Bastio-Timestamp': timestamp.toString(),
  };
}

// Usage
const signedHeaders = signRequest(AGENT_PRIVATE_KEY, AGENT_ID, toolCallBody);
const response = await fetch(`https://api.bastio.com/v1/guard/${proxyId}/tool`, {
  method: 'POST',
  headers: {
    'Authorization': `Bearer ${API_KEY}`,
    'Content-Type': 'application/json',
    ...signedHeaders,
  },
  body: JSON.stringify(toolCallBody),
});

import os
import httpx
from dataclasses import dataclass

@dataclass
class BastioAgent:
    agent_id: str
    private_key: str
    proxy_id: str
    api_key: str

    async def validate_tool(self, session_id: str, tool_call: dict) -> dict:
        """Validate a tool call with agent authentication."""

        body = {
            "session_id": session_id,
            "tool_call": tool_call
        }

        # Sign request
        headers = sign_request(self.private_key, body)
        headers["Authorization"] = f"Bearer {self.api_key}"
        headers["Content-Type"] = "application/json"

        async with httpx.AsyncClient() as client:
            response = await client.post(
                f"https://api.bastio.com/v1/guard/{self.proxy_id}/tool",
                headers=headers,
                json=body
            )
            return response.json()

# Initialize agent
agent = BastioAgent(
    agent_id=os.environ["BASTIO_AGENT_ID"],
    private_key=os.environ["BASTIO_AGENT_PRIVATE_KEY"],
    proxy_id=os.environ["BASTIO_PROXY_ID"],
    api_key=os.environ["BASTIO_API_KEY"]
)

# Use in your agent loop
async def execute_tool_with_validation(session_id: str, tool_call: dict):
    result = await agent.validate_tool(session_id, tool_call)

    if result["action"] == "allow":
        return await execute_tool(tool_call)
    else:
        return f"Tool blocked: {result['message']}"

class BastioAgent {
  constructor(
    private agentId: string,
    private privateKey: string,
    private proxyId: string,
    private apiKey: string
  ) {}

  async validateTool(sessionId: string, toolCall: object): Promise<any> {
    const body = {
      session_id: sessionId,
      tool_call: toolCall,
    };

    const signedHeaders = signRequest(this.privateKey, this.agentId, body);

    const response = await fetch(
      `https://api.bastio.com/v1/guard/${this.proxyId}/tool`,
      {
        method: 'POST',
        headers: {
          'Authorization': `Bearer ${this.apiKey}`,
          'Content-Type': 'application/json',
          ...signedHeaders,
        },
        body: JSON.stringify(body),
      }
    );

    return response.json();
  }
}

// Initialize
const agent = new BastioAgent(
  process.env.BASTIO_AGENT_ID!,
  process.env.BASTIO_AGENT_PRIVATE_KEY!,
  process.env.BASTIO_PROXY_ID!,
  process.env.BASTIO_API_KEY!
);

// Use in your agent
async function executeWithValidation(sessionId: string, toolCall: object) {
  const result = await agent.validateTool(sessionId, toolCall);

  if (result.action === 'allow') {
    return await executeTool(toolCall);
  } else {
    return `Tool blocked: ${result.message}`;
  }
}