Initial files

agents/investigator.py

@@ -0,0 +1,113 @@
+from langchain_community.llms import Ollama
+from langchain_community.chat_models import ChatOllama
+
+from langchain import hub
+
+from agentops.langchain_callback_handler import LangchainCallbackHandler as AgentOpsLangchainCallbackHandler
+
+from langchain.chains.conversation.memory import ConversationBufferWindowMemory
+
+from tools.cve_avd_tool import CVESearchTool
+from tools.misp_tool import MispTool
+from tools.coder_tool import CoderTool
+from tools.mitre_tool import MitreTool
+
+from langchain.agents import initialize_agent, AgentType, load_tools
+from langchain.evaluation import load_evaluator
+
+
+from dotenv import load_dotenv
+import os
+import re
+
+load_dotenv(override=True)
+
+
+llm = Ollama(model="openhermes", base_url=os.getenv('OLLAMA_HOST'), temperature=0.3, num_predict=8192, num_ctx=8192)
+wrn = Ollama(model="openfc", base_url=os.getenv('OLLAMA_HOST'))
+llama3 = Ollama(model="llama3", base_url=os.getenv('OLLAMA_HOST'), temperature=0.3)
+
+command_r = Ollama(model="command-r", base_url=os.getenv('OLLAMA_HOST'), temperature=0.1, num_ctx=8192)
+hermes_llama3 = Ollama(model="adrienbrault/nous-hermes2pro-llama3-8b:q4_K_M", base_url=os.getenv('OLLAMA_HOST'), temperature=0.3, num_ctx=32768)
+yarn_mistral_128k = Ollama(model="yarn-mistral-modified", base_url=os.getenv('OLLAMA_HOST'), temperature=0.1, num_ctx=65536, system="""""")
+
+chat_llm = ChatOllama(model="openhermes", base_url=os.getenv('OLLAMA_HOST'), num_predict=-1)
+
+
+cve_search_tool = CVESearchTool().cvesearch
+fetch_cve_tool = CVESearchTool().fetchcve
+misp_search_tool =  MispTool().search
+misp_search_by_date_tool = MispTool().search_by_date
+misp_search_by_event_id_tool = MispTool().search_by_event_id
+coder_tool = CoderTool().code_generation_tool
+
+get_technique_by_id = MitreTool().get_technique_by_id
+get_technique_by_name = MitreTool().get_technique_by_name
+get_malware_by_name = MitreTool().get_malware_by_name
+get_tactic_by_keyword = MitreTool().get_tactic_by_keyword
+
+tools = [cve_search_tool, fetch_cve_tool, misp_search_tool, misp_search_by_date_tool, misp_search_by_event_id_tool, 
+         coder_tool, get_technique_by_id, get_technique_by_name, get_malware_by_name, get_tactic_by_keyword]
+
+# conversational agent memory
+memory = ConversationBufferWindowMemory(
+    memory_key='chat_history',
+    k=4,
+    return_messages=True
+)
+
+agentops_handler = AgentOpsLangchainCallbackHandler(api_key=os.getenv("AGENTOPS_API_KEY"), tags=['Langchain Example'])
+
+#Error handling
+def _handle_error(error) -> str:
+
+    pattern = r'```(?!json)(.*?)```'
+    match = re.search(pattern, str(error), re.DOTALL)
+    if match: 
+        return "The answer contained a code blob which caused the parsing to fail, i recovered the code blob. Just use it to answer the user question: " + match.group(1)
+    else: 
+        return llm.invoke(f"""Try to summarize and explain the following error into 1 short and consice sentence and give a small indication to correct the error: {error} """)
+
+
+prompt = hub.pull("hwchase17/react-chat-json")
+# create our agent
+conversational_agent = initialize_agent(
+    # agent="chat-conversational-react-description",
+    agent=AgentType.STRUCTURED_CHAT_ZERO_SHOT_REACT_DESCRIPTION,
+    tools=tools,
+    prompt=prompt,
+    llm=llm,
+    verbose=True,
+    max_iterations=5,
+    memory=memory,
+    early_stopping_method='generate',
+    # callbacks=[agentops_handler],
+    handle_parsing_errors=_handle_error,
+    return_intermediate_steps=False,
+    max_execution_time=40,
+)
+
+evaluator = load_evaluator("trajectory", llm=chat_llm)
+
+
+# conversational_agent.agent.llm_chain.prompt.messages[0].prompt.template = """
+# 'Respond to the human as helpfully and accurately as possible. 
+# You should use the tools available to you to help answer the question.
+# Your final answer should be technical, well explained, and accurate.
+# You have access to the following tools:\n\n\n\nUse a json blob to specify a tool by providing an action key (tool name) and an action_input key (tool input).\n\nValid "action" values: "Final Answer" or \n\nProvide only ONE action per $JSON_BLOB, as shown:\n\n```\n{{\n  "action": $TOOL_NAME,\n  "action_input": $INPUT\n}}\n```\n\nFollow this format:\n\nQuestion: input question to answer\nThought: consider previous and subsequent steps\nAction:\n```\n$JSON_BLOB\n```\nObservation: action result\n... (repeat Thought/Action/Observation N times)\nThought: I know what to respond\nAction:\n```\n{{\n  "action": "Final Answer",\n  "action_input": "Final response to human"\n}}\n```\n\nBegin! Reminder to ALWAYS respond with a valid json blob of a single action. Use tools if necessary. Respond directly if appropriate. Format is Action:```$JSON_BLOB```then Observation:.\nThought:'
+# """
+
+template = conversational_agent.agent.llm_chain.prompt.messages[0].prompt.template
+
+conversational_agent.agent.llm_chain.prompt.messages[0].prompt.template = """You are a cyber security analyst, you role is to respond to the human queries in a technical way while providing detailed explanations when providing final answer.""" + template
+
+def invoke(input_text):
+    results = conversational_agent({"input":input_text})
+    # evaluation_result = evaluator.evaluate_agent_trajectory(
+    # prediction=results["output"],
+    # input=results["input"],
+    # agent_trajectory=results["intermediate_steps"],
+    # )
+
+    # print(evaluation_result)
+    return results['output']

