agentshield-sdk 0.2.1
Python SDK for the AgentShield API - Runtime security for agentic AI applications.
Requires Python: >=3.8

# AgentShield SDK

[![PyPI version](https://badge.fury.io/py/agentshield-sdk.svg)](https://badge.fury.io/py/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:

```bash
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).

```bash
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.

```bash
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.

```bash
# 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.

```python
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.
```bash
agentshield register
```

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

```bash
agentshield login
```

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

```bash
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.
```bash
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.

```bash
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](https://owasp.org/www-project-top-10-for-large-language-model-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.

```python
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