dogesec-commons 1.0.2__py3-none-any.whl

dogesec_commons/stixifier/migrations/0004_profile_identity_id.py

@@ -0,0 +1,18 @@
+# Generated by Django 5.1.5 on 2025-05-01 09:48
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('dogesec_stixifier', '0003_rename_ai_content_check_variable_profile_ai_content_check_provider_and_more'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='profile',
+            name='identity_id',
+            field=models.CharField(default=None, max_length=46, null=True),
+        ),
+    ]

dogesec_commons/stixifier/migrations/0005_profile_generate_pdf.py

@@ -0,0 +1,18 @@
+# Generated by Django 5.1.5 on 2025-06-24 12:58
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('dogesec_stixifier', '0004_profile_identity_id'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='profile',
+            name='generate_pdf',
+            field=models.BooleanField(default=False),
+        ),
+    ]

dogesec_commons/stixifier/models.py

@@ -0,0 +1,57 @@
+import txt2stix
+from django.conf import settings
+from django.db import models
+from django.contrib.postgres.fields import ArrayField
+import uuid
+from functools import partial
+import txt2stix.common
+import txt2stix, txt2stix.extractions
+from django.core.exceptions import ValidationError
+
+
+class RelationshipMode(models.TextChoices):
+    AI = "ai", "AI Relationship"
+    STANDARD = "standard", "Standard Relationship"
+
+def validate_extractor(types, name):
+    extractors = txt2stix.extractions.parse_extraction_config(
+            txt2stix.txt2stix.INCLUDES_PATH
+        ).values()
+    for extractor in extractors:
+        if name == extractor.slug and extractor.type in types:
+            return True
+    raise ValidationError(f"{name} does not exist", 400)
+
+
+class Profile(models.Model):
+    id = models.UUIDField(primary_key=True)
+    created = models.DateTimeField(auto_now_add=True)
+    name = models.CharField(max_length=250, unique=True)
+    identity_id = models.CharField(max_length=46, null=True, default=None)
+    extractions = ArrayField(base_field=models.CharField(max_length=256))
+    relationship_mode = models.CharField(choices=RelationshipMode.choices, max_length=20, default=RelationshipMode.STANDARD)
+    extract_text_from_image = models.BooleanField(default=False)
+    defang = models.BooleanField()
+    generate_pdf = models.BooleanField(default=False)
+    ai_settings_relationships = models.CharField(max_length=256, blank=False, null=True)
+    ai_settings_extractions = ArrayField(base_field=models.CharField(max_length=256), default=list)
+    ai_content_check_provider = models.CharField(default=None, null=True, blank=False, max_length=256)
+    ai_summary_provider = models.CharField(max_length=256, blank=False, null=True)
+    ai_create_attack_flow = models.BooleanField(default=True)
+    ignore_image_refs = models.BooleanField(default=True)
+    ignore_link_refs  = models.BooleanField(default=True)
+    ignore_extraction_boundary  = models.BooleanField(default=False)
+
+    #############
+    ignore_embedded_relationships_sro = models.BooleanField(default=True)
+    ignore_embedded_relationships_smo = models.BooleanField(default=True)
+    ignore_embedded_relationships     = models.BooleanField(default=False)
+
+
+    def save(self, *args, **kwargs) -> None:
+        if not self.id:
+            name = self.name
+            if self.identity_id:
+                name = f"{self.name}+{self.identity_id}"
+            self.id = uuid.uuid5(settings.STIXIFIER_NAMESPACE, name)
+        return super().save(*args, **kwargs)

dogesec_commons/stixifier/serializers.py