agents/router.py

@@ -0,0 +1,111 @@
+from langchain_community.llms import Ollama
+
+from langchain import hub
+
+from agentops.langchain_callback_handler import LangchainCallbackHandler as AgentOpsLangchainCallbackHandler
+
+from langchain.chains.conversation.memory import ConversationBufferWindowMemory
+
+from langchain.agents import initialize_agent, AgentType, load_tools
+
+from langchain.tools import StructuredTool, Tool, ShellTool
+
+from dotenv import load_dotenv
+import os
+import re, json
+
+from langchain.agents import create_json_agent
+from langchain.agents.agent_toolkits import JsonToolkit
+from langchain.tools.json.tool import JsonSpec
+
+from .investigator import *
+from .investigator import invoke as investigator_invoke
+
+load_dotenv(override=True)
+
+
+llm = Ollama(model="openhermes", base_url=os.getenv('OLLAMA_HOST'), temperature=0.3, num_predict=8192, num_ctx=8192)
+wrn = Ollama(model="wrn", base_url=os.getenv('OLLAMA_HOST'))
+wrn = Ollama(model="openfc", base_url=os.getenv('OLLAMA_HOST'))
+
+# def get_json_agent(json_path: str):
+#     with open(json_path) as f:
+#         data = json.load(f)
+#     json_spec = JsonSpec(dict_=data, max_value_length=4000)
+#     json_toolkit = JsonToolkit(spec=json_spec)
+
+#     json_agent = create_json_agent(
+#         llm=llm,
+#         toolkit=json_toolkit,
+#         verbose=True
+#     )
+#     return json_agent
+
+# def investigate_agent():
+#     """
+#     This function will help you execute a query to find information about a security event. Just provide the request and get the response.
+#     Parameters:
+#     - request: The request to search for
+#     Returns:
+#     - The response of the search
+#     """
+
+#     def investigate(request: str):
+#         json_agent = get_json_agent("./inventory_prices_dict.json")
+#         result = json_agent.run(
+#             f"""get the price of {inventory_item} from the json file.
+#             Find the closest match to the item you're looking for in that json, e.g.
+#              if you're looking for "mahogany oak table" and that is not in the json, use "table".
+#             Be mindful of the format of the json - there is no list that you can access via [0], so don't try to do that
+#             """)
+#         return result
+
+investigate_tool = Tool(name="Investigate Tool", 
+                        description="This tool will help you execute a query to find information about a security event.(Can be a MISP event, CVE, MITRE attack or technique, malware...) Just provide the request and get the response.", 
+                        func=investigator_invoke)
+
+shell_tool = ShellTool()
+tools = [investigate_tool, shell_tool]
+
+
+memory = ConversationBufferWindowMemory(
+    memory_key='chat_history',
+    k=4,
+    return_messages=True
+)
+
+agent = initialize_agent(
+    agent=AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION,
+    tools=tools,
+    # prompt=prompt,
+    llm=llm,
+    verbose=True,
+    max_iterations=5,
+    memory=memory,
+    early_stopping_method='generate',
+    # return_intermediate_steps=True,
+    handle_parsing_errors=True,
+    max_execution_time=40,
+)
+
+template = agent.agent.llm_chain.prompt.messages[0].prompt.template
+
+# agent.agent.llm_chain.prompt.messages[0].prompt.template = """You are a cyber security analyst called Sonic Cyber Assistant, you were built by a team of engineers at UM6P and DGSSI. you role is to respond to the human queries in a technical way while providing detailed explanations when providing final answer.
+#  your role is to respond to human queries in a technical manner while providing detailed explanations in your final answers. You have a set of tools at your disposal to assist in answering questions. Always delegate investigative tasks to the Investigate Tool, which will perform the investigation and provide results for you to use in your responses. If the Investigate Tool's response contains important information, include it in your answer. If it does not, use the response to formulate your answer. For executing commands, use the Shell Tool and provide the output to the user. Preserve any code blocks and links in your responses as they may contain important information. If a question is unclear, ask the user for clarification. When faced with multiple questions, answer each one separately and sequentially. Never answer questions that are not related to cybersecurity.
+# """
+agent.agent.llm_chain.prompt.messages[0].prompt.template = """You are a cyber security analyst called Sonic Cyber Assistant, you were built by a team of engineers at UM6P and DGSSI. you role is to respond to the human queries in a technical way while providing detailed explanations when providing final answer.
+You are provided with a set of tools to help you answer the questions. Use the tools to help you answer the questions.
+Always delegate the investigation to the Investigate Tool. The Investigate Tool will perform the investigation and provide the results, which you will use to answer the user's question. If the Investigate Tool's response contains some important information, answer the user's question while providing the information. If the Investigate Tool's response does not contain important information, use the Investigate Tool's response to answer the user's question.
+If the user asked you to execute a command, use the Shell Tool to execute the command and provide the output to the user.
+Also try to preserve any code blocks in the response as well as links, as they may contain important information.
+If the question is not clear, ask the user to clarify the question.
+One important thing to remember is that if the question is composed of multiple questions, answer each question separately in a sequential manner.
+NEVER ANSWER QUESTIONS THAT ARE NOT RELATED TO CYBERSECURITY.
+"""
+# print(agent.agent.llm_chain.prompt.messages[0].prompt.template)
+
+def invoke(input_text):
+    return agent({"input":input_text})
+
+def generate_title(input_text):
+    return llm.invoke(f"Generate a title for the following question: {input_text}, the title should be short and concise.")

