diff --git a/agentic_security/http_spec.py b/agentic_security/http_spec.py index dcb938d..2ef594e 100644 --- a/agentic_security/http_spec.py +++ b/agentic_security/http_spec.py @@ -18,16 +18,34 @@ class Modality(Enum): def encode_image_base64_by_url(url: str = "https://github.com/fluidicon.png") -> str: """Encode image data to base64 from a URL""" - response = httpx.get(url) - encoded_content = base64.b64encode(response.content).decode("utf-8") - return "data:image/jpeg;base64," + encoded_content + try: + response = httpx.get(url) + encoded_content = base64.b64encode(response.content).decode("utf-8") + return "data:image/jpeg;base64," + encoded_content + except httpx.TimeoutException as e: + raise ValueError(f"Image fetch timed out: {str(e)}") + except httpx.HTTPStatusError as e: + raise ValueError(f"HTTP error {e.response.status_code}: {str(e)}") + except httpx.RequestError as e: + raise ValueError(f"Image fetch failed: {str(e)}") + except Exception as e: + raise ValueError(f"Unexpected error fetching image: {str(e)}") def encode_audio_base64_by_url(url: str) -> str: """Encode audio data to base64 from a URL""" - response = httpx.get(url) - encoded_content = base64.b64encode(response.content).decode("utf-8") - return "data:audio/mpeg;base64," + encoded_content + try: + response = httpx.get(url) + encoded_content = base64.b64encode(response.content).decode("utf-8") + return "data:audio/mpeg;base64," + encoded_content + except httpx.TimeoutException as e: + raise ValueError(f"Audio fetch timed out: {str(e)}") + except httpx.HTTPStatusError as e: + raise ValueError(f"HTTP error {e.response.status_code}: {str(e)}") + except httpx.RequestError as e: + raise ValueError(f"Audio fetch failed: {str(e)}") + except Exception as e: + raise ValueError(f"Unexpected error fetching audio: {str(e)}") class InvalidHTTPSpecError(Exception): @@ -58,16 +76,25 @@ def timeout(self): async def _probe_with_files(self, files): transport = httpx.AsyncHTTPTransport(retries=settings_var("network.retry", 3)) - async with httpx.AsyncClient(transport=transport) as client: - response = await client.request( - method=self.method, - url=self.url, - headers=self.headers, - files=files, - timeout=self.timeout(), - ) - - return response + try: + async with httpx.AsyncClient(transport=transport) as client: + response = await client.request( + method=self.method, + url=self.url, + headers=self.headers, + files=files, + timeout=self.timeout(), + ) + response.raise_for_status() + return response + except httpx.TimeoutException as e: + raise ValueError(f"Request timed out: {str(e)}") + except httpx.HTTPStatusError as e: + raise ValueError(f"HTTP error {e.response.status_code}: {str(e)}") + except httpx.RequestError as e: + raise ValueError(f"Request failed: {str(e)}") + except Exception as e: + raise ValueError(f"Unexpected error: {str(e)}") def validate(self, prompt, encoded_image, encoded_audio, files) -> None: if self.has_files and not files: @@ -102,16 +129,26 @@ async def probe( content = content.replace("<>", encoded_audio) transport = httpx.AsyncHTTPTransport(retries=settings_var("network.retry", 3)) - async with httpx.AsyncClient(transport=transport) as client: - response = await client.request( - method=self.method, - url=self.url, - headers=self.headers, - content=content, - timeout=self.timeout(), - ) - - return response + try: + async with httpx.AsyncClient(transport=transport) as client: + response = await client.request( + method=self.method, + url=self.url, + headers=self.headers, + content=content, + timeout=self.timeout(), + ) + response.raise_for_status() + return response + + except httpx.TimeoutException as e: + raise ValueError(f"Request timed out: {str(e)}") + except httpx.HTTPStatusError as e: + raise ValueError(f"HTTP error {e.response.status_code}: {str(e)}") + except httpx.RequestError as e: + raise ValueError(f"Request failed: {str(e)}") + except Exception as e: + raise ValueError(f"Unexpected error: {str(e)}") async def verify(self) -> httpx.Response: match self: diff --git a/agentic_security/mcp/main.py b/agentic_security/mcp/main.py index 0e754dd..5397fb8 100644 --- a/agentic_security/mcp/main.py +++ b/agentic_security/mcp/main.py @@ -16,9 +16,19 @@ async def verify_llm(spec: str) -> dict: """Verify an LLM model specification using the FastAPI server.""" url = f"{AGENTIC_SECURITY}/verify" - async with httpx.AsyncClient() as client: - response = await client.post(url, json={"spec": spec}) - return response.json() + try: + async with httpx.AsyncClient(timeout=10.0) as client: + response = await client.post(url, json={"spec": spec}) + response.raise_for_status() + return response.json() + except httpx.TimeoutException as e: + return {"error": f"Request timed out: {str(e)}"} + except httpx.HTTPStatusError as e: + return {"error": f"HTTP error {e.response.status_code}: {str(e)}"} + except httpx.RequestError as e: + return {"error": f"Request failed: {str(e)}"} + except Exception as e: + return {"error": f"Unexpected error: {str(e)}"} @mcp.tool() @@ -39,36 +49,77 @@ async def start_scan( "probe_datasets": [], "secrets": {}, } - async with httpx.AsyncClient() as client: - response = await client.post(url, json=payload) - return response.json() + try: + async with httpx.AsyncClient(timeout=10.0) as client: + response = await client.post(url, json=payload) + response.raise_for_status() + return response.json() + except httpx.TimeoutException as e: + return {"error": f"Request timed out: {str(e)}"} + except httpx.HTTPStatusError as e: + return {"error": f"HTTP error {e.response.status_code}: {str(e)}"} + except httpx.RequestError as e: + return {"error": f"Request failed: {str(e)}"} + except Exception as e: + return {"error": f"Unexpected error: {str(e)}"} @mcp.tool() async def stop_scan() -> dict: """Stop an ongoing scan via the FastAPI server.""" url = f"{AGENTIC_SECURITY}/stop" - async with httpx.AsyncClient() as client: - response = await client.post(url) - return response.json() + try: + async with httpx.AsyncClient(timeout=10.0) as client: + response = await client.post(url) + response.raise_for_status() + return response.json() + except httpx.TimeoutException as e: + return {"error": f"Request timed out: {str(e)}"} + except httpx.HTTPStatusError as e: + return {"error": f"HTTP error {e.response.status_code}: {str(e)}"} + except httpx.RequestError as e: + return {"error": f"Request failed: {str(e)}"} + except Exception as e: + return {"error": f"Unexpected error: {str(e)}"} @mcp.tool() async def get_data_config() -> list: """Retrieve data configuration from the FastAPI server.""" url = f"{AGENTIC_SECURITY}/v1/data-config" - async with httpx.AsyncClient() as client: - response = await client.get(url) - return response.json() + try: + async with httpx.AsyncClient(timeout=10.0) as client: + response = await client.get(url) + response.raise_for_status() + return response.json() + except httpx.TimeoutException as e: + return {"error": f"Request timed out: {str(e)}"} + except httpx.HTTPStatusError as e: + return {"error": f"HTTP error {e.response.status_code}: {str(e)}"} + except httpx.RequestError as e: + return {"error": f"Request failed: {str(e)}"} + except Exception as e: + return {"error": f"Unexpected error: {str(e)}"} @mcp.tool() async def get_spec_templates() -> list: """Retrieve data configuration from the FastAPI server.""" url = f"{AGENTIC_SECURITY}/v1/llm-specs" - async with httpx.AsyncClient() as client: - response = await client.get(url) - return response.json() + try: + async with httpx.AsyncClient(timeout=10.0) as client: + response = await client.get(url) + response.raise_for_status() + return response.json() + + except httpx.TimeoutException as e: + return {"error": f"Request timed out: {str(e)}"} + except httpx.HTTPStatusError as e: + return {"error": f"HTTP error {e.response.status_code}: {str(e)}"} + except httpx.RequestError as e: + return {"error": f"Request failed: {str(e)}"} + except Exception as e: + return {"error": f"Unexpected error: {str(e)}"} # Run the MCP server diff --git a/pip b/pip new file mode 100644 index 0000000..e69de29