@@ -0,0 +1,192 @@
+import argparse
+import contextlib
+from functools import partial
+import uuid
+from rest_framework import serializers
+
+from . import conf
+from .models import Profile
+from rest_framework import serializers
+import txt2stix.extractions
+import txt2stix.txt2stix
+from urllib.parse import urljoin
+from django.conf import settings
+from django.contrib.postgres.fields import ArrayField
+from rest_framework.validators import ValidationError
+from dogesec_commons.utils.serializers import CommonErrorSerializer
+
+from drf_spectacular.utils import OpenApiResponse, OpenApiExample
+
+from drf_spectacular.utils import OpenApiResponse, OpenApiExample
+
+from django.db import models
+
+
+def validate_model(model):
+    if not model:
+        return None
+    try:
+        extractor = txt2stix.txt2stix.parse_model(model)
+    except BaseException as e:
+        raise ValidationError(f"invalid model: {model}")
+    return model
+
+def validate_extractor(typestr, types, name):
+    extractors = txt2stix.extractions.parse_extraction_config(
+            txt2stix.txt2stix.INCLUDES_PATH
+        )
+    if  name not in extractors or extractors[name].type not in types:
+        raise ValidationError(f"`{name}` is not a valid {typestr}", 400)
+    
+def validate_stix_id(stix_id: str, type: str):
+    type_part, _, id_part = stix_id.partition('--')
+    if type_part != type:
+        raise ValidationError(f"Invalid STIX ID for type `{type}`")
+    with contextlib.suppress(Exception):
+        uuid.UUID(id_part)
+        return stix_id
+    raise ValidationError("Invalid STIX ID")
+
+
+def uses_ai(slugs):
+    extractors = txt2stix.extractions.parse_extraction_config(
+            txt2stix.txt2stix.INCLUDES_PATH
+        )
+    ai_based_extractors = []
+    for slug in slugs:
+        if extractors[slug].type == 'ai':
+            ai_based_extractors.append(slug)
+    
+    if ai_based_extractors:
+        raise ValidationError(f'AI based extractors `{ai_based_extractors}` used when `ai_settings_extractions` is not configured')
+
+class ProfileSerializer(serializers.ModelSerializer):
+    id = serializers.UUIDField(read_only=True)
+    identity_id = serializers.CharField(
+        max_length=46,
+        validators=[lambda stix_id: validate_stix_id(stix_id, "identity")],
+        allow_null=True,
+        required=False,
+        help_text="STIX Identity ID (e.g `identity--19686d47-3a50-48a0-8ef0-f3e0f8a4bd99`)"
+    )
+
+    ai_settings_relationships = serializers.CharField(
+        validators=[validate_model],
+        help_text='(required if AI relationship enabled): passed in format `provider:model`. Can only pass one model at this time.',
+        allow_null=True,
+        required=False,
+    )
+    ai_settings_extractions = serializers.ListField(
+        child=serializers.CharField(max_length=256, validators=[validate_model]),
+        help_text='(required if AI extractions enabled) passed in format provider[:model] e.g. openai:gpt4o. Can pass more than one value to get extractions from multiple providers. model part is optional',
+        required=False,
+    )
+    ai_summary_provider = serializers.CharField(
+        validators=[validate_model],
+        help_text='you can optionally get an AI model to produce a summary of the blog. You must pass the request in format `provider:model`. model part is optional',
+        allow_null=True,
+        required=False,
+    )
+    ai_content_check_provider = serializers.CharField(
+        max_length=256, validators=[validate_model],
+        allow_null=True,
+        required=False,
+        help_text='check content before proceeding'
+    )
+    ai_create_attack_flow = serializers.BooleanField(required=False, help_text="should create attack-flow (default is `false`)", default=True)
+    extractions = serializers.ListField(
+        min_length=1,
+        child=serializers.CharField(max_length=256, validators=[partial(validate_extractor, 'extractor', ["ai", "pattern", "lookup"])]),
+        help_text="extraction id(s)",
+    )
+    defang = serializers.BooleanField(help_text='If the text should be defanged before processing')
+
+
+    ignore_embedded_relationships     = serializers.BooleanField(required=False, help_text="applies to SDO and SCO types (default is `false`)")
+    ignore_embedded_relationships_sro = serializers.BooleanField(required=False, help_text="sets wether to ignore embedded refs on `relationship` object types (default is `true`)")
+    ignore_embedded_relationships_smo = serializers.BooleanField(required=False, help_text="sets wether to ignore embedded refs on SMO object types (`marking-definition`, `extension-definition`, `language-content`) (default is `true`)")
+    generate_pdf = serializers.BooleanField(required=False, help_text="Whether or not to generate pdf file for input, applies to both stixify and obstracts (default is `false`)")
+
+    class Meta:
+        model = Profile
+        fields = "__all__"
+
+    def validate(self, attrs):
+        if not attrs.get('ai_settings_relationships'):
+            if attrs['relationship_mode'] == 'ai':
+                raise ValidationError('`ai_settings_relationships` is required when `relationship_mode == "ai"`')
+            if attrs['ai_create_attack_flow']:
+                raise ValidationError('`ai_settings_relationships` is required when `ai_create_attack_flow == true`')
+        if not attrs.get('ai_settings_extractions'):
+            uses_ai(attrs['extractions'])
+        return super().validate(attrs)
+
+
+
+DEFAULT_400_ERROR = OpenApiResponse(
+    CommonErrorSerializer,
+    "The server did not understand the request",
+    [
+        OpenApiExample(
+            "http400",
+            {"message": " The server did not understand the request", "code": 400},
+        )
+    ],
+)
+
+
+DEFAULT_404_ERROR = OpenApiResponse(
+    CommonErrorSerializer,
+    "Resource not found",
+    [
+        OpenApiExample(
+            "http404",
+            {
+                "message": "The server cannot find the resource you requested",
+                "code": 404,
+            },
+        )
+    ],
+)
+
+
+##
+
+
+
+class Txt2stixExtractorSerializer(serializers.Serializer):
+    id = serializers.CharField(label='The `id` of the extractor')
+    name = serializers.CharField()
+    type = serializers.CharField()
+    description = serializers.CharField(required=False, allow_null=True)
+    notes = serializers.CharField(required=False, allow_null=True)
+    file = serializers.CharField(required=False, allow_null=True)
+    created = serializers.CharField(required=False, allow_null=True)
+    modified = serializers.CharField(required=False, allow_null=True)
+    created_by = serializers.CharField(required=False, allow_null=True)
+    version = serializers.CharField()
+    stix_mapping = serializers.CharField(required=False, allow_null=True)
+    dogesec_web = serializers.BooleanField(required=False, allow_null=True)
+
+    @classmethod
+    def all_extractors(cls, types):
+        retval = {}
+        extractors = txt2stix.extractions.parse_extraction_config(
+            txt2stix.txt2stix.INCLUDES_PATH
+        ).values()
+        for extractor in extractors:
+            if extractor.type in types:
+                retval[extractor.slug] = cls.cleanup_extractor(extractor)
+                if extractor.file:
+                    retval[extractor.slug]["file"] = urljoin(conf.TXT2STIX_INCLUDE_URL, str(extractor.file.relative_to(txt2stix.txt2stix.INCLUDES_PATH)))
+        return retval
+    
+    @classmethod
+    def cleanup_extractor(cls, dct: dict):
+        KEYS = cls(data={}).get_fields()
+        retval = {"id": dct["slug"]}
+        for key in KEYS:
+            if key in dct:
+                retval[key] = dct[key]
+        return retval
+