app.py

@@ -0,0 +1,53 @@
+import streamlit as st
+from agents import investigator, router
+# from agents_openai_fc import investigator
+
+st.title('Cyber Hunter!')
+st.caption("🚀 A streamlit chatbot powered by OpenAI LLM")
+
+if "messages" not in st.session_state:
+    st.session_state["messages"] = [{"role": "assistant", "content": "How can I help you?"}]
+
+for msg in st.session_state.messages:
+    st.chat_message(msg["role"]).write(msg["content"])
+
+if prompt := st.chat_input():
+
+    st.session_state.messages.append({"role": "user", "content": prompt})
+    st.chat_message("user").write(prompt)
+    response = router.invoke(prompt)
+    msg = response["output"]
+    print("Title : ", router.generate_title(prompt))
+
+    # print(response['intermediate_steps'])
+    
+    # If there's an existing assistant message, update it with the new response
+    if st.session_state.messages[-1]["role"] == "assistant":
+        st.session_state.messages[-1]["content"] = msg
+    else:
+        st.session_state.messages.append({"role": "assistant", "content": msg})
+
+    st.chat_message("assistant").write(msg)
+
+# Add a button to regenerate response
+if st.button("Regenerate Response"):
+    st.session_state.messages.pop()  # Remove the latest assistant message
+    prompt = st.session_state.messages[-1]["content"]  # Retrieve user's last input
+    with st.spinner("Searching..."):
+        response = router.invoke(prompt)
+        msg = response["output"]
+        
+        # If there's an existing assistant message, update it with the new response
+        if st.session_state.messages[-1]["role"] == "assistant":
+            st.session_state.messages[-1]["content"] = msg
+        else:
+            st.session_state.messages.append({"role": "assistant", "content": msg})
+
+        st.chat_message("assistant").write(msg)
+
+if st.button("Clear Chat"):
+    st.session_state.messages = [{"role": "assistant", "content": "How can I help you?"}]
+    router.memory.clear()
+    for msg in st.session_state.messages:
+        st.chat_message(msg["role"]).write(msg["content"])
+    st.success("Chat cleared!")

