diff --git a/src/cai/tools/reconnaissance/c99.py b/src/cai/tools/reconnaissance/c99.py index 724dd278..65da70f3 100644 --- a/src/cai/tools/reconnaissance/c99.py +++ b/src/cai/tools/reconnaissance/c99.py @@ -18,6 +18,7 @@ Environment: """ import json +import logging import os from typing import Any, Dict, List, Optional, Union, Literal @@ -26,6 +27,8 @@ from dotenv import load_dotenv from cai.sdk.agents import function_tool +logger = logging.getLogger(__name__) + JSONType = Union[Dict[str, Any], List[Any], str, int, float, bool, None] @@ -61,16 +64,25 @@ def _call_c99_api(endpoint: str, params: Dict[str, Any]) -> Optional[JSONType]: try: response = requests.get(base_url, params=query_params, timeout=60) - except Exception: # pylint: disable=broad-except + response.raise_for_status() + except requests.exceptions.Timeout: + logger.warning("C99 API request timed out for endpoint: %s", endpoint) return None - - if response.status_code != 200: + except requests.exceptions.ConnectionError: + logger.warning("C99 API connection error for endpoint: %s", endpoint) + return None + except requests.exceptions.HTTPError as e: + logger.warning("C99 API HTTP error for endpoint %s: %s", endpoint, e) + return None + except requests.exceptions.RequestException as e: + logger.warning("C99 API request failed for endpoint %s: %s", endpoint, e) return None try: return response.json() - except Exception: # pylint: disable=broad-except + except ValueError: # If JSON parsing fails, fall back to raw text. + logger.debug("C99 API response is not valid JSON for endpoint: %s, falling back to raw text", endpoint) return response.text diff --git a/src/cai/tools/reconnaissance/shodan.py b/src/cai/tools/reconnaissance/shodan.py index c67dbcfe..74a9e1af 100644 --- a/src/cai/tools/reconnaissance/shodan.py +++ b/src/cai/tools/reconnaissance/shodan.py @@ -5,12 +5,15 @@ This module provides functions to search Shodan for information about hosts, services, and vulnerabilities using the Shodan API. """ +import logging import os import requests from typing import Dict, List, Optional, Any from dotenv import load_dotenv from cai.sdk.agents import function_tool +logger = logging.getLogger(__name__) + @function_tool def shodan_search(query: str, limit: int = 10) -> str: @@ -112,10 +115,8 @@ def _perform_shodan_search(query: str, limit: int = 10) -> List[Dict[str, Any]]: } try: - response = requests.get(base_url, params=params) - - if response.status_code != 200: - return [] + response = requests.get(base_url, params=params, timeout=30) + response.raise_for_status() data = response.json() @@ -124,7 +125,20 @@ def _perform_shodan_search(query: str, limit: int = 10) -> List[Dict[str, Any]]: return data["matches"][:limit] - except Exception: + except requests.exceptions.Timeout: + logger.warning("Shodan API request timed out for query: %s", query) + return [] + except requests.exceptions.ConnectionError: + logger.warning("Shodan API connection error for query: %s", query) + return [] + except requests.exceptions.HTTPError as e: + logger.warning("Shodan API HTTP error: %s", e) + return [] + except (ValueError, KeyError) as e: + logger.warning("Shodan API response parsing error: %s", e) + return [] + except requests.exceptions.RequestException as e: + logger.warning("Shodan API request failed: %s", e) return [] @@ -149,14 +163,25 @@ def _get_shodan_host_info(ip: str) -> Optional[Dict[str, Any]]: params = {"key": api_key} try: - response = requests.get(base_url, params=params) - - if response.status_code != 200: - return None + response = requests.get(base_url, params=params, timeout=30) + response.raise_for_status() return response.json() - except Exception: + except requests.exceptions.Timeout: + logger.warning("Shodan API request timed out for IP: %s", ip) + return None + except requests.exceptions.ConnectionError: + logger.warning("Shodan API connection error for IP: %s", ip) + return None + except requests.exceptions.HTTPError as e: + logger.warning("Shodan API HTTP error for IP %s: %s", ip, e) + return None + except ValueError as e: + logger.warning("Shodan API response parsing error for IP %s: %s", ip, e) + return None + except requests.exceptions.RequestException as e: + logger.warning("Shodan API request failed for IP %s: %s", ip, e) return None