txt2detection 1.0.15__py3-none-any.whl

txt2detection/models.py

@@ -0,0 +1,427 @@
+import enum
+import json
+import re
+import typing
+import uuid
+import requests
+from slugify import slugify
+from datetime import date as dt_date
+from typing import Any, ClassVar, List, Literal, Optional, Union
+from uuid import UUID
+from stix2extensions.data_source import DataSource
+
+import jsonschema
+from pydantic import BaseModel, Field, computed_field, field_validator
+from pydantic_core import PydanticCustomError, core_schema
+import yaml
+
+from stix2 import (
+    MarkingDefinition,
+)
+
+
+if typing.TYPE_CHECKING:
+    from txt2detection.bundler import Bundler
+
+UUID_NAMESPACE = uuid.UUID("a4d70b75-6f4a-5d19-9137-da863edd33d7")
+
+TAG_PATTERN = re.compile(r"^[a-z0-9_-]+\.[a-z0-9._-]+$")
+
+MITRE_TACTIC_MAP = {
+    "initial-access": "TA0001",
+    "execution": "TA0002",
+    "persistence": "TA0003",
+    "privilege-escalation": "TA0004",
+    "defense-evasion": "TA0005",
+    "credential-access": "TA0006",
+    "discovery": "TA0007",
+    "lateral-movement": "TA0008",
+    "collection": "TA0009",
+    "exfiltration": "TA0010",
+    "command-and-control": "TA0011",
+    "impact": "TA0040",
+}
+
+
+class TLP_LEVEL(enum.Enum):
+    CLEAR = MarkingDefinition(
+        spec_version="2.1",
+        id="marking-definition--94868c89-83c2-464b-929b-a1a8aa3c8487",
+        created="2022-10-01T00:00:00.000Z",
+        definition_type="TLP:CLEAR",
+        extensions={
+            "extension-definition--60a3c5c5-0d10-413e-aab3-9e08dde9e88d": {
+                "extension_type": "property-extension",
+                "tlp_2_0": "clear",
+            }
+        },
+    )
+    GREEN = MarkingDefinition(
+        spec_version="2.1",
+        id="marking-definition--bab4a63c-aed9-4cf5-a766-dfca5abac2bb",
+        created="2022-10-01T00:00:00.000Z",
+        definition_type="TLP:GREEN",
+        extensions={
+            "extension-definition--60a3c5c5-0d10-413e-aab3-9e08dde9e88d": {
+                "extension_type": "property-extension",
+                "tlp_2_0": "green",
+            }
+        },
+    )
+    AMBER = MarkingDefinition(
+        spec_version="2.1",
+        id="marking-definition--55d920b0-5e8b-4f79-9ee9-91f868d9b421",
+        created="2022-10-01T00:00:00.000Z",
+        definition_type="TLP:AMBER",
+        extensions={
+            "extension-definition--60a3c5c5-0d10-413e-aab3-9e08dde9e88d": {
+                "extension_type": "property-extension",
+                "tlp_2_0": "amber",
+            }
+        },
+    )
+    AMBER_STRICT = MarkingDefinition(
+        spec_version="2.1",
+        id="marking-definition--939a9414-2ddd-4d32-a0cd-375ea402b003",
+        created="2022-10-01T00:00:00.000Z",
+        definition_type="TLP:AMBER+STRICT",
+        extensions={
+            "extension-definition--60a3c5c5-0d10-413e-aab3-9e08dde9e88d": {
+                "extension_type": "property-extension",
+                "tlp_2_0": "amber+strict",
+            }
+        },
+    )
+    RED = MarkingDefinition(
+        spec_version="2.1",
+        id="marking-definition--e828b379-4e03-4974-9ac4-e53a884c97c1",
+        created="2022-10-01T00:00:00.000Z",
+        definition_type="TLP:RED",
+        extensions={
+            "extension-definition--60a3c5c5-0d10-413e-aab3-9e08dde9e88d": {
+                "extension_type": "property-extension",
+                "tlp_2_0": "red",
+            }
+        },
+    )
+
+    @classmethod
+    def levels(cls):
+        return dict(
+            clear=cls.CLEAR,
+            green=cls.GREEN,
+            amber=cls.AMBER,
+            amber_strict=cls.AMBER_STRICT,
+            red=cls.RED,
+        )
+
+    @classmethod
+    def values(cls):
+        return [
+            cls.CLEAR.value,
+            cls.GREEN.value,
+            cls.AMBER.value,
+            cls.AMBER_STRICT.value,
+            cls.RED.value,
+        ]
+
+    @classmethod
+    def get(cls, level: "str|TLP_LEVEL"):
+        if isinstance(level, cls):
+            return level
+        level = level.lower()
+        level = level.replace("+", "_").replace("-", "_")
+        if level not in cls.levels():
+            raise Exception(f"unsupported tlp level: `{level}`")
+        return cls.levels()[level]
+
+    @property
+    def name(self):
+        return super().name.lower()
+
+
+class Statuses(enum.StrEnum):
+    stable = enum.auto()
+    test = enum.auto()
+    experimental = enum.auto()
+    deprecated = enum.auto()
+    unsupported = enum.auto()
+
+
+class Level(enum.StrEnum):
+    informational = enum.auto()
+    low = enum.auto()
+    medium = enum.auto()
+    high = enum.auto()
+    critical = enum.auto()
+
+
+class SigmaTag(str):
+    @classmethod
+    def __get_pydantic_core_schema__(
+        cls,
+        _source: type[Any],
+        _handler,
+    ) -> core_schema.CoreSchema:
+        return core_schema.no_info_after_validator_function(
+            cls._validate, core_schema.str_schema()
+        )
+
+    @classmethod
+    def __get_pydantic_json_schema__(cls, core_schema: core_schema.CoreSchema, handler):
+        field_schema = handler(core_schema)
+        field_schema.update(
+            type="string", pattern=TAG_PATTERN.pattern, format="sigma-tag"
+        )
+        return field_schema
+
+    @classmethod
+    def _validate(cls, input_value: str, /) -> str:
+        if not TAG_PATTERN.match(input_value):
+            raise PydanticCustomError(
+                "value_error",
+                "value is not a valid SIGMA tag: {reason}",
+                {
+                    "reason": f"Must be in format namespace.value and match pattern {TAG_PATTERN.pattern}"
+                },
+            )
+        return input_value
+
+
+class RelatedRule(BaseModel):
+    id: UUID
+    type: Literal["derived", "obsolete", "merged", "renamed", "similar"]
+
+
+class BaseDetection(BaseModel):
+    title: str
+    description: str
+    detection: dict
+    logsource: dict
+    status: Statuses = Statuses.experimental
+    falsepositives: list[str]
+    tags: list[str]
+    level: Level
+    _custom_id = None
+    _extra_data: dict
+    sigma_json_schema: ClassVar = requests.get(
+        "https://github.com/SigmaHQ/sigma-specification/raw/refs/heads/main/json-schema/sigma-detection-rule-schema.json"
+    ).json()
+
+    def model_post_init(self, __context):
+        self.tags = self.tags or []
+        self._extra_data = dict()
+        return super().model_post_init(__context)
+
+    @property
+    def detection_id(self):
+        return str(self._custom_id or getattr(self, "id", None) or uuid.uuid4())
+
+    @detection_id.setter
+    def detection_id(self, custom_id):
+        self._custom_id = custom_id.split("--")[-1]
+
+    @property
+    def tlp_level(self):
+        return tlp_from_tags(self.tags)
+
+    @tlp_level.setter
+    def tlp_level(self, level):
+        set_tlp_level_in_tags(self.tags, level)
+
+    def set_labels(self, labels):
+        self.tags.extend(labels)
+
+    def set_extra_data_from_bundler(self, bundler: "Bundler"):
+        raise NotImplementedError("this class should no longer be in use")
+
+    def make_rule(self, bundler: "Bundler"):
+        self.set_extra_data_from_bundler(bundler)
+        self.tags = list(dict.fromkeys(self.tags))
+
+        rule = dict(
+            id=self.detection_id,
+            **self.model_dump(
+                exclude=["indicator_types", "id"], mode="json", by_alias=True
+            ),
+        )
+        for k, v in list(rule.items()):
+            if not v:
+                rule.pop(k, None)
+
+        self.validate_rule_with_json_schema(rule)
+        if getattr(self, "date", 0):
+            rule.update(date=self.date)
+        if getattr(self, "modified", 0):
+            rule.update(modified=self.modified)
+        return yaml.dump(rule, sort_keys=False, indent=4)
+
+    def validate_rule_with_json_schema(self, rule):
+        jsonschema.validate(
+            rule,
+            self.sigma_json_schema,
+        )
+
+    @property
+    def external_references(self):
+        refs = []
+        for attr in ["level", "status", "license"]:
+            if attr_val := getattr(self, attr, None):
+                refs.append(dict(source_name=f"sigma-{attr}", description=attr_val))
+        return refs
+
+    @property
+    def mitre_attack_ids(self):
+        retval = []
+        for i, label in enumerate(self.tags):
+            label = label.replace("_", "-").lower()
+            namespace, _, label_id = label.partition(".")
+            if namespace == "attack":
+                retval.append(MITRE_TACTIC_MAP.get(label_id, label_id.upper()))
+        return retval
+
+    @property
+    def cve_ids(self):
+        retval = []
+        for label in self.tags:
+            namespace, _, label_id = label.partition(".")
+            if namespace == "cve":
+                retval.append(namespace.upper() + "-" + label_id)
+        return retval
+
+    def make_data_source(self):
+        return DataSource(
+            category=self.logsource.get("category"),
+            product=self.logsource.get("product"),
+            service=self.logsource.get("service"),
+            definition=self.logsource.get("definition"),
+        )
+
+
+class AIDetection(BaseDetection):
+    indicator_types: list[str] = Field(default_factory=list)
+
+    def to_sigma_rule_detection(self, bundler):
+        rule_dict = {
+            **self.model_dump(exclude=["indicator_types"]),
+            **dict(
+                date=bundler.report.created.date(),
+                modified=bundler.report.modified.date(),
+                id=uuid.uuid4(),
+            ),
+        }
+        try:
+            return SigmaRuleDetection.model_validate(rule_dict)
+        except Exception as e:
+            raise ValueError(
+                dict(message="validate ai output failed", error=e, content=rule_dict)
+            )
+
+
+class SigmaRuleDetection(BaseDetection):
+    title: str
+    id: Optional[UUID] = None
+    related: Optional[list[RelatedRule]] = None
+    name: Optional[str] = None
+    taxonomy: Optional[str] = None
+    status: Optional[Statuses] = None
+    description: Optional[str] = None
+    license: Optional[str] = None
+    author: Optional[str] = None
+    references: Optional[List[str]] = Field(default_factory=list)
+    date: Optional["dt_date"] = Field(alias="date", default=None)
+    modified: Optional["dt_date"] = None
+    logsource: dict
+    detection: dict
+    fields: Optional[List[str]] = None
+    falsepositives: Optional[List[str]] = None
+    level: Optional[Level] = None
+    tags: Optional[List[SigmaTag]] = Field(default_factory=list)
+    scope: Optional[List[str]] = None
+    _indicator_types: list = None
+
+    @property
+    def detection_id(self):
+        return str(self.id)
+
+    @property
+    def indicator_types(self):
+        return self._indicator_types
+
+    @indicator_types.setter
+    def indicator_types(self, types):
+        self._indicator_types = types
+
+    @detection_id.setter
+    def detection_id(self, new_id):
+        if self.id and str(self.id) != str(new_id):
+            self.related = self.related or []
+            self.related.append(RelatedRule(id=self.id, type="renamed"))
+        self.id = new_id
+
+    @field_validator("tags", mode="after")
+    @classmethod
+    def validate_tlp(cls, tags: list[str]):
+        tlps = []
+        for tag in tags:
+            if tag.startswith("tlp."):
+                tlps.append(tag)
+        if len(tlps) > 1:
+            raise ValueError(
+                f"tag must not contain more than one tag in tlp namespace. Got {tlps}"
+            )
+        return tags
+
+    @field_validator("modified", mode="after")
+    @classmethod
+    def validate_modified(cls, modified, info):
+        if info.data.get("date") == modified:
+            return None
+        return modified
+
+    def set_extra_data_from_bundler(self, bundler: "Bundler"):
+        if not bundler:
+            return
+
+        if not self.date:
+            from .utils import as_date
+
+            self.date = as_date(bundler.created)
+
+        self.set_labels(bundler.labels)
+        self.tlp_level = bundler.tlp_level.name
+        self.author = bundler.report.created_by_ref
+        self.license = bundler.license
+        self.references = bundler.reference_urls
+
+
+class DetectionContainer(BaseModel):
+    success: bool
+    detections: list[Union[BaseDetection, AIDetection, SigmaRuleDetection]]
+
+
+class DataContainer(BaseModel):
+    detections: DetectionContainer
+    navigator_layer: dict = Field(default=None)
+    observables: list[dict] = Field(default=None)
+    cves: dict[str, str] = Field(default_factory=dict)
+    attacks: dict[str, str] = Field(default_factory=dict)
+
+
+def tlp_from_tags(tags: list[SigmaTag]):
+    for tag in tags:
+        ns, _, level = tag.partition(".")
+        if ns != "tlp":
+            continue
+        if tlp_level := TLP_LEVEL.get(level.replace("-", "_")):
+            return tlp_level
+    return None
+
+
+def set_tlp_level_in_tags(tags: list[SigmaTag], level):
+    level = str(level)
+    for i, tag in enumerate(tags):
+        if tag.startswith("tlp."):
+            tags.remove(tag)
+    tags.append("tlp." + level.replace("_", "-"))
+    return tags

