Feature/ava

Makefile

@@ -0,0 +1,21 @@
+SHELL := /bin/bash
+
+.DEFAULT_GOAL := help
+
+SRC_DIRS := sigmaiq/llm examples
+PYTHON_FILES := $(shell find $(SRC_DIRS) -type f -name "*.py")
+
+.PHONY: help
+help:
+	@echo "Usage:"
+	@awk 'BEGIN {FS = ":.*?## "} /^[a-zA-Z0-9_-]+:.*?## / {printf "  %-20s%s\n", $$1, $$2}' $(MAKEFILE_LIST)
+
+.PHONY: format
+format: ## Format Python files
+	@echo "Formatting Python files..."
+	black $(PYTHON_FILES)
+
+.PHONY: ruff
+ruff: ## Run Ruff linter
+	@echo "Running Ruff linter..."
+	ruff check $(PYTHON_FILES)

examples/custom_field_mappings.py

@@ -1,46 +1,47 @@
-# %% This example shows how to use the SigmAIQ pySigma wrapper to provide custom field mappings for a backend
-# %% This will allow you to translate specific field names to custom field names during rule translation
-
-# %% Import SigmAIQ
-from sigmaiq import SigmAIQBackend, SigmAIQPipeline
-
-# %% Import pprint for pretty printing, and copy for copying rules
-from pprint import pprint
-from copy import copy
-
-# %% A basic Sigma Rule in YAML str to convert to a query.
-# %% SigmAIQ also accepts a rule in JSON/Dict format, SigmaRule objects, and SigmaCollection objects
-
-sigma_rule = """
-title: whoami Command
-description: Detects a basic whoami commandline execution
-logsource:
-    product: windows
-    category: process_creation
-detection:
-    selection1:
-        - CommandLine|contains: 'whoami.exe'
-    condition: selection1
-"""
-
-# %% Create SigmAIQ backend translate the rule to a Microsoft 365 Defender query
-sigmaiq_backend = SigmAIQBackend(backend="splunk").create_backend()
-query = sigmaiq_backend.translate(copy(sigma_rule))  # Returns List of queries
-
-print("\nM365Defender Query: ", end="\n\n")
-pprint(query[0])
-print("\n-------------------")
-
-# %% Create custom field mappings
-# %% This will map the CommandLine field to a custom field named "CustomCommandLine"
-custom_field_mappings = {"CommandLine": "CustomCommandLine"}
-my_custom_pipeline = SigmAIQPipeline.from_fieldmap(custom_field_mappings, priority=0).create_pipeline()
-
-# %% Create SigmAIQ backend translate the rule to a Microsoft 365 Defender query with our custom field mappings
-sigmaiq_backend = SigmAIQBackend(backend="splunk", processing_pipeline=my_custom_pipeline).create_backend()
-
-query = sigmaiq_backend.translate(copy(sigma_rule))  # Returns List of queries
-
-print("\nM365Defender Query with Custom Fieldmappings: ", end="\n\n")
-pprint(query[0])
-print("\n-------------------")
+# %% This example shows how to use the SigmAIQ pySigma wrapper to provide custom field mappings for a backend
+# %% This will allow you to translate specific field names to custom field names during rule translation
+
+from copy import copy
+
+# %% Import pprint for pretty printing, and copy for copying rules
+from pprint import pprint
+
+# %% Import SigmAIQ
+from sigmaiq import SigmAIQBackend, SigmAIQPipeline
+
+# %% A basic Sigma Rule in YAML str to convert to a query.
+# %% SigmAIQ also accepts a rule in JSON/Dict format, SigmaRule objects, and SigmaCollection objects
+
+sigma_rule = """
+title: whoami Command
+description: Detects a basic whoami commandline execution
+logsource:
+    product: windows
+    category: process_creation
+detection:
+    selection1:
+        - CommandLine|contains: 'whoami.exe'
+    condition: selection1
+"""
+
+# %% Create SigmAIQ backend translate the rule to a Microsoft 365 Defender query
+sigmaiq_backend = SigmAIQBackend(backend="splunk").create_backend()
+query = sigmaiq_backend.translate(copy(sigma_rule))  # Returns List of queries
+
+print("\nM365Defender Query: ", end="\n\n")
+pprint(query[0])
+print("\n-------------------")
+
+# %% Create custom field mappings
+# %% This will map the CommandLine field to a custom field named "CustomCommandLine"
+custom_field_mappings = {"CommandLine": "CustomCommandLine"}
+my_custom_pipeline = SigmAIQPipeline.from_fieldmap(custom_field_mappings, priority=0).create_pipeline()
+
+# %% Create SigmAIQ backend translate the rule to a Microsoft 365 Defender query with our custom field mappings
+sigmaiq_backend = SigmAIQBackend(backend="splunk", processing_pipeline=my_custom_pipeline).create_backend()
+
+query = sigmaiq_backend.translate(copy(sigma_rule))  # Returns List of queries
+
+print("\nM365Defender Query with Custom Fieldmappings: ", end="\n\n")
+pprint(query[0])
+print("\n-------------------")