dogesec_commons/stixifier/stixifier.py

@@ -0,0 +1,252 @@
+import io
+import json
+import logging
+import os
+from pathlib import Path
+import shutil
+import uuid
+from attr import dataclass
+
+from dogesec_commons.stixifier.summarizer import parse_summarizer_model
+
+from ..objects import db_view_creator
+from . import models
+import tempfile
+from file2txt.converter import get_parser_class
+from txt2stix import get_include_path
+from txt2stix.stix import txt2stixBundler
+from txt2stix.ai_extractor import BaseAIExtractor
+from stix2arango.stix2arango import Stix2Arango
+from django.conf import settings
+from txt2stix.ai_extractor.utils import DescribesIncident
+
+
+from file2txt.converter import Fanger, get_parser_class
+from file2txt.parsers.core import BaseParser
+import txt2stix.utils
+import txt2stix.txt2stix
+import txt2stix.extractions
+
+
+def all_extractors(names, _all=False):
+    retval = {}
+    extractors = txt2stix.extractions.parse_extraction_config(
+        get_include_path()
+    ).values()
+    for extractor in extractors:
+        if _all or extractor.slug in names:
+            retval[extractor.slug] = extractor
+    return retval
+
+
+@dataclass
+class ReportProperties:
+    name: str = None
+    identity: dict = None
+    tlp_level: str = None
+    confidence: int = None
+    labels: list[str] = None
+    created: str = None
+    kwargs: dict = {}
+
+
+class StixifyProcessor:
+    def __init__(
+        self,
+        file: io.FileIO,
+        profile: models.Profile,
+        job_id: uuid.UUID,
+        post=None,
+        file2txt_mode="html",
+        report_id=None,
+        base_url=None,
+        always_extract=False,
+    ) -> None:
+        self.job_id = str(job_id)
+        self.extra_data = dict()
+        self.report_id = report_id
+        self.profile = profile
+        self.collection_name = "stixify"
+        self.tmpdir = Path(tempfile.mkdtemp(prefix="stixify-"))
+        self.file2txt_mode = file2txt_mode
+        self.md_images = []
+        self.processed_image_base_url = ""
+        self.base_url = base_url
+        self.incident: DescribesIncident = None
+        self.summary = None
+
+        self.filename = self.tmpdir / Path(file.name).name
+        self.filename.write_bytes(file.read())
+
+        self.task_name = f"{self.profile.name}/{job_id}/{self.report_id}"
+        self.always_extract = always_extract
+
+    def setup(self, /, report_prop: ReportProperties, extra={}):
+        self.extra_data.update(extra)
+        self.report_prop = report_prop
+
+    def file2txt(self):
+        parser_class = get_parser_class(self.file2txt_mode, self.filename.name)
+        converter: BaseParser = parser_class(
+            self.filename,
+            self.file2txt_mode,
+            self.profile.extract_text_from_image,
+            settings.GOOGLE_VISION_API_KEY,
+            base_url=self.base_url,
+        )
+        output = converter.convert(
+            processed_image_base_url=self.processed_image_base_url
+        )
+        if self.profile.defang:
+            output = Fanger(output).defang()
+        for name, img in converter.images.items():
+            img_file = io.BytesIO()
+            img_file.name = name
+            img.save(img_file, format="png")
+            self.md_images.append(img_file)
+
+        self.output_md = output
+        self.md_file = self.tmpdir / f"post_md_{self.report_id or 'file'}.md"
+        self.md_file.write_text(self.output_md)
+
+    def txt2stix(self):
+        extractors = all_extractors(self.profile.extractions)
+        extractors_map = {}
+        for extractor in extractors.values():
+            if extractors_map.get(extractor.type):
+                extractors_map[extractor.type][extractor.slug] = extractor
+            else:
+                extractors_map[extractor.type] = {extractor.slug: extractor}
+
+        self.bundler = txt2stixBundler(
+            self.report_prop.name,
+            identity=self.report_prop.identity,
+            tlp_level=self.report_prop.tlp_level,
+            confidence=self.report_prop.confidence,
+            labels=self.report_prop.labels,
+            description=self.output_md,
+            extractors=extractors,
+            report_id=self.report_id,
+            created=self.report_prop.created,
+            **self.report_prop.kwargs,
+        )
+        self.extra_data["_stixify_report_id"] = str(self.bundler.report.id)
+        input_text = txt2stix.utils.remove_links(
+            self.output_md,
+            self.profile.ignore_image_refs,
+            self.profile.ignore_link_refs,
+        )
+        ai_extractors = [
+            txt2stix.txt2stix.parse_model(model_str)
+            for model_str in self.profile.ai_settings_extractions
+        ]
+        self.txt2stix_data = txt2stix.txt2stix.run_txt2stix(
+            self.bundler,
+            input_text,
+            extractors_map,
+            ai_content_check_provider=self.profile.ai_content_check_provider
+            and txt2stix.txt2stix.parse_model(self.profile.ai_content_check_provider),
+            ai_create_attack_flow=self.profile.ai_create_attack_flow,
+            input_token_limit=settings.INPUT_TOKEN_LIMIT,
+            ai_settings_extractions=ai_extractors,
+            ai_settings_relationships=self.profile.ai_settings_relationships
+            and txt2stix.txt2stix.parse_model(self.profile.ai_settings_relationships),
+            relationship_mode=self.profile.relationship_mode,
+            ignore_extraction_boundary=self.profile.ignore_extraction_boundary,
+            always_extract=self.always_extract,
+        )
+        self.incident = self.txt2stix_data.content_check
+        return self.bundler
+
+    def summarize(self):
+        if self.profile.ai_summary_provider:
+            logging.info(
+                f"summarizing report {self.report_id} using `{self.profile.ai_summary_provider}`"
+            )
+            try:
+                report = self.bundler.report
+                summary_extractor = parse_summarizer_model(
+                    self.profile.ai_summary_provider
+                )
+                self.summary = summary_extractor.summarize(self.output_md)
+                summary_note_obj = {
+                    "type": "note",
+                    "spec_version": "2.1",
+                    "id": report.id.replace("report", "note"),
+                    "created": report.created,
+                    "modified": report.modified,
+                    "created_by_ref": report.created_by_ref,
+                    "external_references": [
+                        {
+                            "source_name": "txt2stix_ai_summary_provider",
+                            "external_id": self.profile.ai_summary_provider,
+                        },
+                    ],
+                    "abstract": f"AI Summary: {report.name}",
+                    "content": self.summary,
+                    "object_refs": [report.id],
+                    "object_marking_refs": report.object_marking_refs,
+                    "labels": report.labels,
+                    "confidence": report.confidence,
+                }
+
+                self.bundler.add_ref(summary_note_obj)
+                self.bundler.add_ref(
+                    self.bundler.new_relationship(
+                        summary_note_obj["id"],
+                        report.id,
+                        relationship_type="summary-of",
+                        description=f"AI generated summary for {report.name}",
+                        external_references=summary_note_obj["external_references"],
+                    )
+                )
+            except BaseException as e:
+                print(f"got err {e}")
+                logging.info(f"got err {e}", exc_info=True)
+        return self.summary
+
+    def process(self) -> str:
+        logging.info(f"running file2txt on {self.task_name}")
+        self.file2txt()
+        logging.info(f"running txt2stix on {self.task_name}")
+        bundler = self.txt2stix()
+        self.summarize()
+        self.write_bundle(bundler)
+        logging.info(f"uploading {self.task_name} to arangodb via stix2arango")
+        self.upload_to_arango()
+        return bundler.report.id
+
+    def write_bundle(self, bundler: txt2stixBundler):
+        bundle = json.loads(bundler.to_json())
+        self.bundle = json.dumps(bundle, indent=4)
+        self.bundle_file = self.tmpdir / f"bundle_{self.report_id}.json"
+        self.bundle_file.write_text(self.bundle)
+
+    def upload_to_arango(self):
+        s2a = Stix2Arango(
+            file=str(self.bundle_file),
+            database=settings.ARANGODB_DATABASE,
+            collection=self.collection_name,
+            stix2arango_note=f"stixifier-report--{self.report_id}",
+            ignore_embedded_relationships=self.profile.ignore_embedded_relationships,
+            ignore_embedded_relationships_smo=self.profile.ignore_embedded_relationships_smo,
+            ignore_embedded_relationships_sro=self.profile.ignore_embedded_relationships_sro,
+            host_url=settings.ARANGODB_HOST_URL,
+            username=settings.ARANGODB_USERNAME,
+            password=settings.ARANGODB_PASSWORD,
+        )
+        s2a.arangodb_extra_data.update(self.extra_data)
+        db_view_creator.link_one_collection(
+            s2a.arango.db,
+            settings.ARANGODB_DATABASE_VIEW,
+            f"{self.collection_name}_edge_collection",
+        )
+        db_view_creator.link_one_collection(
+            s2a.arango.db,
+            settings.ARANGODB_DATABASE_VIEW,
+            f"{self.collection_name}_vertex_collection",
+        )
+        s2a.run()
+
+    def __del__(self):
+        shutil.rmtree(self.tmpdir)