template.env

@@ -0,0 +1,4 @@
+OLLAMA_HOST=""
+MISP_URL=""
+MISP_KEY=""
+ENV="dev"

tools/coder_tool.py

@@ -0,0 +1,30 @@
+from langchain.tools import tool
+from langchain_community.llms import Ollama
+
+import os
+from dotenv import load_dotenv
+load_dotenv(override=True)
+
+wrn = Ollama(model="wrn", base_url=os.getenv('OLLAMA_HOST'), num_predict=512, temperature=0.2,
+             system="""
+    You are a coder and you are trying to generate a code snippet based on a given prompt.
+    The code snippet should be in the programming language that's asked for.
+    Don't Wrap the function in a markdown code block. Return it as a text.
+""")
+
+
+class CoderTool():
+  @tool("Code Generation Tool")
+  def code_generation_tool(instruction: str, language: str = "python"):
+    """The code generation tool is a tool that can generate code snippets based on a given instruction. 
+    It uses a language model to generate code snippets that are relevant to the given instruction.
+    Parameters:
+    - instruction: The instruction for which the code snippet should be generated.
+    - language: The programming language in which the code snippet should be generated. Default is python.
+    Returns:
+    - A code snippet generated based on the given instruction.
+    """
+
+    response = wrn.invoke(instruction)
+    response = response.replace("```", "")
+    return f"'{response}'"

tools/cve_avd_tool.py

