txt2detection 1.1.3__py3-none-any.whl

txt2detection/__init__.py

@@ -0,0 +1,2 @@
+from .__main__ import run_txt2detection
+from . import observables, bundler, models, utils

txt2detection/__main__.py

@@ -0,0 +1,343 @@
+import argparse
+from datetime import UTC, datetime
+
+from dataclasses import dataclass
+import io
+import json
+import os
+from pathlib import Path
+import logging
+import re
+import sys
+import uuid
+from pydantic import ValidationError
+from stix2 import Identity
+import yaml
+
+from txt2detection import credential_checker
+from txt2detection.ai_extractor.base import BaseAIExtractor
+from txt2detection.models import (
+    TAG_PATTERN,
+    DetectionContainer,
+    Level,
+    SigmaRuleDetection,
+)
+from txt2detection.utils import validate_token_count
+
+
+def configureLogging():
+    # Configure logging
+    stream_handler = logging.StreamHandler()  # Log to stdout and stderr
+    stream_handler.setLevel(logging.INFO)
+    logging.basicConfig(
+        level=logging.DEBUG,  # Set the desired logging level
+        format=f"%(asctime)s [%(levelname)s] %(message)s",
+        handlers=[stream_handler],
+        datefmt="%d-%b-%y %H:%M:%S",
+    )
+
+    return logging.root
+
+
+configureLogging()
+
+
+def setLogFile(logger, file: Path):
+    file.parent.mkdir(parents=True, exist_ok=True)
+    logger.info(f"Saving log to `{file.absolute()}`")
+    handler = logging.FileHandler(file, "w")
+    handler.formatter = logging.Formatter(
+        fmt="%(levelname)s %(asctime)s - %(message)s", datefmt="%d-%b-%y %H:%M:%S"
+    )
+    handler.setLevel(logging.DEBUG)
+    logger.addHandler(handler)
+    logger.info("=====================txt2detection======================")
+
+
+from .bundler import Bundler
+import shutil
+
+
+from .utils import STATUSES, as_date, make_identity, valid_licenses, parse_model
+
+
+def parse_identity(str):
+    return Identity(**json.loads(str))
+
+
+@dataclass
+class Args:
+    input_file: str
+    input_text: str
+    name: str
+    tlp_level: str
+    labels: list[str]
+    created: datetime
+    use_identity: Identity
+    ai_provider: BaseAIExtractor
+    report_id: uuid.UUID
+    external_refs: dict[str, str]
+    reference_urls: list[str]
+
+
+def parse_created(value):
+    """Convert the created timestamp to a datetime object."""
+    try:
+        return datetime.strptime(value, "%Y-%m-%dT%H:%M:%S").replace(tzinfo=UTC)
+    except ValueError:
+        raise argparse.ArgumentTypeError(
+            "Invalid date format. Use YYYY-MM-DDTHH:MM:SS."
+        )
+
+
+def parse_ref(value):
+    m = re.compile(r"(.+?)=(.+)").match(value)
+    if not m:
+        raise argparse.ArgumentTypeError("must be in format key=value")
+    return dict(source_name=m.group(1), external_id=m.group(2))
+
+
+def parse_label(label: str):
+    if not TAG_PATTERN.match(label):
+        raise argparse.ArgumentTypeError(
+            "Invalid label format. Must follow sigma tag format {namespace}.{label}"
+        )
+    namespace, _, _ = label.partition(".")
+    if namespace in ["tlp"]:
+        raise argparse.ArgumentTypeError(f"Unsupported tag namespace `{namespace}`")
+    return label
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Convert text file to detection format."
+    )
+    mode = parser.add_subparsers(
+        title="process-mode", dest="mode", description="mode to use"
+    )
+    file = mode.add_parser("file", help="process a file input using ai")
+    text = mode.add_parser("text", help="process a text argument using ai")
+    sigma = mode.add_parser("sigma", help="process a sigma file without ai")
+    check_credentials = mode.add_parser(
+        "check-credentials",
+        help="show status of external services with respect to credentials",
+    )
+
+    for mode_parser in [file, text, sigma]:
+        mode_parser.add_argument(
+            "--report_id", type=uuid.UUID, help="report_id to use for generated report"
+        )
+        mode_parser.add_argument(
+            "--name",
+            required=True,
+            help="Name of file, max 72 chars. Will be used in the STIX Report Object created.",
+        )
+        mode_parser.add_argument(
+            "--tlp_level",
+            choices=["clear", "green", "amber", "amber_strict", "red"],
+            help="Options are clear, green, amber, amber_strict, red. Default is clear if not passed.",
+        )
+        mode_parser.add_argument(
+            "--labels",
+            type=parse_label,
+            action="extend",
+            nargs="+",
+            help="Comma-separated list of labels. Case-insensitive (will be converted to lower-case). Allowed a-z, 0-9.",
+        )
+        mode_parser.add_argument(
+            "--created",
+            type=parse_created,
+            help="Explicitly set created time in format YYYY-MM-DDTHH:MM:SS.sssZ. Default is current time.",
+        )
+        mode_parser.add_argument(
+            "--use_identity",
+            type=parse_identity,
+            help="Pass a full STIX 2.1 identity object (properly escaped). Validated by the STIX2 library. Default is SIEM Rules identity.",
+        )
+        mode_parser.add_argument(
+            "--ai_provider",
+            required=False,
+            type=parse_model,
+            help="(required): defines the `provider:model` to be used. Select one option.",
+            metavar="provider[:model]",
+        )
+        mode_parser.add_argument(
+            "--external_refs",
+            type=parse_ref,
+            help="pass additional `external_references` entry (or entries) to the report object created. e.g --external_ref author=dogesec link=https://dkjjadhdaj.net",
+            default=[],
+            metavar="{source_name}={external_id}",
+            action="extend",
+            nargs="+",
+        )
+        mode_parser.add_argument(
+            "--reference_urls",
+            help="pass additional `external_references` url entry (or entries) to the report object created.",
+            default=[],
+            metavar="{url}",
+            action="extend",
+            nargs="+",
+        )
+        mode_parser.add_argument(
+            "--license",
+            help="Valid SPDX license for the rule",
+            default=None,
+            metavar="[LICENSE]",
+            choices=valid_licenses(),
+        )
+        mode_parser.add_argument(
+            "--create_attack_navigator_layer",
+            help="Create navigator layer",
+            action="store_true",
+            default=False,
+        )
+
+    file.add_argument(
+        "--input_file",
+        help="The file to be converted. Must be .txt",
+        type=lambda x: Path(x).read_text(),
+    )
+    text.add_argument("--input_text", help="The text to be converted")
+    sigma.add_argument(
+        "--sigma_file",
+        help="The sigma file to be converted. Must be .yml",
+        type=lambda x: Path(x).read_text(),
+    )
+    sigma.add_argument(
+        "--status",
+        help="If passed, will overwrite any existing `status` recorded in the rule",
+        choices=STATUSES,
+    )
+    sigma.add_argument(
+        "--level",
+        help="If passed, will overwrite any existing `level` recorded in the rule",
+        choices=Level._member_names_,
+    )
+
+    args: Args = parser.parse_args()
+    if args.mode == "check-credentials":
+        statuses = credential_checker.check_statuses(test_llms=True)
+        credential_checker.format_statuses(statuses)
+        sys.exit(0)
+
+    if args.mode != "sigma":
+        assert args.ai_provider, "--ai_provider is required in file or txt mode"
+
+    if args.mode == "file":
+        args.input_text = args.input_file
+
+    args.input_text = getattr(args, "input_text", "")
+    if not args.report_id:
+        args.report_id = Bundler.generate_report_id(
+            args.use_identity.id if args.use_identity else None, args.created, args.name
+        )
+
+    return args
+
+
+def run_txt2detection(
+    name,
+    identity,
+    tlp_level,
+    input_text: str,
+    labels: list[str],
+    report_id: str | uuid.UUID,
+    ai_provider: BaseAIExtractor,
+    create_attack_navigator_layer=False,
+    **kwargs,
+) -> Bundler:
+    if not kwargs.get("sigma_file"):
+        validate_token_count(
+            int(os.getenv("INPUT_TOKEN_LIMIT", 0)), input_text, ai_provider
+        )
+
+    if sigma := kwargs.get("sigma_file"):
+        detection = get_sigma_detections(sigma, name=name)
+        if not identity and detection.author:
+            identity = make_identity(detection.author)
+        kwargs.update(
+            reference_urls=kwargs.setdefault("reference_urls", [])
+            + detection.references
+        )
+        if not kwargs.get("created"):
+            # only consider rule.date and rule.modified if user does not pass --created
+            kwargs.update(
+                created=detection.date,
+                modified=detection.modified,
+            )
+        kwargs['license'] = kwargs.get('license') or detection.license
+        detection.level = kwargs.get("level") or detection.level
+        detection.status = kwargs.get("status") or detection.status
+        detection.date = as_date(kwargs.get("created"))
+        detection.modified = as_date(kwargs.get("modified"))
+        detection.references = kwargs["reference_urls"]
+        detection.detection_id = str(report_id).removeprefix("report--")
+        bundler = Bundler(
+            name or detection.title,
+            identity,
+            tlp_level or detection.tlp_level or "clear",
+            detection.description,
+            (labels or []) + detection.tags,
+            report_id=report_id,
+            **kwargs,
+        )
+        detections = DetectionContainer(success=True, detections=[])
+        detections.detections.append(detection)
+    else:
+        bundler = Bundler(
+            name, identity, tlp_level, input_text, labels, report_id=report_id, **kwargs
+        )
+        detections = ai_provider.get_detections(input_text)
+    bundler.bundle_detections(detections)
+    if create_attack_navigator_layer:
+        bundler.create_attack_navigator()
+    return bundler
+
+
+def get_sigma_detections(sigma: str, name=None) -> SigmaRuleDetection:
+    obj = yaml.safe_load(io.StringIO(sigma))
+    if not isinstance(obj, dict):
+        raise ValueError(
+            f"bad sigma input file. expected object/dict, got {type(obj)}."
+        )
+    if name:
+        obj["title"] = name
+    return SigmaRuleDetection.model_validate(obj)
+
+
+def main(args: Args):
+
+    setLogFile(logging.root, Path(f"logs/log-{args.report_id}.log"))
+    logging.info(f"starting argument: {json.dumps(sys.argv[1:])}")
+    kwargs = args.__dict__
+    kwargs["identity"] = args.use_identity
+    try:
+        bundler = run_txt2detection(**kwargs)
+    except (ValidationError, ValueError) as e:
+        logging.error(f"Validate sigma file failed: {str(e)}")
+        if isinstance(e, ValidationError):
+            full_error = e.json(indent=4)
+            logging.debug(f"Validate sigma file failed: {full_error}", exc_info=True)
+        sys.exit(19)
+
+    output_dir = Path("./output") / str(bundler.bundle.id)
+    shutil.rmtree(output_dir, ignore_errors=True)
+    rules_dir = output_dir / "rules"
+    rules_dir.mkdir(exist_ok=True, parents=True)
+
+    output_path = output_dir / "bundle.json"
+    data_path = output_dir / f"data.json"
+    output_path.write_text(bundler.to_json())
+    data_path.write_text(bundler.data.model_dump_json(indent=4))
+    for obj in bundler.bundle["objects"]:
+        if obj["type"] != "indicator" or obj["pattern_type"] != "sigma":
+            continue
+        rule_id: str = obj["id"].replace("indicator--", "")
+        rule_path = rules_dir / ("rule--" + rule_id + ".yml")
+        nav_path = rules_dir / f"attack-enterprise-navigator-layer-rule--{rule_id}.json"
+        rule_path.write_text(obj["pattern"])
+        if rule_nav := (
+            bundler.data.navigator_layer and bundler.data.navigator_layer.get(rule_id)
+        ):
+            nav_path.write_text(json.dumps(rule_nav, indent=4))
+    logging.info(f"Writing bundle output to `{output_path}`")