dogesec_commons/stixifier/summarizer.py

@@ -0,0 +1,62 @@
+from typing import Type
+from llama_index.core import PromptTemplate
+from txt2stix.ai_extractor import ALL_AI_EXTRACTORS, BaseAIExtractor
+from llama_index.core.response_synthesizers import SimpleSummarize
+from rest_framework.validators import ValidationError
+
+prompt = PromptTemplate("""
+<persona>
+
+You are a cyber-security threat intelligence analyst responsible for analysing intelligence. You have a deep understanding of cybersecurity concepts and threat intelligence. You are responsible for simplifying long intelligence reports into concise summaries for other to quickly understand the contents.
+
+</persona>
+
+<requirement>
+
+Using the MARKDOWN of the report provided in <document>, provide an executive summary of it containing no more than one paragraphs.
+
+IMPORTANT: the output should be structured as markdown text.
+IMPORTANT: do not put output in code block
+
+</requirement>
+
+<accuracy>
+
+Think about your answer first before you respond.
+
+</accuracy>
+                        
+<document>
+{context_str}
+</document>
+
+""")
+
+def get_provider(klass: Type[BaseAIExtractor]):
+    class SummarizerSession(klass, provider="~"+klass.provider, register=False):
+        system_prompt = """
+            You are a cyber-security threat intelligence analyst responsible for analysing intelligence.
+            You have a deep understanding of cybersecurity concepts and threat intelligence.
+            You are responsible for simplifying long intelligence reports into concise summaries for other to quickly understand the contents.
+        """
+        def summarize(self, text):
+            summarizer = SimpleSummarize(llm=self.llm, text_qa_template=prompt)
+            return summarizer.get_response('', text_chunks=[text])
+    SummarizerSession.__name__ = klass.__name__.replace('Extractor', 'Summarizer')
+    return SummarizerSession
+
+
+def parse_summarizer_model(value: str):
+    try:
+        splits = value.split(':', 1)
+        provider = splits[0]
+        if provider not in ALL_AI_EXTRACTORS:
+            raise ValidationError(f"invalid summary provider in `{value}`, must be one of [{list(ALL_AI_EXTRACTORS)}]")
+        provider = get_provider(ALL_AI_EXTRACTORS[provider])
+        if len(splits) == 2:
+            return provider(model=splits[1])
+        return provider()
+    except ValidationError:
+        raise
+    except BaseException as e:
+        raise ValidationError(f'invalid model: {value}') from e