Metadata-Version: 2.4
Name: agentshield-sdk
Version: 0.1.2
Summary: Python SDK for the AgentShield API - Runtime security for agentic AI applications.
Author-email: AgentShield Admin <hello@sanjayck.com>
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Topic :: Security
Classifier: Topic :: Software Development :: Libraries :: Python Modules
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.8
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Requires-Python: >=3.8
Description-Content-Type: text/markdown
Requires-Dist: httpx>=0.20.0
Requires-Dist: pydantic>=2.0

# AgentShield SDK

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

```bash
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:**

➡️ **hello@sanjayck.com** ⬅️

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:

```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 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](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:* 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.

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