agentshield-sdk 0.2.1

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

AgentShield SDK

Python SDK and CLI 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, including its API key for authentication.
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.
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

The agentshield-sdk package includes both the Python SDK for your agent code and the agentshield command-line tool for managing your user account and API keys.

Install the SDK using pip:

pip install agentshield-sdk

2. Account Registration & API Key Generation (Using the CLI)

AgentShield requires an API key to authenticate requests from your agent to the security service. You can now manage your account and generate API keys yourself using the included agentshield CLI tool.

(a) Register an Account:

Open your terminal and run the register command. You will be prompted for a username, email, and password (min 8 characters).

agentshield register

(b) Log In:

Once registered, log in to your account. This will save a session token locally (typically in ~/.agentshield/config.json) allowing you to manage keys.

agentshield login

(c) Generate an API Key:

Now you can generate one or more API keys for your agents. You can optionally add a description to help you remember what each key is for.

# Example with description
agentshield keys create -d "My Test Agent v1 Key"

# Example without description
agentshield keys create

---> IMPORTANT: <--- The command will output your new API key (e.g., ask_...). Copy this key immediately and store it securely (e.g., in an environment variable or secret manager). You cannot retrieve the full key again after this step. You can list your keys later (agentshield keys list) to see their ID, prefix, and description, but not the full key.

3. Using the SDK in Your Agent

Now that you have generated an API key via the CLI, you can use it to initialize the AgentShield client in your Python agent code.

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 via the 'agentshield keys create' command
# Best practice: Load key from environment variable or secure storage
# Example: export AGENTSHEILD_API_KEY="ask_YOUR_GENERATED_KEY_HERE"
api_key = os.environ.get("AGENTSHEILD_API_KEY")

# The URL where the AgentShield API service is deployed
# Use the official public URL unless self-hosting
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 not api_key:
    print("ERROR: Please set the AGENTSHEILD_API_KEY environment variable.")
    # In a real agent, handle missing key appropriately
    exit()

# --- Initialize Client ---
# You only need to do this once for your agent instance
try:
    shield = AgentShield(
        api_key=api_key,
        service_url=service_url,
        agent_id=agent_id,
        timeout=15.0 # Optional: Default timeout for calls TO AgentShield (seconds)
    )
except ValueError as e:
    print(f"ERROR: Configuration error initializing AgentShield: {e}")
    # Handle configuration errors
    exit()

# --- Example: Guarded GET request ---
async def make_guarded_request():
    url_to_check = "[https://httpbin.org/get](https://httpbin.org/get)" # A test URL
    print(f"\nChecking if GET request to {url_to_check} is allowed...")

    try:
        # This method first calls AgentShield's /evaluate endpoint.
        # If allowed, it then proceeds to make the actual GET request using httpx.
        response = await shield.guarded_get(
            url=url_to_check,
            headers={"X-Custom-Header": "MyAgent"}, # 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 the target
        # 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}")


# --- Example: Safe Execute Check ---
async def check_code_execution():
    safe_code = "print('This is safe code.')\nx = 1 + 2"
    dangerous_code = "import os\nos.system('echo potentially harmful')"

    print(f"\nChecking if safe code is allowed...")
    try:
        # This calls /evaluate but DOES NOT run the code
        await shield.safe_execute(code_snippet=safe_code)
        print("Check ALLOWED for safe code. (Agent would execute it now)")
        # Agent would proceed to execute 'safe_code' here using its own method
    except AgentShieldError as e:
        print(f"Check BLOCKED for safe code unexpectedly! Reason: {e} (Details: {e.policy_details})")
    except Exception as e:
        print(f"An unexpected error occurred: {e}")


    print(f"\nChecking if dangerous code is allowed...")
    try:
        # This calls /evaluate but DOES NOT run the code
        await shield.safe_execute(code_snippet=dangerous_code)
        print("Check ALLOWED for dangerous code unexpectedly! (Agent should NOT execute this)")
    except AgentShieldError as e:
        print(f"Check BLOCKED for dangerous code as expected. Reason: {e} (Details: {e.policy_details})")
        # Agent should NOT execute 'dangerous_code'
    except Exception as e:
        print(f"An unexpected error occurred: {e}")

# Run the examples
async def run_examples():
    await make_guarded_request()
    await check_code_execution()

if __name__ == "__main__":
    # Ensure you have a valid API key set before running the test
    if api_key:
        asyncio.run(run_examples())
    else:
        print("Skipping example run: API key not configured.")

CLI Tool Reference (agentshield)

The agentshield-sdk package includes a command-line tool, agentshield, for managing your user account and API keys required by the SDK.

Configuration:

• The CLI interacts with the deployed AgentShield API service. By default, it uses https://agentshield-api-service-338863748406.us-central1.run.app.

• Login sessions are stored locally (typically ~/.agentshield/config.json).

Commands:

• agentshield --help: Shows the main help message and lists all commands.
• agentshield --version: Shows the installed package version.

Account Management:

• agentshield register: Register a new user account. Prompts for username, email, and password.

agentshield register

• agentshield login: Log in to your account. Prompts for email and password. Saves a session token locally.

agentshield login

• agentshield logout: Log out by deleting the locally stored session token.

agentshield logout

API Key Management (Requires Login):

• agentshield keys list: List all your existing API keys, showing their ID, prefix, status, creation date, and description.

agentshield keys list

• agentshield keys create [-d DESCRIPTION]: Generate a new API key.
  • Use the optional -d or --description flag to add a description (e.g., "Key for production agent").
  • Important: The full API key is shown only once upon creation. Copy and save it securely.

# Create with description
agentshield keys create -d "My Production Agent Key"

# Create without description
agentshield keys create

• agentshield keys delete <KEY_ID>: Permanently delete an API key using its numeric ID (obtained from keys list). You will be asked for confirmation.

agentshield keys delete 123

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: Any policies added, removed, or modifed will be updated in this documentation)

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