txt2detection/ai_extractor/__init__.py

@@ -0,0 +1,16 @@
+import logging
+
+import dotenv
+
+from .base import _ai_extractor_registry as ALL_AI_EXTRACTORS
+
+from .base import BaseAIExtractor
+
+class ModelError(Exception):
+    pass
+
+for path in ["openai", "anthropic", "gemini", "deepseek", "openrouter"]:
+    try:
+        __import__(__package__ + "." + path)
+    except Exception as e:
+        logging.warning("%s not supported, please install missing modules", path, exc_info=True)

txt2detection/ai_extractor/anthropic.py

@@ -0,0 +1,12 @@
+
+import logging
+import os
+from .base import BaseAIExtractor
+from llama_index.llms.anthropic import Anthropic
+
+
+class AnthropicAIExtractor(BaseAIExtractor, provider="anthropic"):
+    def __init__(self, **kwargs) -> None:
+        kwargs.setdefault('temperature', float(os.environ.get('TEMPERATURE', 0.0)))
+        self.llm = Anthropic(max_tokens=4096, system_prompt=self.system_prompt, **kwargs)
+        super().__init__()

txt2detection/ai_extractor/base.py

@@ -0,0 +1,72 @@
+import logging
+from typing import Type
+from llama_index.core.program import LLMTextCompletionProgram
+
+import textwrap
+from llama_index.core.llms.llm import LLM
+
+from txt2detection.ai_extractor import prompts
+
+from txt2detection.ai_extractor.utils import ParserWithLogging
+from txt2detection.models import DetectionContainer, DetectionContainer
+from llama_index.core.utils import get_tokenizer
+
+
+_ai_extractor_registry: dict[str, "Type[BaseAIExtractor]"] = {}
+
+
+class BaseAIExtractor:
+    llm: LLM
+    system_prompt = textwrap.dedent(
+        """
+    <persona>
+
+        You are a cyber-security detection engineering tool responsible for analysing intelligence reports provided in text files and writing SIGMA detection rules to detect the content being described in the reports.
+
+        You have a deep understanding of cybersecurity tools like SIEMs and XDRs, as well as threat intelligence concepts.
+
+        IMPORTANT: You must always deliver your work as a computer-parsable output in JSON format. All output from you will be parsed with pydantic for further processing.
+        
+    </persona>
+    """
+    )
+
+    def get_detections(self, input_text) -> DetectionContainer:
+        logging.info("getting detections")
+
+        return LLMTextCompletionProgram.from_defaults(
+            output_parser=ParserWithLogging(DetectionContainer),
+            prompt=prompts.SIEMRULES_PROMPT,
+            verbose=True,
+            llm=self.llm,
+        )(document=input_text)
+
+    def __init__(self, *args, **kwargs) -> None:
+        pass
+
+    def count_tokens(self, input_text):
+        logging.info(
+            "unsupported model `%s`, estimating using llama-index's default tokenizer",
+            self.extractor_name,
+        )
+        return len(get_tokenizer()(input_text))
+
+    def __init_subclass__(cls, /, provider, register=True, **kwargs):
+        super().__init_subclass__(**kwargs)
+        if register:
+            cls.provider = provider
+            _ai_extractor_registry[provider] = cls
+
+    @property
+    def extractor_name(self):
+        return f"{self.provider}:{self.llm.model}"
+
+    def check_credential(self):
+        try:
+            return "authorized" if self._check_credential() else "unauthorized"
+        except:
+            return "unknown"
+
+    def _check_credential(self):
+        self.llm.complete("say 'hi'")
+        return True