@@ -0,0 +1,143 @@
+import requests
+from langchain.tools import tool
+import json
+
+from datetime import datetime, timedelta
+
+def get_current_formatted_date():
+    # Get the current date and time
+    now = datetime.now()
+
+    # Format the date and time in the desired format
+    formatted_date = now.strftime('%Y-%m-%dT%H:%M:%S.%f')[:-3]
+    return formatted_date
+
+def get_yesterday_formatted_date():
+    # Get the current date and time
+    now = datetime.now()
+    # Calculate yesterday's date and time
+    yesterday = now - timedelta(days=2)
+    # Format the date and time in the desired format
+    formatted_date = yesterday.strftime('%Y-%m-%dT%H:%M:%S.%f')[:-3]
+    return formatted_date
+
+def format_cve(cve_response):
+    """
+    Takes a dictionary representing the API response and formats each CVE entry into a structured prompt for LLM.
+    
+    Parameters:
+    - cve_response: A dictionary representing the API response containing CVE information.
+    
+    Returns:
+    - A list of strings, each a formatted prompt for LLM based on the CVE entries in the response.
+    """
+    formatted_prompts = "Use these CVE Search Results to answer user question:\n\n"
+    
+    if len(cve_response['vulnerabilities']) == 0:
+
+        formatted_prompts = "No CVEs found matching the search criteria."
+        return formatted_prompts
+    
+    
+    for vulnerability in cve_response['vulnerabilities']:
+        cve = vulnerability.get('cve', {})
+        # prompt = f"Explain {cve.get('id', 'N/A')} in simple terms:\n\n"
+        prompt = f"- CVE ID: {cve.get('id', 'N/A')}\n"
+        prompt += f"- Status: {cve.get('vulnStatus', 'Unknown')}\n"
+        
+        descriptions = cve.get('descriptions', [])
+        description_text = descriptions[0].get('value', 'No description available.') if descriptions else "No description available."
+        prompt += f"- Description: {description_text}\n"
+        
+        if 'metrics' in cve and 'cvssMetricV2' in cve['metrics']:
+            cvss_metrics = cve['metrics']['cvssMetricV2'][0]
+            prompt += f"- CVSS Score: {cvss_metrics.get('cvssData', {}).get('baseScore', 'Not available')} ({cvss_metrics.get('baseSeverity', 'Unknown')})\n"
+        else:
+            prompt += "- CVSS Score: Not available\n"
+        
+        configurations = cve.get('configurations', {})
+        for conf in configurations:
+            nodes = conf.get('nodes', [])
+            affected_configs = []
+            for node in nodes:
+                for cpe_match in node.get('cpeMatch', []):
+                    if cpe_match.get('vulnerable', False):
+                        affected_configs.append(cpe_match.get('criteria', 'Not specified'))
+            prompt += f"- Affected Configurations: {', '.join(affected_configs) if affected_configs else 'Not specified'}\n"
+        
+        references = cve.get('references', [])
+        ref_urls = ', '.join([ref.get('url', 'No URL') for ref in references])
+        prompt += f"- References: {ref_urls if references else 'No references available.'}\n"
+        
+        
+        formatted_prompts += prompt+"\n\n"
+        print(formatted_prompts)
+
+    # formatted_prompts += "\nSummarize the vulnerability, its impact, and any known mitigation strategies."
+    
+    return formatted_prompts
+
+class CVESearchTool():
+  @tool("CVE search Tool")
+  def cvesearch(keyword: str, date: str = None):
+    """
+    Searches for CVEs based on a keyword or phrase and returns the results in JSON format.
+    
+    Parameters:
+    - keyword (str): A word or phrase to search for in the CVE descriptions.
+    - date (str): An optional date to include in the search query.
+    
+    Returns:
+    - JSON: A list of CVEs matching the keyword search.
+    """
+
+    if date:
+        keyword = f"{keyword} {date}"
+    # Encode the spaces in the keyword(s) as "%20" for the URL
+    keyword_encoded = keyword.replace(" ", "%20")
+    # keyword_encoded = keyword_encoded.join(" 2023")
+    
+    # Construct the URL for the API request
+    url = f"https://services.nvd.nist.gov/rest/json/cves/2.0?keywordSearch={keyword_encoded}&resultsPerPage=5"
+    
+    try:
+        # Send the request to the NVD API
+        response = requests.get(url)
+        # Check if the request was successful
+        if response.status_code == 200:
+            # Return the JSON response
+            return format_cve(response.json())
+        else:
+            return {"error": "Failed to fetch data from the NVD API.", "status_code": response.status_code}
+    except Exception as e:
+        return {"error": str(e)}
+
+  @tool("Fetch last 5 CVEs")
+  def fetchcve(keyword: str):
+        """
+        Fetches the last 5 CVEs based on a keyword or phrase and returns the results in JSON format.
+        Use this exclusively to get the latest CVEs, or for todays CVEs.
+        Parameters:
+        - keyword (str): A word or phrase to search for in the CVE descriptions.
+        - date (str): An optional date to include in the search query.
+        Returns:
+        - JSON: A list of CVEs matching the keyword search.
+        """
+        # Encode the spaces in the keyword(s) as "%20" for the URL
+        keyword_encoded = keyword.replace(" ", "%20")
+        # keyword_encoded = keyword_encoded.join(" 2023")
+        
+        # Construct the URL for the API request
+        url = f"https://services.nvd.nist.gov/rest/json/cves/2.0/?pubStartDate={get_yesterday_formatted_date()}&pubEndDate={get_current_formatted_date()}&keywordSearch={keyword_encoded}&resultsPerPage=3"
+        print(url)
+        try:
+            # Send the request to the NVD API
+            response = requests.get(url)
+            # Check if the request was successful
+            if response.status_code == 200:
+                # Return the JSON response
+                return format_cve(response.json())
+            else:
+                return {"error": "Failed to fetch data from the NVD API.", "status_code": response.status_code}
+        except Exception as e:
+            return {"error": str(e)}

tools/cve_search_tool.py

