agentshield-sdk 0.1.2

Python SDK for the AgentShield API - Runtime security for agentic AI applications.

AgentShield SDK

Python SDK for the AgentShield API - Runtime security for agentic AI applications.

What is AgentShield and When is it Useful?

Modern AI agents often need to interact with the outside world (like calling APIs or running code) to complete tasks. However, allowing an AI agent unrestricted ability to perform these actions can be risky. It might:

• Call malicious or unintended APIs.
• Leak sensitive data through API calls.
• Execute harmful code generated by itself or based on malicious input.
• Access internal systems or cloud metadata it shouldn't (SSRF attacks).
• Run commands that damage the system.

AgentShield acts as a safety checkpoint for your AI agent.

It's useful whenever your agent is about to perform an action that could potentially carry risk. Specifically:

1. Before Making External API Calls: If your agent generates a URL and intends to make an HTTP request (GET, POST, etc.), AgentShield can check if that URL is allowed based on security policies (e.g., block internal IPs, block non-HTTPS, block known malicious domains).
2. Before Executing Generated Code: If your agent generates a snippet of code (e.g., Python code) and intends to execute it, AgentShield can check if that code contains potentially dangerous patterns (e.g., importing os, using eval, accessing sensitive files) based on security policies.

Think of it like a security guard asking "Are you sure you should do that?" before the agent takes a potentially dangerous step.

How it Works (The Checkpoint Process)

AgentShield works proactively, meaning it performs its check before the risky action happens. Here's the flow when using the SDK:

1. Intention: Your agent decides it needs to perform an action (e.g., "I need to GET data from https://api.example.com/data" or "I need to run this code: import os; os.remove('file')").
2. SDK Call: Instead of directly making the API call or running the code, your agent calls an AgentShield SDK function:
   • shield.guarded_get(url=...) for API calls.
   • shield.safe_execute(code_snippet=...) for code execution checks.