examples/llm_basic_usage.py

@@ -1,36 +1,37 @@
-# %% This example will demonstrate how to use SigmAIQ to perform the following:
-# %% 1. Download the latest Sigma Rule package release
-# %% 2. Create embeddings of the Sigma Rules in the package
-# %% 3. Create and save a VectorDB of the Sigma Rule embeddings
-# %% 4. Use a similarity search on the VectorDB to find Sigma Rules similar to a provided query
-from pprint import pprint
-
-# %% NOTE, this example uses OpenAI for embeddings. Ensure you have an OpenAI API key set in your environment variable
-# %% OPENAI_API_KEY
-
-# %% Also ensure you have installed the correct requirements with:
-# `pip install -r requirements/common.txt -r requirements/llm.txt`
-
-
-# %% Import SigmAIQ LLM and OpenAIEmbeddings
-from sigmaiq.llm.base import SigmaLLM
-
-# %% Create a SigmaLLM object with default settings. See the class docstring for more information
-sigma_llm = SigmaLLM()
-
-# %% The `create_sigma_vectordb()` method will automatically do all the work for you :) (only run this once)
-sigma_llm.create_sigma_vectordb(save=True)  # Save locally to disk
-
-# %% Run a similarity search on the vectordb for encoded powershell commands and print top 3 results
-query = "Encoded powershell commands"
-matching_rules = sigma_llm.simple_search(query, k=3)
-for matching_rule in matching_rules:
-    print(matching_rule.page_content, end="\n\n-------------------\n\n")
-
-# %% You can also load an existing vector store from disk (recommended)
-sigma_llm.load_sigma_vectordb()
-
-query = "certutil"
-matching_rules = sigma_llm.simple_search(query, k=3)
-for matching_rule in matching_rules:
-    print(matching_rule.page_content, end="\n\n-------------------\n\n")
+# %% This example will demonstrate how to use SigmAIQ to perform the following:
+# %% 1. Download the latest Sigma Rule package release
+# %% 2. Create embeddings of the Sigma Rules in the package
+# %% 3. Create and save a VectorDB of the Sigma Rule embeddings
+# %% 4. Use a similarity search on the VectorDB to find Sigma Rules similar to a provided query
+
+# %% NOTE, this example uses OpenAI for embeddings. Ensure you have an OpenAI API key set in your environment variable
+# %% OPENAI_API_KEY
+
+# %% Also ensure you have installed the correct requirements with:
+# `pip install -r requirements/common.txt -r requirements/llm.txt`
+
+
+# %% Create a SigmaLLM object with default settings. See the class docstring for more information
+from langchain_openai import OpenAIEmbeddings
+
+# %% Import SigmAIQ LLM and OpenAIEmbeddings
+from sigmaiq.llm.base import SigmaLLM
+
+sigma_llm = SigmaLLM(embedding_model=OpenAIEmbeddings(model="text-embedding-3-large"))
+
+# %% The `create_sigma_vectordb()` method will automatically do all the work for you :) (only run this once)
+sigma_llm.create_sigma_vectordb(save=True)  # Save locally to disk
+
+# %% Run a similarity search on the vectordb for encoded powershell commands and print top 3 results
+query = "Encoded powershell commands"
+matching_rules = sigma_llm.simple_search(query, k=3)
+for matching_rule in matching_rules:
+    print(matching_rule.page_content, end="\n\n-------------------\n\n")
+
+# %% You can also load an existing vector store from disk (recommended)
+sigma_llm.load_sigma_vectordb()
+
+query = "certutil"
+matching_rules = sigma_llm.simple_search(query, k=3)
+for matching_rule in matching_rules:
+    print(matching_rule.page_content, end="\n\n-------------------\n\n")