@@ -0,0 +1,38 @@
+import requests
+from bs4 import BeautifulSoup
+from langchain.tools import tool
+import json
+
+class CVESearchTool():
+  @tool("CVE search Tool")
+  def cvesearch(keyword: str):
+    "CVE (Common Vulnerabilities and Exposures) search tool is a useful tool to search for known security vulnerabilities and exposures in various software products, systems, and devices. It helps users to identify specific vulnerabilities by searching through the CVE database, which contains detailed information about vulnerabilities."
+    url = f"https://cve.mitre.org/cgi-bin/cvekey.cgi?keyword={keyword}"
+    
+    # Fetch the HTML content
+    response = requests.get(url)
+    if response.status_code == 200:
+        html_content = response.content
+
+        # Parse HTML
+        soup = BeautifulSoup(html_content, 'html.parser')
+
+        # Find CVE records
+        cves = soup.find_all('td', {'valign': 'top', 'nowrap': 'nowrap'})
+
+        # Create a dictionary to store CVEs and descriptions
+        cve_dict = {}
+
+        # Iterate through CVE records
+        for cve in cves:
+            cve_id = cve.text.strip()  # Extract CVE ID
+            description = cve.find_next('td').text.strip()  # Extract Description
+            cve_dict[cve_id] = description
+
+        # Convert dictionary to JSON string
+        json_string = json.dumps(cve_dict, indent=4)
+        # return json_string
+        return json_string
+    else:
+        print("Failed to fetch the page:", response.status_code)
+        return None

tools/elastic_tool.py

@@ -0,0 +1,57 @@
+import requests
+from bs4 import BeautifulSoup
+from elasticsearch import Elasticsearch
+from langchain.tools import tool
+
+es = Elasticsearch(
+      "https://localhost:9200",
+      basic_auth=("elastic","dVJI85*y60R3ZVbECj1w"),
+      ca_certs="/Volumes/macOS/Projects/PFE UM6P/elasticsearch-8.12.1/config/certs/http_ca.crt"
+    )
+
+
+class EventSearchTool():
+    @tool("Event search Tool")
+    def search(keyword: str):
+      """Useful tool to search for an indicator of compromise or an security event
+      Parameters:
+      - keyword: The keyword to search for
+      Returns:
+      - A list of events that match the keyword
+      """
+
+      
+      # if not es.ping():
+      #   raise "ElasticNotReachable"
+      
+      query = {
+          "match": {"value": {
+              "query": keyword
+          }}
+      }
+
+      # Execute the search query
+      res = es.search(size=5, index="all_events_full", query=query, knn=None, _source=["event_id", "event_title", "event_date", "category", "attribute_tags", "type", "value"])
+      hits = res["hits"]["hits"]
+      events = [x['_source'] for x in hits]
+
+      return events
+    
+
+    @tool("Event search by event_id Tool")
+    def get_event_by_id(id:str):
+      """Useful tool to search for an event by its id, and return the full event details
+      Parameters:
+      - id: The event id to search for
+      Returns:
+      - The full details of the event with the specified id
+      """
+
+      if not es.ping():
+        raise "ElasticNotReachable"
+      res = es.search(index="all_events_full", query={"match": {"event_id": id}}, _source=["event_id", "event_title", "event_date", "category", "attribute_tags", "type", "value"])
+      hits = res["hits"]["hits"]
+      events = [x['_source'] for x in hits]
+
+      return events
+      

tools/log_tool.py

@@ -0,0 +1 @@
+## TO_IMPLEMENT_LATER

tools/misp_tool.py

@@ -0,0 +1,63 @@
+from langchain.tools import tool
+# import pymisp
+
+
+from pymisp import PyMISP
+from dotenv import load_dotenv
+import os
+
+load_dotenv(override=True)
+
+URL = os.getenv('MISP_URL')
+KEY = os.getenv('MISP_KEY')
+verify_cert = False
+
+print(URL, KEY)
+
+misp = PyMISP(url=URL, key=KEY, ssl=verify_cert)
+
+class MispTool():
+    @tool("MISP search Tool by keyword")
+    def search(keyword: str):
+      """Useful tool to search for an indicator of compromise or an security event by keyword
+      Parameters:
+      - keyword: The keyword to search for
+      Returns:
+      - A list of events that match the keyword
+      """
+
+      events = misp.search(controller='attributes', value=keyword, limit=5, metadata=True, include_event_tags=False, include_context=False, return_format='json', sg_reference_only=True)
+      
+      if len(events['Attribute']) == 0:
+        return "No events found matching the search criteria."
+      
+      results = """Answer user question using these search results:\n\n"""
+      return results + str(events)
+    
+    @tool("MISP search Tool by date")
+    def search_by_date(date_from: str = None, date_to: str = None):
+      """Useful tool to retrieve events that match a specific date or date range, use this if you know the date of the event
+      Parameters:
+      - date_from: The start date of the event
+      - date_to: The end date of the event
+      Not necessary to provide both dates, you can provide one or the other
+
+      Returns:
+      - A list of events that match the date or date range
+      """
+
+      events = misp.search(controller='attributes',date_from=date_from, date_to=date_to, limit=5)
+      return events
+
+    @tool("MISP search Tool by event_id")
+    def search_by_event_id(event_id: str | int):
+      """Useful tool to retrieve events by their ID, use this if you know the ID of the event.
+      Parameters:
+      - event_id: The ID of the event
+      Returns:
+      - A list of events that match the event ID
+      """
+
+      events = misp.search(controller='attributes', eventid=event_id, limit=1)
+      return events
+    