3. Evaluation Request: The SDK takes the input provided (the url string for guarded_get or the code_snippet string for safe_execute) and sends it to the deployed AgentShield API's /evaluate endpoint.
4. Policy Check: The AgentShield API service receives the intended action details. It compares the provided URL or code snippet against its list of currently active security policies (see below).
5. Decision: Based on the policies, the API makes a decision: "allow" or "block".
6. Response to SDK: The API sends this decision back to the SDK.
7. Action Control:
   • If the decision was "allow": The SDK function completes successfully. For guarded_get, it then proceeds to make the actual HTTP request to the target URL. For safe_execute, it simply returns, indicating the check passed (it doesn't run the code itself).
   • If the decision was "block": The SDK function raises an AgentShieldError, immediately stopping the process. The risky API call is not made, and the code check indicates it should not be executed. The error includes details about which policy caused the block.

This entire check happens between the agent's intention to act and the actual execution of that action.

Getting Started

1. Installation

Install the SDK using pip:

pip install agentshield-sdk

2. Get an API Key

AgentShield requires an API key for authenticating requests to its /evaluate endpoint.

To obtain your API key, please contact the AgentShield administrator:

➡️ [e-mail] ⬅️

You will receive an API key string (likely starting with ask_...). Keep this key secure, as it cannot be retrieved again.

3. Basic Usage

Here's a minimal example of how to configure the client and use it to guard an API call:

import os
import asyncio
import httpx  # SDK uses httpx; import it if handling its specific exceptions
from agentshield_sdk.client import AgentShield, AgentShieldError

# --- Configuration ---

# Your unique API Key obtained from the administrator
# Best practice: Load key from environment variable or secure storage
api_key = os.environ.get("AGENTSHEILD_API_KEY", "YOUR_API_KEY_HERE")

# The URL where the AgentShield API service is deployed
# Use the official public URL provided
service_url = "[https://agentshield-api-service-338863748406.us-central1.run.app](https://agentshield-api-service-338863748406.us-central1.run.app)"

# Optional: An identifier for the agent using the SDK
agent_id = "my_example_agent_v1"

if api_key == "YOUR_API_KEY_HERE":
    print("ERROR: Please set the AGENTSHEILD_API_KEY environment variable or replace the placeholder in the script.")
    exit()

# --- Initialize Client ---
# You only need to do this once for your agent instance
shield = AgentShield(
    api_key=api_key,
    service_url=service_url,
    agent_id=agent_id,
    timeout=15.0 # Optional: Default timeout for API calls to AgentShield (seconds)
)

# --- Example: Guarded GET request ---
async def make_guarded_request():
    url_to_check = "[https://api.thirdparty.com/some/data](https://api.thirdparty.com/some/data)" # The URL your agent wants to call
    print(f"Checking if GET request to {url_to_check} is allowed...")

    try:
        # This method first calls AgentShield's /evaluate endpoint.
        # If allowed, it then performs the actual GET request using httpx.
        response = await shield.guarded_get(
            url_to_check,
            headers={"Authorization": "Bearer SOME_OTHER_TOKEN"}, # Example: Headers for the *target* URL
            timeout=20.0 # Optional: Timeout for the call to the *target* URL
        )

        # If we reach here, the request was allowed by AgentShield
        print(f"Request allowed by AgentShield! Status Code from target: {response.status_code}")

        # Process the actual response from api.thirdparty.com
        # data = response.json()
        # print(f"Data received: {data}")

    except AgentShieldError as e:
        # If blocked by AgentShield policy or API communication fails
        print(f"Request BLOCKED or SDK Error! Reason: {e} (Details: {e.policy_details})")
        # Agent should NOT proceed with the request
    except ValueError as e:
        # Handle invalid input like bad URLs passed to the SDK
        print(f"Input Error: {e}")
    except httpx.TimeoutException:
        # Handle timeout contacting the *target* URL (after allow)
        print(f"Request to {url_to_check} timed out.")
    except httpx.RequestError as e:
        # Handle network errors contacting the *target* URL (after allow)
        print(f"Network error contacting {url_to_check}: {e}")
    except Exception as e:
        # Handle any other unexpected errors
        print(f"An unexpected error occurred: {e}")

# Run the example
# Note: This requires an event loop running, common in async applications.
# If running standalone, use asyncio.run()
if __name__ == "__main__":
    # Ensure you have a valid API key set before running the test
    if api_key != "YOUR_API_KEY_HERE":
        asyncio.run(make_guarded_request())
    else:
        print("Skipping example run: API key not configured.")

Default Security Policies (Active)

AgentShield comes with a set of default blocking policies enabled to provide baseline protection. These checks are performed before the agent executes an API call (guarded_get) or runs code (safe_execute).

Here's a summary of what the currently active default policies detect and block, with links to relevant OWASP Top 10 for LLM Applications items:

Code Execution Risks Blocked by Default:

• OS Interaction: Blocks code importing the os module or using os.system/os.popen. (LLM05, LLM06)
• Subprocess Usage: Blocks code importing or using the subprocess module. (LLM05, LLM06)
• Risky File Operations: Blocks code using shutil or specific os file deletion/modification functions. (LLM06)
• Socket Operations: Blocks code importing or using the socket module. (LLM06, LLM03)
• Sensitive File Access: Blocks code attempting to open() common sensitive file paths (like /etc/passwd, SSH keys, cloud credentials, .env). (LLM02, LLM06)
• Command Injection Patterns: Blocks code containing basic command chaining (; or && followed by risky commands like rm, wget, curl). (LLM01, LLM05, LLM06)
• Hardcoded Credentials: Blocks code containing patterns resembling database connection strings with passwords or common secret variable assignments. (LLM02)
• eval() Usage: Blocks code using the eval() function. (LLM05, LLM06)
• exec() Usage: Blocks code using the exec() function. (LLM05, LLM06)
• pickle Usage: Blocks code importing or using pickle.load/loads to prevent unsafe deserialization. (LLM03, LLM04)

API Call Risks Blocked by Default:

• Localhost/Metadata Access (SSRF): Blocks API calls targeting localhost, 127.0.0.1, or common cloud provider metadata service endpoints (GCP, AWS, Azure). (LLM06)
• Private IP Access (SSRF): Blocks API calls targeting common private IP ranges (e.g., 10.x.x.x, 192.168.x.x). (LLM06)
• Plain HTTP URLs: Blocks API calls using unencrypted http://. (LLM02, LLM06)

(Note: A default policy attempting to block URL path traversal (../) was temporarily disabled due to matching issues and may be revisited later. The administrator can also add, remove, or modify policies.)

SDK Client Reference

Initialization

Instantiate the client with your API key and the service URL.

from agentshield_sdk.client import AgentShield

shield = AgentShield(  
    api_key: str,          # Your mandatory API key from the admin  
    service_url: str,      # Mandatory URL of the AgentShield API service  
    agent_id: str = "sdk_agent_default", # Optional identifier for your agent  
    timeout: float = 10.0  # Optional default timeout for calls TO the AgentShield API itself  
)

async guarded_get(url: str, **kwargs) -> httpx.Response

Checks if a GET request to the specified url is permitted by AgentShield policies before executing it using httpx.

• Parameters:
  • url (str): The target URL the agent intends to call.
  • **kwargs: Any additional keyword arguments accepted by httpx.get (e.g., headers, params, timeout). The timeout kwarg here applies to the request to the target URL.
• Returns: An httpx.Response object from the target URL if the request is allowed by AgentShield and the request to the target is successful.
• Raises:
  • AgentShieldError: If the action is blocked by an AgentShield policy or if communication with the AgentShield API fails (check policy_details attribute for more info).
  • ValueError: If the provided url format is invalid.
  • httpx.TimeoutException: If the request to the target URL times out.
  • httpx.RequestError: For other network errors contacting the target URL (after being allowed).

async safe_execute(code_snippet: str)

Checks if executing a Python code_snippet (string) is permitted by AgentShield policies.

• Warning: This SDK function DOES NOT EXECUTE the code for safety reasons. It only performs the security check against the AgentShield API. The calling agent is responsible for execution after this check passes.
• Parameters:
  • code_snippet (str): The Python code the agent intends to execute.
• Returns: None if the check indicates the code is allowed.
• Raises:
  • AgentShieldError: If code execution is blocked by an AgentShield policy or if communication with the Agent