txt2detection/observables.py

@@ -0,0 +1,161 @@
+import re, validators
+from typing import Any, Dict, List
+from stix2 import parse as parse_stix, parse_observable
+
+# Mapping of key regex patterns to STIX observable types
+STIX_PATTERNS_KEYS = {
+    "ipv4-addr": r"(?i)\b(ip|ipv4)\b",
+    "ipv6-addr": r"(?i)\bipv6\b",
+    "email-addr": r"(?i)\bemail\b",
+    "url": r"(?i)\b(url|uri)\b",
+    "directory": r"(?i)\b(directory|path)\b",
+    "domain-name": r"(?i)\bdomain\b",
+    "hostname": r"(?i)\bhost\b",
+    "file.hashes.MD5": r"(?i)\bmd5\b",
+    "file.hashes.SHA-1": r"(?i)\bsha1\b",
+    "file.hashes.SHA-256": r"(?i)\bsha256\b",
+    "file.hashes.SHA-512": r"(?i)\bsha512\b",
+    "file.hashes.SSDEEP": r"(?i)\bssdeep\b",
+    "mac-addr": r"(?i)\bmac\b",
+    "user-account": r"(?i)\buser\b",
+    "windows-registry-key": r"(?i)\bregistry\b",
+    "x509-certificate": r"(?i)\bx509\b",
+}
+
+# Mapping of value regex patterns to STIX observable types
+STIX_PATTERNS_VALUES = {
+    "ipv4-addr": [r"\b(?:(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)(?:\.|$)){4}\b"],
+    "ipv6-addr": [r"\b(?:[A-F0-9]{1,4}:){7}[A-F0-9]{1,4}\b"],
+    "email-addr": [r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}\b"],
+    "url": [r"\bhttps?://[^\s/$.?#].[^\x00\s]*\b"],
+    "directory": [r"(?:[A-Za-z]:)?(?:\\\\[^\\\\:*?\"<>|\r\n]+)+\\\\?"],
+    "domain-name": [r"\b(?:[a-zA-Z0-9-]+\.)+[a-zA-Z]{2,}\b"],
+    "hostname": [r"\b[a-zA-Z0-9-]{1,63}(\.[a-zA-Z0-9-]{1,63})*\b"],
+    "file.hashes.MD5": [r"\b[a-fA-F0-9]{32}\b"],
+    "file.hashes.SHA-1": [r"\b[a-fA-F0-9]{40}\b"],
+    "file.hashes.SHA-256": [r"\b[a-fA-F0-9]{64}\b"],
+    "file.hashes.SHA-512": [r"\b[a-fA-F0-9]{128}\b"],
+    "file.hashes.SSDEEP": [r"\b\d{1,}:[A-Za-z0-9/+]{10,}:[A-Za-z0-9/+]{10,}\b"],
+    "mac-addr": [r"\b([0-9A-Fa-f]{2}[:-]){5}([0-9A-Fa-f]{2})\b"],
+    "user-account": [r"\b[A-Za-z0-9._%-]{3,}\\\\?[A-Za-z0-9._%-]{3,}\b"],
+    "windows-registry-key": [r"HK\w{0,2}_[A-Z_]+\\.*"],
+    "x509-certificate": [r"-----BEGIN CERTIFICATE-----.+?-----END CERTIFICATE-----"],
+}
+
+
+def filter_out(type, value: str):
+    match type:
+        case "ipv4-addr":
+            return validators.ipv4(value)
+        case "ipv6-addr":
+            return validators.ipv6(value)
+        case "email-addr":
+            return validators.email(value)
+        case "url":
+            return validators.url(value)
+        case "domain-name":
+            return validators.domain(value, consider_tld=True)
+        case "file.hashes.MD5":
+            return validators.hashes.md5(value)
+        case "file.hashes.SHA-1":
+            return validators.hashes.sha1(value)
+        case "file.hashes.SHA-256":
+            return validators.hashes.sha256(value)
+        case "file.hashes.SHA-512":
+            return validators.hashes.sha512(value)
+        case "file.hashes.SSDEEP":
+            pass
+        case "mac-addr":
+            return validators.mac_address(value)
+        case "user-account":
+            pass
+
+        case "windows-registry-key":
+            print(value)
+            ns, _, _ = value.partition("\\")
+            return ns in [
+                "HKEY_CLASSES_ROOT",
+                "HKCR",
+                "HKEY_CURRENT_USER",
+                "HKCU",
+                "HKEY_LOCAL_MACHINE",
+                "HKLM",
+                "HKEY_USERS",
+                "HKU",
+                "HKEY_CURRENT_CONFIG",
+                "HKCC",
+                "HKEY_PERFORMANCE_DATA",
+                "HKEY_DYN_DATA",
+            ]
+        case _:
+            return False
+    return False
+
+
+def find_stix_observables(detection: Any, matches: List[str] = None) -> List[str]:
+    if matches is None:
+        matches = []
+
+    if isinstance(detection, dict):
+        for key, value in detection.items():
+            for stix_type, key_pattern in STIX_PATTERNS_KEYS.items():
+                value_patterns = STIX_PATTERNS_VALUES.get(stix_type, [])
+                if re.search(key_pattern, key, re.IGNORECASE):
+                    for pattern in value_patterns:
+                        if isinstance(value, str) and re.search(
+                            pattern, value, re.IGNORECASE
+                        ):
+                            if filter_out(stix_type, value):
+                                matches.append((stix_type, value))
+                find_stix_observables(value, matches)
+            find_stix_observables(value, matches)
+    elif isinstance(detection, list):
+        for item in detection:
+            find_stix_observables(item, matches)
+    elif isinstance(detection, str):
+        for stix_type, value_patterns in STIX_PATTERNS_VALUES.items():
+            for pattern in value_patterns:
+                if re.search(pattern, detection, re.IGNORECASE):
+                    if filter_out(stix_type, detection):
+                        matches.append((stix_type, detection))
+    return matches
+
+
+def to_stix_object(observable_type: str, value):
+    match observable_type:
+        case (
+            "ipv4-addr"
+            | "ipv6-addr"
+            | "email-addr"
+            | "url"
+            | "domain-name"
+            | "mac-addr"
+        ):
+            return parse_observable(
+                dict(
+                    type=observable_type,
+                    value=value,
+                    spec_version="2.1",
+                )
+            )
+        case (
+            "file.hashes.MD5"
+            | "file.hashes.SHA-1"
+            | "file.hashes.SHA-256"
+            | "file.hashes.SHA-512"
+            | "file.hashes.SSDEEP"
+        ):
+            _, _, hash_type = observable_type.rpartition(".")
+            return parse_observable(
+                dict(type="file", spec_version="2.1", hashes={hash_type: value})
+            )
+
+        case "windows-registry-key":
+            return parse_observable(
+                dict(
+                    type=observable_type,
+                    spec_version="2.1",
+                    key=value,
+                )
+            )
+    return None