txt2detection/ai_extractor/deepseek.py

@@ -0,0 +1,20 @@
+import logging
+import os
+
+from .base import BaseAIExtractor
+from llama_index.llms.deepseek import DeepSeek
+
+
+class DeepseekExtractor(BaseAIExtractor, provider="deepseek"):
+    def __init__(self, **kwargs) -> None:
+        kwargs.setdefault("temperature", float(os.environ.get("TEMPERATURE", 0.0)))
+        kwargs.setdefault("model", "deepseek-chat")
+        self.llm = DeepSeek(system_prompt=self.system_prompt, **kwargs)
+        super().__init__()
+
+    def count_tokens(self, text):
+        try:
+            return len(self.llm._tokenizer.encode(text))
+        except Exception as e:
+            logging.warning(e)
+            return super().count_tokens(text)

txt2detection/ai_extractor/gemini.py

@@ -0,0 +1,18 @@
+   
+import os
+from .base import BaseAIExtractor
+from llama_index.llms.gemini import Gemini
+
+
+class GeminiAIExtractor(BaseAIExtractor, provider="gemini"):
+    def __init__(self, **kwargs) -> None:
+        kwargs.setdefault('temperature', float(os.environ.get('TEMPERATURE', 0.0)))
+        self.llm = Gemini(max_tokens=4096, **kwargs)
+        super().__init__()
+
+    def count_tokens(self, text):
+        return self.llm._model.count_tokens(text).total_tokens
+    
+    @property
+    def extractor_name(self):
+        return f"{self.provider}:{self.llm.model}"

