tsumarios/CTI_Agent_Example

1from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool, tool2import datetime3import re4import requests5import pytz6import yaml7from tools.final_answer import FinalAnswerTool8 9from Gradio_UI import GradioUI10 11# CTI tools12@tool13def threat_lookup(threat_name: str) -> str:14    """Fetches basic information about a known cyber threat.15    Args:16        threat_name: The name of the threat to search.17    """18    try:19        url = f"https://cve.circl.lu/api/cve/{threat_name}"20        response = requests.get(url, timeout=5)  # Added timeout21        if response.status_code == 200:22            data = response.json()23            if 'id' in data and 'summary' in data:24                return f"CVE: {data['id']}\nSummary: {data['summary']}"25            else:26                return f"No detailed information found for '{threat_name}'."27        else:28            return f"Threat '{threat_name}' not found (HTTP {response.status_code})."29    except requests.exceptions.Timeout:30        return "Request timed out. Try again later."31    except Exception as e:32        return f"Error fetching threat details: {str(e)}"33 34 35 36@tool37def check_ip_reputation(ip_address: str) -> str:38    """Checks if an IP address has been reported as malicious.39    Args:40        ip_address: The IP address to check.41    """42    try:43        url = f"https://www.abuseipdb.com/check/{ip_address}"44        return f"Check the reputation of {ip_address} here: {url}"45    except Exception as e:46        return f"Error checking IP reputation: {str(e)}"47 48 49@tool50def extract_iocs(text: str) -> dict:51    """Extracts Indicators of Compromise (IP addresses, domains, and hashes) from text.52    Args:53        text: A block of text containing potential IOCs.54    """55    ip_pattern = r"\b(?:\d{1,3}\.){3}\d{1,3}\b"56    domain_pattern = r"\b[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}(?<!\.)\b"57    hash_pattern = r"\b[a-fA-F0-9]{32,64}\b"58 59    iocs = {60        "IPs": re.findall(ip_pattern, text),61        "Domains": re.findall(domain_pattern, text),62        "Hashes": re.findall(hash_pattern, text)63    }64    65    return {key: value for key, value in iocs.items() if value}66 67 68# Other tools from the example template69@tool70def get_current_time_in_timezone(timezone: str) -> str:71    """A tool that fetches the current local time in a specified timezone.72    Args:73        timezone: A string representing a valid timezone (e.g., 'America/New_York').74    """75    try:76        # Create timezone object77        tz = pytz.timezone(timezone)78        # Get current time in that timezone79        local_time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")80        return f"The current local time in {timezone} is: {local_time}"81    except Exception as e:82        return f"Error fetching time for timezone '{timezone}': {str(e)}"83 84 85final_answer = FinalAnswerTool()86model = HfApiModel(87max_tokens=2096,88temperature=0.5,89model_id='https://wxknx1kg971u7k1n.us-east-1.aws.endpoints.huggingface.cloud',# it is possible that this model may be overloaded (it is indeed D:)90custom_role_conversions=None,91)92 93# Import tool from Hub94image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)95 96# DuckDuckGoSearch tool97search_tool = DuckDuckGoSearchTool()98 99with open("prompts.yaml", 'r') as stream:100    prompt_templates = yaml.safe_load(stream)101    102agent = CodeAgent(103    model=model,104    tools=[threat_lookup, check_ip_reputation, extract_iocs, get_current_time_in_timezone, 105           search_tool, image_generation_tool, final_answer], ## add your tools here (don't remove final answer)106    max_steps=6,107    verbosity_level=1,108    grammar=None,109    planning_interval=None,110    name=None,111    description=None,112    prompt_templates=prompt_templates113)114 115 116GradioUI(agent).launch()