txt2detection/utils.py

@@ -0,0 +1,100 @@
+from datetime import date, datetime
+from functools import lru_cache
+from types import SimpleNamespace
+import uuid
+import requests
+from .ai_extractor import ALL_AI_EXTRACTORS, BaseAIExtractor, ModelError
+import logging
+
+import enum
+import logging
+import requests
+from stix2 import Identity
+
+from .models import UUID_NAMESPACE
+
+
+class DetectionLanguage(SimpleNamespace):
+    pass
+
+
+def parse_model(value: str):
+    splits = value.split(":", 1)
+    provider = splits[0]
+    if provider not in ALL_AI_EXTRACTORS:
+        raise NotImplementedError(
+            f"invalid AI provider in `{value}`, must be one of {list(ALL_AI_EXTRACTORS)}"
+        )
+    provider = ALL_AI_EXTRACTORS[provider]
+    try:
+        if len(splits) == 2:
+            return provider(model=splits[1])
+        return provider()
+    except Exception as e:
+        raise ModelError(f"Unable to initialize model `{value}`") from e
+
+
+def make_identity(name, namespace=None, created_by_ref=None, object_marking_refs=None):
+    from .bundler import Bundler
+
+    if isinstance(namespace, str):
+        namespace = uuid.UUID(namespace)
+    namespace = namespace or UUID_NAMESPACE
+    return Identity(
+        id="identity--" + str(uuid.uuid5(namespace, f"{name}")),
+        name=name,
+        created_by_ref=created_by_ref or Bundler.default_identity.id,
+        created=datetime(2020, 1, 1),
+        modified=datetime(2020, 1, 1),
+        object_marking_refs=object_marking_refs
+        or [
+            "marking-definition--94868c89-83c2-464b-929b-a1a8aa3c8487",
+            "marking-definition--a4d70b75-6f4a-5d19-9137-da863edd33d7",
+        ],
+    )
+
+
+def validate_token_count(max_tokens, input, extractor: BaseAIExtractor):
+    logging.info("INPUT_TOKEN_LIMIT = %d", max_tokens)
+    token_count = extractor.count_tokens(input)
+    logging.info("TOKEN COUNT FOR %s: %d", extractor.extractor_name, token_count)
+    if token_count > max_tokens:
+        raise Exception(
+            f"{extractor.extractor_name}: input_file token count ({token_count}) exceeds INPUT_TOKEN_LIMIT ({max_tokens})"
+        )
+
+
+@lru_cache(maxsize=5)
+def get_licenses(date):
+    resp = requests.get(
+        "https://github.com/spdx/license-list-data/raw/refs/heads/main/json/licenses.json"
+    )
+    return {l["licenseId"]: l["name"] for l in resp.json()["licenses"]}
+
+
+def valid_licenses():
+    return get_licenses(datetime.now().date().isoformat())
+
+
+def remove_rule_specific_tags(tags):
+    labels = []
+    for tag in tags:
+        namespace, _, label = tag.partition(".")
+        if namespace in ["attack", "cve", "tlp"]:
+            continue
+        labels.append(tag)
+    return labels
+
+@lru_cache()
+def load_stix_object_from_url(url):
+    resp = requests.get(url)
+    return resp.json()
+
+
+def as_date(d: "date|datetime"):
+    if isinstance(d, datetime):
+        return d.date()
+    return d
+
+
+STATUSES = ["stable", "test", "experimental", "deprecated", "unsupported"]