tools/mitre_tool.py

@@ -0,0 +1,49 @@
+import requests
+from stix2 import MemoryStore, Filter
+from taxii2client.v20 import Server # only specify v20 if your installed version is >= 2.0.0
+
+from langchain.tools import tool
+
+
+def get_data_from_branch(domain, branch="master"):
+    """get the ATT&CK STIX data from MITRE/CTI. Domain should be 'enterprise-attack', 'mobile-attack' or 'ics-attack'. Branch should typically be master."""
+    BASE_URL = f"https://raw.githubusercontent.com/mitre/cti/{branch}/{domain}/{domain}.json"
+    stix_json = requests.get(BASE_URL).json()
+    return MemoryStore(stix_data=stix_json["objects"])
+
+store = {
+        "enterprise": get_data_from_branch("enterprise-attack"),
+        "mobile": get_data_from_branch("mobile-attack"),
+        "ics": get_data_from_branch("ics-attack")
+}
+
+class MitreTool():
+
+    @tool("MITRE Technique search by ID")
+    def get_technique_by_id(domain: str, technique_id: str):
+        """Get the technique by its ID. Domain should be 'enterprise', 'mobile' or 'ics'
+        Techniques represent 'how' an adversary achieves a tactical goal by performing an action. For example, an adversary may dump credentials to achieve credential access.
+        """
+        result = store[domain].query([Filter('external_references.external_id', '=', technique_id)])
+        return result if result else "No technique found with that ID"
+    
+    @tool("MITRE Technique search by name")
+    def get_technique_by_name(domain: str, technique_name: str):
+        """Get the technique by its name. Domain should be 'enterprise', 'mobile' or 'ics'
+        Techniques represent 'how' an adversary achieves a tactical goal by performing an action. For example, an adversary may dump credentials to achieve credential access."""
+        result = store[domain].query([Filter('name', 'contains', technique_name), Filter('type', '=', 'attack-pattern')])
+        return result if result else "No technique found with that name"
+    
+    @tool("MITRE Malware search by name")
+    def get_malware_by_name(domain: str, malware_name: str):
+        """Get the malware by its name. Domain should be 'enterprise', 'mobile' or 'ics'
+        Malware represents software used to achieve a tactical goal by performing an action. For example, an adversary may use malware to achieve initial access."""
+        result = store[domain].query([Filter('name', 'contains', malware_name), Filter('type', '=', 'malware')])
+        return result if result else "No malware found with that name"
+    
+    @tool("MITRE Technique search by keyword")
+    def get_tactic_by_keyword(domain: str, keyword: str):
+        """Search for tactics/techniques by a keyword. Domain should be 'enterprise', 'mobile' or 'ics'
+        Tactics represent the "why" of an ATT&CK technique or sub-technique. It is the adversary's tactical goal: the reason for performing an action. For example, an adversary may want to achieve credential access."""
+        result = store[domain].query([Filter('description', 'contains', keyword)], Filter('type', '=', 'attack-pattern'))
+        return result[0] if result else "No tactics/techniques matches the keyword you provided"

tools/scraper_tools.py

@@ -0,0 +1,31 @@
+import requests
+from bs4 import BeautifulSoup
+from langchain.tools import tool
+
+class ScraperTool():
+  @tool("Scraper Tool")
+  def scrape(url: str):
+    "Useful tool to scrap a website content, use to learn more about a given url."
+
+    headers = {
+      'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.3'}
+
+    response = requests.get(url, headers=headers)
+    
+    # Check if the request was successful
+    if response.status_code == 200:
+        # Parse the HTML content of the page
+        soup = BeautifulSoup(response.text, 'html.parser')
+
+        article = soup.find(id='insertArticle')
+        
+        if article:
+            # Extract and print the text from the article
+            text = (article.get_text(separator=' ', strip=True))
+        else:
+            print("Article with specified ID not found.")
+        
+        return text
+    else:
+        print("Failed to retrieve the webpage")
+