txt2detection/ai_extractor/openai.py

@@ -0,0 +1,18 @@
+import logging
+import os
+from .base import BaseAIExtractor
+from llama_index.llms.openai import OpenAI
+
+
+class OpenAIExtractor(BaseAIExtractor, provider="openai"):
+    def __init__(self, **kwargs) -> None:
+        kwargs.setdefault("temperature", float(os.environ.get("TEMPERATURE", 0.0)))
+        self.llm = OpenAI(system_prompt=self.system_prompt, **kwargs, max_tokens=4096)
+        super().__init__()
+
+    def count_tokens(self, text):
+        try:
+            return len(self.llm._tokenizer.encode(text))
+        except Exception as e:
+            logging.warning(e)
+            return super().count_tokens(text)

txt2detection/ai_extractor/openrouter.py

@@ -0,0 +1,20 @@
+import logging
+import os
+from .base import BaseAIExtractor
+from llama_index.llms.openrouter import OpenRouter
+
+
+class OpenRouterExtractor(BaseAIExtractor, provider="openrouter"):
+    def __init__(self, **kwargs) -> None:
+        kwargs.setdefault("temperature", float(os.environ.get("TEMPERATURE", 0.0)))
+        self.llm = OpenRouter(
+            system_prompt=self.system_prompt, max_tokens=4096, **kwargs
+        )
+        super().__init__()
+
+    def count_tokens(self, text):
+        try:
+            return len(self.llm._tokenizer.encode(text))
+        except Exception as e:
+            logging.warning(e)
+            return super().count_tokens(text)