examples/llm_rule_translation_and_creation.py

@@ -1,59 +1,110 @@
-# %% This example will demonstrate how to create a Sigma langchain agent chatbot, which can perform various tasks like
-# %% automatically translate a rule for you, and create new rules from a users input.
-
-# %% Import required SigmAIQ classes and functions
-from sigmaiq.llm.toolkits.base import create_sigma_agent
-from sigmaiq.llm.base import SigmaLLM
-
-# %% Ensure we have our Sigma vector store setup with our base LLM class
-sigma_llm = SigmaLLM()
-
-try:
-    sigma_llm.load_sigma_vectordb()
-except Exception as e:
-    print(e)
-    print("Creating new Sigma VectorDB")
-    sigma_llm.create_sigma_vectordb(save=True)
-
-# %% Create a Sigma Agent Executor, and pass it our Sigma VectorDB
-sigma_agent_executor = create_sigma_agent(sigma_vectorstore=sigma_llm.sigmadb)
-
-# %% RULE TRANSLATION
-# %% Have the agent automatically translate a Sigma rule to a Splunk query with the splunk_cim_dm pipeline
-
-sigma_rule = r"""
-title: whoami Command
-description: Detects a basic whoami commandline execution
-logsource:
-    product: windows
-    category: process_creation
-detection:
-    selection1:
-        - CommandLine|contains: 'whoami.exe'
-    condition: selection1
-"""
-
-user_input = (
-    "Translate the following Sigma rule to a Splunk query using the 'splunk_cim_dm' pipeline: \n\n" + sigma_rule
-)
-
-answer = sigma_agent_executor.invoke({"input": user_input})
-print("\nRULE TRANSLATION:", end="\n\n")
-print(f"Question:\n {user_input}", end="\n\n")
-print(f"Answer: \n")
-print(answer.get("output"), end="\n\n")
-
-# %% RULE CREATION
-# %% The agent will take the user input, look up similar Sigma Rules in the Sigma vector store, then create a brand
-# %% new rule based on the context of the users input and the similar Sigma Rules.
-
-user_input = (
-    "Create a Windows process creation Sigma Rule for certutil downloading a file "
-    "from definitely-not-malware.com, then translate it to a Microsoft 365 Defender query."
-)
-
-answer = sigma_agent_executor.invoke({"input": user_input})
-print("\nRULE CREATION:", end="\n\n")
-print(f"Question:\n {user_input}", end="\n\n")
-print(f"Answer: \n")
-print(answer.get("output"), end="\n\n")
+# %% This example will demonstrate how to create a Sigma langchain agent chatbot, which can perform various tasks like
+# %% automatically translate a rule for you, and create new rules from a users input.
+import asyncio
+
+from langchain_openai import OpenAIEmbeddings
+
+from sigmaiq.llm.base import SigmaLLM
+
+# %% Import required SigmAIQ classes and functions
+from sigmaiq.llm.toolkits.base import create_sigma_agent
+
+# %% Ensure we have our Sigma vector store setup with our base LLM class
+sigma_llm = SigmaLLM(embedding_model=OpenAIEmbeddings(model="text-embedding-3-large"))
+
+try:
+    sigma_llm.load_sigma_vectordb()
+except Exception as e:
+    print(e)
+    print("Creating new Sigma VectorDB")
+    sigma_llm.create_sigma_vectordb(save=True)
+
+# %% Create a Sigma Agent Executor, and pass it our Sigma VectorDB
+sigma_agent_executor = create_sigma_agent(sigma_vectorstore=sigma_llm.sigmadb)
+
+
+async def main():
+    print("\n--------\nRULE TRANSLATION\n--------\n")
+    # %% RULE TRANSLATION
+    # %% Have the agent automatically translate a Sigma rule to a Splunk query with the splunk_cim_dm pipeline
+    user_input = (
+        "Convert this Sigma rule to a Splunk query using the 'splunk_cim_dm' pipeline: \n\n"
+        + "title: whoami Command\n"
+        + "description: Detects a basic whoami commandline execution\n"
+        + "logsource:\n"
+        + "    product: windows\n"
+        + "    category: process_creation\n"
+        + "detection:\n"
+        + "    selection1:\n"
+        + "        - CommandLine|contains: 'whoami.exe'\n"
+        + "    condition: selection1"
+    )
+
+    answer = await sigma_agent_executor.ainvoke({"input": user_input})
+    print(f"\n\nQUESTION:\n {user_input}", end="\n\n")
+    print("ANSWER: \n")
+    print(answer.get("output"), end="\n\n")
+
+    # %% RULE SEARCHING
+    # %% The agent will find official Sigma rules in the Sigma vector store based on the context of the users input.
+    print("\n--------\nRULE SEARCHING\n--------\n")
+
+    user_input = (
+        "What Sigma Rule can detect a process creation event where the parent process is word.exe and the child "
+        "process is cmd.exe?"
+    )
+
+    answer = await sigma_agent_executor.ainvoke({"input": user_input})
+    print(f"QUESTION:\n {user_input}", end="\n\n")
+    print("ANSWER: \n")
+    print(answer.get("output"), end="\n\n")
+
+    # %% QUERY TO SIGMA RULE CONVERSION
+    # %% If you already have a query and want to convert it back to a Sigma Rule
+    print("\n--------\nQUERY TO SIGMA RULE \n--------\n")
+
+    user_input = """
+Convert this Microsoft365Defender query into a Sigma Rule:
+DeviceProcessEvents
+| where ((ProcessCommandLine contains "powershell.exe -enc" or ProcessCommandLine contains "cmd.exe /c" or ProcessCommandLine contains "rundll32.exe") or (InitiatingProcessFolderPath endswith "\\explorer.exe" or InitiatingProcessFolderPath endswith "\\svchost.exe")) and (not((ProcessCommandLine contains "schtasks" or ProcessCommandLine contains "tasklist")))
+
+"""
+
+    answer = await sigma_agent_executor.ainvoke({"input": user_input})
+    print(f"\n\nQUESTION:\n {user_input}", end="\n\n")
+    print("ANSWER: \n")
+    print(answer.get("output"), end="\n\n")
+
+    # %% RULE CREATION
+    # %% The agent will take the user input, look up similar Sigma Rules in the Sigma vector store, then create a brand
+    # %% new rule based on the context of the users input and the similar Sigma Rules.
+    print("\n--------\nRULE CREATION\n--------\n")
+
+    user_input = (
+        "Create a Windows process creation Sigma Rule for certutil downloading a file "
+        "from definitely-not-malware.com, then translate it to a Microsoft 365 Defender query."
+    )
+
+    answer = await sigma_agent_executor.ainvoke({"input": user_input})
+    print(f"QUESTION:\n {user_input}", end="\n\n")
+    print("ANSWER: \n")
+    print(answer.get("output"), end="\n\n")
+
+    # %% RULE CREATION + TRANSLATION
+    # %% The agent will take the user input, look up similar Sigma Rules in the Sigma vector store, then create a brand
+    # %% new rule based on the context of the users input and the similar Sigma Rules. Then convert it to a M365 query
+    print("\n--------\nRULE CREATION + TRANSLATION\n--------\n")
+
+    user_input = (
+        "Create a Windows process creation Sigma Rule for certutil downloading a file "
+        "from definitely-not-malware.com, then translate it to a Microsoft 365 Defender query."
+    )
+
+    answer = await sigma_agent_executor.ainvoke({"input": user_input})
+    print(f"QUESTION:\n {user_input}", end="\n\n")
+    print("ANSWER: \n")
+    print(answer.get("output"), end="\n\n")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())