txt2detection/ai_extractor/prompts.py

@@ -0,0 +1,121 @@
+from llama_index.core import PromptTemplate, ChatPromptTemplate
+import textwrap
+from llama_index.core.base.llms.types import ChatMessage, MessageRole
+
+
+SIEMRULES_PROMPT = ChatPromptTemplate(
+    [
+        ChatMessage.from_str(
+            """
+**Persona:**  
+
+You are an expert in cybersecurity threat detection. Given a structured security report, generate a Sigma rule following the Sigma specification.  
+
+## **Requirements:**  
+Return the result as a **JSON output**, ensuring that each dictionary represents a Sigma rule with the following **AI-generated fields**:  
+
+### **Meta Properties (Generated by AI Only)**  
+- `"title"`: A concise, descriptive title for the rule.  
+- `"description"`: A summary of the rule, explaining its purpose and detection logic.  
+- `"tags"`: **Generated by AI**, including:  
+  - ATT&CK Technique IDs, ATT&CK Sub-technique IDs and CVE IDs relevant to the report.
+- `"falsepositives"`: Please describe situations where this detection rule might trigger false positive detections. Put each situation description as a new list item
+- `"logsources"`: Valid sigma rule logsource
+- `"detection"`: Valid sigma rule detection
+- `"indicator_types"`: One or more STIX v2.1 indicator.indicator_types
+- `"level"`: select SIGMA level for the rule
+
+                         
+### **Indicator Types**
+- `"anomalous-activity"`: Unexpected, or unusual activity that may not necessarily be malicious or indicate compromise. This type of activity may include reconnaissance-like behavior such as port scans or version identification, network behavior anomalies, and asset and/or user behavioral anomalies.
+- `"anonymization"`: Suspected anonymization tools or infrastructure (proxy, TOR, VPN, etc.).
+- `"benign"`: Activity that is not suspicious or malicious in and of itself, but when combined with other activity may indicate suspicious or malicious behavior.
+- `"compromised"`: Assets that are suspected to be compromised.
+- `"malicious-activity"`: Patterns of suspected malicious objects and/or activity.
+- `"attribution"`: Patterns of behavior that indicate attribution to a particular Threat Actor or Campaign.
+- `"unknown"`: There is not enough information available to determine the type of indicator.
+
+### **Level**
+The level field contains one of five string values. It describes the criticality of a triggered rule. While low and medium level events have an informative character, events with high and critical level should lead to immediate reviews by security analysts.
+- informational: Rule is intended for enrichment of events, e.g. by tagging them. No case or alerting should be triggered by such rules because it is expected that a huge amount of events will match these rules.
+- low: Notable event but rarely an incident. Low rated events can be relevant in high numbers or combination with others. Immediate reaction shouldn't be necessary, but a regular review is recommended.
+- medium: Relevant event that should be reviewed manually on a more frequent basis.
+- high: Relevant event that should trigger an internal alert and requires a prompt review.
+- critical: Highly relevant event that indicates an incident. Critical events should be reviewed immediately. It is used only for cases in which probability borders certainty.
+
+
+### **Tags***
+
+Tags are to follow the format <namespace>.<value>.
+IMPORTANT: select 0 or more from this section and only include valid and most appropriate tags
+
+#### Namespace: attack
+
+ATT&CK is either a [technique], a [group], a [software] or a [tactic].
+
+* *attack.t1234*: Refers to a [technique](https://attack.mitre.org/wiki/All_Techniques)
+* *attack.g1234*: Refers to a [group](https://attack.mitre.org/wiki/Groups)
+* *attack.s1234*: Refers to [software](https://attack.mitre.org/wiki/Software)
+
+Tactics:
+
+* attack.initial-access
+* attack.execution
+* attack.persistence
+* attack.privilege-escalation
+* attack.defense-evasion
+* attack.credential-access
+* attack.discovery
+* attack.lateral-movement
+* attack.collection
+* attack.exfiltration
+* attack.command-and-control
+* attack.impact
+
+#### Namespace: cve
+Only include from this section, CVEs that are explicitly mentioned in input document. Do not attempt to make up any random CVE-ID
+
+Use the CVE tag from [MITRE](https://cve.mitre.org) in lower case separated by dots. Example tag: `cve.2021-44228`.
+                         
+### **Detection:**  
+
+You need to generate a structured `detection` object representing a set of **search identifiers** used for detecting patterns in log data. Follow these rules when constructing `detection`:  
+
+1. **Overall Structure:**  
+   - The root object must contain a **required** field called `"condition"`, which is a string that defines the relationship between different search identifiers (e.g., `"selection1 OR selection2"`).  
+   - Additional properties (besides `"condition"`) represent **search identifiers**, which can be either a **list** or a **map**.  
+
+2. **Search Identifier Format:**  
+   - A **list** (`array`) can contain:  
+     - **Strings** (e.g., `"error_code_500"`, `"user_login_failed"`).  
+     - **Integers** (e.g., `404`, `500`).  
+     - **Objects** where all values are strings.  
+   - A **map** (`object`) where all values are strings.  
+
+3. **Example Detection Data:**  
+```json
+{
+  "condition": "selection OR event_selection",
+  "selection": ["error_code_500", 404, {"key": "value"}],
+  "event_selection": {
+    "event_type": "failed_login",
+    "source_ip": "192.168.1.1"
+  }=
+}
+```
+
+Make sure your response follows this format and adheres to the rules above.
+
+
+## **Additional Instructions**  
+- Ensure the `"tags"` field includes relevant ATT&CK and CVE references based on the report content.  
+- Return a **valid JSON output** without YAML formatting for seamless processing.
+"""
+        ),
+        ChatMessage.from_str(
+            "Taking the entire input of my next message, analyze and return appropriate response",
+            MessageRole.USER,
+        ),
+        ChatMessage.from_str("{document}", MessageRole.USER),
+    ]
+)

txt2detection/ai_extractor/utils.py

@@ -0,0 +1,21 @@
+import io
+import logging
+import typing
+import json_repair
+
+from llama_index.core.output_parsers import PydanticOutputParser
+
+
+if typing.TYPE_CHECKING:
+    from txt2detection.bundler import Bundler
+
+
+class ParserWithLogging(PydanticOutputParser):
+    def parse(self, text: str):
+        f = io.StringIO()
+        print("\n" * 5 + "=================start=================", file=f)
+        print(text, file=f)
+        print("=================close=================" + "\n" * 5, file=f)
+        logging.debug(f.getvalue())
+        repaired_json = json_repair.repair_json(text)
+        return super().parse(repaired_json)