txt2stix 0.0.4__py3-none-any.whl → 1.0.1__py3-none-any.whl

txt2stix/ai_extractor/base.py

@@ -6,7 +6,7 @@ import textwrap
 from llama_index.core import PromptTemplate
 from llama_index.core.llms.llm import LLM
 
-from txt2stix.ai_extractor.prompts import DEFAULT_CONTENT_CHECKER_TEMPL, DEFAULT_EXTRACTION_TEMPL, DEFAULT_RELATIONSHIP_TEMPL, DEFAULT_SYSTEM_PROMPT, ATTACK_FLOW_PROMPT_TEMPL
+from txt2stix.ai_extractor.prompts import DEFAULT_CONTENT_CHECKER_WITH_SUMMARY_TEMPL, DEFAULT_EXTRACTION_TEMPL, DEFAULT_RELATIONSHIP_TEMPL, DEFAULT_SYSTEM_PROMPT, ATTACK_FLOW_PROMPT_TEMPL
 from txt2stix.ai_extractor.utils import AttackFlowList, DescribesIncident, ExtractionList, ParserWithLogging, RelationshipList, get_extractors_str
 from llama_index.core.utils import get_tokenizer
 
@@ -19,7 +19,7 @@ class BaseAIExtractor():
 
     relationship_template = DEFAULT_RELATIONSHIP_TEMPL
 
-    content_check_template = DEFAULT_CONTENT_CHECKER_TEMPL
+    content_check_template = DEFAULT_CONTENT_CHECKER_WITH_SUMMARY_TEMPL
 
     def _get_extraction_program(self):
         return LLMTextCompletionProgram.from_defaults(
@@ -63,7 +63,8 @@ class BaseAIExtractor():
         return self._get_relationship_program()(relationship_types=relationship_types, input_file=input_text, extractions=extractions)
     
     def extract_objects(self, input_text, extractors) -> ExtractionList:
-        return self._get_extraction_program()(extractors=get_extractors_str(extractors), input_file=input_text)
+        extraction_list = self._get_extraction_program()(extractors=get_extractors_str(extractors), input_file=input_text)
+        return extraction_list.model_dump().get('extractions', [])
     
     def __init__(self, *args, **kwargs) -> None:
         pass

txt2stix/ai_extractor/prompts.py

@@ -86,7 +86,7 @@ DEFAULT_RELATIONSHIP_TEMPL = PromptTemplate(textwrap.dedent(
 """
 ))
 
-DEFAULT_CONTENT_CHECKER_TEMPL = PromptTemplate("""
+DEFAULT_CONTENT_CHECKER_WITH_SUMMARY_TEMPL = PromptTemplate("""
 <persona>
     You are a cyber security threat intelligence analyst.
     Your job is to review reports that describe a cyber security incidents and/or threat intelligence.
@@ -120,6 +120,12 @@ DEFAULT_CONTENT_CHECKER_TEMPL = PromptTemplate("""
     * `indicator_of_compromise`
     * `ttp`
 </incident_classification>
+<summary>
+    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: This `summary` is different from explanation.
+    IMPORTANT: You are to simplify the long intelligence reports into concise summaries for other to quickly understand the contents.
+</summary>
 """)
 
 ATTACK_FLOW_PROMPT_TEMPL = ChatPromptTemplate([

txt2stix/ai_extractor/utils.py

@@ -34,6 +34,7 @@ class DescribesIncident(BaseModel):
     describes_incident: bool = Field(description="does the <document> include malware analysis, APT group reports, data breaches and vulnerabilities?")
     explanation: str = Field(description="Two or three sentence summary of the incidents it describes OR summary of what it describes instead of an incident")
     incident_classification : list[str] = Field(description="All the valid incident classifications that describe this document/report")
+    summary: str = Field(description="executive summary of the document containing no more than one paragraphs.")
 
 class AttackFlowItem(BaseModel):
     position : int = Field(description="order of object starting at 0")

txt2stix/bundler.py

@@ -6,6 +6,7 @@ from stix2 import (
     MarkingDefinition,
     Relationship,
     Bundle,
+    Note,
 )
 from stix2.parsing import dict_to_stix2, parse as parse_stix
 from stix2.serialization import serialize
@@ -110,6 +111,8 @@ class TLP_LEVEL(enum.Enum):
 
     @classmethod
     def get(cls, level):
+        if isinstance(level, str):
+            level = level.replace('-', '_').replace('+', '_')
         if isinstance(level, cls):
             return level
         return cls.levels()[level]
@@ -261,7 +264,7 @@ class txt2stixBundler:
     def load_stix_object_from_url(url):
         resp = requests.get(url)
         return dict_to_stix2(resp.json())
-    
+
     def add_ref(self, sdo, is_report_object=True):
         self.add_extension(sdo)
         sdo_id = sdo["id"]
@@ -277,22 +280,19 @@ class txt2stixBundler:
                 sdo_value = v
                 break
         else:
-            if refs := sdo.get('external_references', []):
-                sdo_value = refs[0]['external_id']
+            if refs := sdo.get("external_references", []):
+                sdo_value = refs[0]["external_id"]
             else:
                 sdo_value = "{NOTEXTRACTED}"
 
-
         self.id_value_map[sdo_id] = sdo_value
 
-
     def add_indicator(self, extracted_dict, add_standard_relationship):
         extractor = self.all_extractors[extracted_dict["type"]]
         stix_mapping = extractor.stix_mapping
         extracted_value = extracted_dict["value"]
         extracted_id = extracted_dict["id"]
 
-
         indicator = self.new_indicator(extractor, stix_mapping, extracted_value)
         # set id so it doesn''t need to be created in build_observables
         if extracted_dict.get("indexes"):
@@ -340,7 +340,7 @@ class txt2stixBundler:
                 },
             ],
         }
-        
+
         return indicator
 
     def add_ai_relationship(self, gpt_out):
@@ -413,11 +413,44 @@ class txt2stixBundler:
         return "indicator--" + str(
             uuid.uuid5(UUID_NAMESPACE, f"txt2stix+{self.identity['id']}+{self.report_md5}+{stix_mapping}+{value}")
         )
-    
+
+    def add_summary(self, summary, ai_summary_provider):
+        summary_note_obj = Note(
+            type="note",
+            spec_version="2.1",
+            id=self.report.id.replace("report", "note"),
+            created=self.report.created,
+            modified=self.report.modified,
+            created_by_ref=self.report.created_by_ref,
+            external_references=[
+                {
+                    "source_name": "txt2stix_ai_summary_provider",
+                    "external_id": ai_summary_provider,
+                },
+            ],
+            abstract=f"AI Summary: {self.report.name}",
+            content=summary,
+            object_refs=[self.report.id],
+            object_marking_refs=self.report.object_marking_refs,
+            labels=self.report.get('labels'),
+            confidence=self.report.get('confidence')
+        )
+
+        self.add_ref(summary_note_obj)
+        self.add_ref(
+            self.new_relationship(
+                summary_note_obj["id"],
+                self.report.id,
+                relationship_type="summary-of",
+                description=f"AI generated summary for {self.report.name}",
+                external_references=summary_note_obj["external_references"],
+            )
+        )
+
     @property
     def flow_objects(self):
         return self._flow_objects
-    
+
     @flow_objects.setter
     def flow_objects(self, objects):
         for obj in objects:
@@ -425,4 +458,4 @@ class txt2stixBundler:
                 continue
             is_report_object = obj['type'] != "extension-definition"
             self.add_ref(obj, is_report_object=is_report_object)
-        self._flow_objects = objects
+        self._flow_objects = objects

txt2stix/txt2stix.py

@@ -14,7 +14,7 @@ from txt2stix.ai_extractor.utils import DescribesIncident
 from txt2stix.attack_flow import parse_flow
 
 
-from .utils import Txt2StixData, remove_links
+from .utils import RELATIONSHIP_TYPES, Txt2StixData, remove_links
 
 from .common import UUID_NAMESPACE, FatalException
 
@@ -99,8 +99,6 @@ def parse_extractors_globbed(type, all_extractors, names):
                 pattern.load_extractor(extractor)
             filtered_extractors[extractor.type] =  extraction_processor
             extraction_processor[extractor_name] = extractor
-        except KeyError:
-            raise argparse.ArgumentTypeError(f"no such {type} slug `{extractor_name}`")
         except BaseException as e:
             raise argparse.ArgumentTypeError(f"{type} `{extractor_name}`: {e}")
     return filtered_extractors
@@ -137,22 +135,23 @@ def parse_args():
 
     inf_arg  = parser.add_argument("--input_file", "--input-file", required=True, help="The file to be converted. Must be .txt", type=Path)
     parser.add_argument("--ai_content_check_provider", required=False, type=parse_model, help="Use an AI model to check wether the content of the file contains threat intelligence. Paticularly useful to weed out vendor marketing.")
+    parser.add_argument("--always_extract", default=True, type=parse_bool, help="Whether to always extract or not depending on output of ai_content_check_provider. Default, extracts even when content_check returns describes_incident=False")
     name_arg = parser.add_argument("--name", required=True, help="Name of the file, max 124 chars", default="stix-out")
     parser.add_argument("--created", required=False, default=datetime.now(), help="Allow user to optionally pass --created time in input, which will hardcode the time used in created times")
     parser.add_argument("--ai_settings_extractions", required=False, type=parse_model, help="(required if AI extraction enabled): passed in format provider:model e.g. openai:gpt4o. Can pass more than one value to get extractions from multiple providers.", metavar="provider[:model]", nargs='+')
     parser.add_argument("--ai_settings_relationships", required=False, type=parse_model, help="(required if AI relationship enabled): passed in format `provider:model`. Can only pass one model at this time.", metavar="provider[:model]")
     parser.add_argument("--labels", type=parse_labels)
-    parser.add_argument("--relationship_mode", choices=["ai", "standard"], required=True)
+    rmode_arg = parser.add_argument("--relationship_mode", choices=["ai", "standard"], required=True)
     parser.add_argument("--report_id", type=uuid.UUID, required=False, help="id to use instead of automatically generated `{name}+{created}`", metavar="VALID_UUID")
     parser.add_argument("--confidence", type=range_type(0,100), default=None, help="value between 0-100. Default if not passed is null.", metavar="[0-100]")
     parser.add_argument("--tlp_level", "--tlp-level", choices=TLP_LEVEL.levels().keys(), default="clear", help="TLP level, default is clear")
-    parser.add_argument("--use_extractions", "--use-extractions", default={}, type=functools.partial(parse_extractors_globbed, "extractor", all_extractors),  help="Specify extraction types from the default/local extractions .yaml file", metavar="EXTRACTION1,EXTRACTION2")
+    extractions_arg = parser.add_argument("--use_extractions", "--use-extractions", default={}, type=functools.partial(parse_extractors_globbed, "extractor", all_extractors),  help="Specify extraction types from the default/local extractions .yaml file", metavar="EXTRACTION1,EXTRACTION2")
     parser.add_argument("--use_identity", "--use-identity", help="Specify an identity file id (e.g., {\"type\":\"identity\",\"name\":\"demo\",\"identity_class\":\"system\"})", metavar="[stix2 identity json]", type=parse_stix)
     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='+')
     parser.add_argument('--ignore_image_refs', default=True, type=parse_bool)
     parser.add_argument('--ignore_link_refs', default=True, type=parse_bool)
     parser.add_argument("--ignore_extraction_boundary", default=False, type=parse_bool, help="default if not passed is `false`, but if set to `true` will ignore boundary capture logic for extractions")
-    parser.add_argument('--ai_create_attack_flow', default=False, action='store_true', help="create attack flow for attack objects in report/bundle")
+    aflow_arg = parser.add_argument('--ai_create_attack_flow', default=False, action='store_true', help="create attack flow for attack objects in report/bundle")
 
     args = parser.parse_args()
     if not args.input_file.exists():
@@ -161,13 +160,13 @@ def parse_args():
         raise argparse.ArgumentError(name_arg, "max 124 characters")
 
     if args.relationship_mode == 'ai' and not args.ai_settings_relationships:
-        parser.error("relationship_mode is set to AI, --ai_settings_relationships is required")
+        raise argparse.ArgumentError(rmode_arg, "relationship_mode is set to AI, --ai_settings_relationships is required")
 
     if args.ai_create_attack_flow and not args.ai_settings_relationships:
-        parser.error("--ai_create_attack_flow requires --ai_settings_relationships")
+        raise argparse.ArgumentError(aflow_arg, "--ai_create_attack_flow requires --ai_settings_relationships")
     #### process --use-extractions 
     if args.use_extractions.get('ai') and not args.ai_settings_extractions:
-        parser.error("ai based extractors are passed, --ai_settings_extractions is required")
+        raise argparse.ArgumentError(extractions_arg, "ai based extractors are passed, --ai_settings_extractions is required")
 
     args.all_extractors  = all_extractors
     return args
@@ -218,7 +217,6 @@ def extract_all(bundler: txt2stixBundler, extractors_map, text_content, ai_extra
             logging.info("running extractor: %s", extractor.extractor_name)
             try:
                 ai_extracts = extractor.extract_objects(text_content, extractors_map["ai"].values())
-                ai_extracts = ai_extracts.model_dump().get('extractions', [])
                 bundler.process_observables(ai_extracts)
                 all_extracts[f"ai-{extractor.extractor_name}"] = ai_extracts
             except BaseException as e:
@@ -231,8 +229,7 @@ def extract_relationships_with_ai(bundler: txt2stixBundler, text_content, all_ex
     relationships = None
     try:
         all_extracts = list(itertools.chain(*all_extracts.values()))
-        relationship_types = (INCLUDES_PATH/"helpers/stix_relationship_types.txt").read_text().splitlines()
-        relationships = ai_extractor_session.extract_relationships(text_content, all_extracts, relationship_types)
+        relationships = ai_extractor_session.extract_relationships(text_content, all_extracts, RELATIONSHIP_TYPES)
         relationships = relationships.model_dump()
         log_notes(relationships, "Relationships")
         bundler.process_relationships(relationships['relationships'])
@@ -277,6 +274,7 @@ def run_txt2stix(bundler: txt2stixBundler, preprocessed_text: str, extractors_ma
         logging.info(retval.content_check.model_dump_json())
         for classification in retval.content_check.incident_classification:
             bundler.report.labels.append(f'txt2stix:{classification}'.lower())
+        bundler.add_summary(retval.content_check.summary, model.extractor_name)
 
     if should_extract or always_extract:
         if extractors_map.get("ai"):
@@ -313,7 +311,6 @@ def main():
 
         bundler = txt2stixBundler(args.name, args.use_identity, args.tlp_level, input_text, args.confidence, args.all_extractors, args.labels, created=args.created, report_id=args.report_id, external_references=args.external_refs)
         log_notes(sys.argv, "Config")
-        convo_str = None
 
         data = run_txt2stix(
             bundler, preprocessed_text, args.use_extractions,

txt2stix/utils.py

@@ -1,6 +1,5 @@
 import os
 import pkgutil
-import re
 from pathlib import Path
 from typing import Dict
 from pydantic import BaseModel, Field
@@ -45,7 +44,6 @@ class ImageLinkRemover(MarkdownRenderer):
                 del img['src']
         return soup.decode()
 
-import tldextract
 
 
 class Txt2StixData(BaseModel):
@@ -81,6 +79,9 @@ def validate_file_mimetype(file_name):
     return FILE_EXTENSIONS.get(ext)
 
 
+
+
 TLDs = [tld.lower() for tld in read_included_file('helpers/tlds.txt').splitlines()]
 REGISTRY_PREFIXES = [key.lower() for key in read_included_file('helpers/windows_registry_key_prefix.txt').splitlines()]
-FILE_EXTENSIONS = dict(line.lower().split(',') for line in read_included_file('helpers/mimetype_filename_extension_list.csv').splitlines())
+FILE_EXTENSIONS = dict(line.lower().split(',') for line in read_included_file('helpers/mimetype_filename_extension_list.csv').splitlines())
+RELATIONSHIP_TYPES = [x for x in read_included_file('helpers/stix_relationship_types.txt').splitlines() if x and not x.startswith('#')]

{txt2stix-0.0.4.dist-info → txt2stix-1.0.1.dist-info}/METADATA

@@ -1,10 +1,12 @@
 Metadata-Version: 2.4
 Name: txt2stix
-Version: 0.0.4
+Version: 1.0.1
 Summary: txt2stix is a Python script that is designed to identify and extract IoCs and TTPs from text files, identify the relationships between them, convert them to STIX 2.1 objects, and output as a STIX 2.1 bundle.
 Project-URL: Homepage, https://github.com/muchdogesec/txt2stix
 Project-URL: Issues, https://github.com/muchdogesec/txt2stix/issues
-Author-email: DOGESEC <support@dogesec.com>
+Project-URL: dogesec HQ, https://dogesec.com
+Author: dogesec
+Maintainer: dogesec
 License-File: LICENSE
 Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Operating System :: OS Independent
@@ -16,7 +18,6 @@ Requires-Dist: llama-index-core>=0.12.42
 Requires-Dist: llama-index-llms-anthropic>=0.7.2
 Requires-Dist: llama-index-llms-deepseek>=0.1.2
 Requires-Dist: llama-index-llms-gemini>=0.5.0
-Requires-Dist: llama-index-llms-openai-like>=0.4.0
 Requires-Dist: llama-index-llms-openai>=0.4.5
 Requires-Dist: llama-index-llms-openrouter>=0.3.2
 Requires-Dist: mistune>=3.0.2

{txt2stix-0.0.4.dist-info → txt2stix-1.0.1.dist-info}/RECORD

@@ -1,23 +1,23 @@
 txt2stix/__init__.py,sha256=Sm_VT913IFuAZ6dJEdVz3baPwC5VYtHySVfBAOUG92w,803
 txt2stix/attack_flow.py,sha256=WWlukuQYrGW1SJ1DnhfROYC5Ck4WYqNifgmtiuyDg7E,4177
-txt2stix/bundler.py,sha256=FRcvuWbAn8lycGRMdR3abnurTmlA6K2TYwr1XlK8SUQ,15464
+txt2stix/bundler.py,sha256=EVTcVgZyVMwb6XjNQ3Gyj7zm44UErXo9wbVr2JGsjQQ,16797
 txt2stix/common.py,sha256=ISnGNKqJPE1EcfhL-x_4G18mcwt1urmorkW-ru9kV-0,585
 txt2stix/extractions.py,sha256=ExynKWSeuWOC0q6i4SuU1NkeNw7uoOm6xu0YtJRVaiE,2058
 txt2stix/indicator.py,sha256=c6S0xx0K8JM-PT_Qd1PlN_ZlDXdnEwiRS8529iUp3yg,30774
 txt2stix/lookups.py,sha256=h42YVtYUkWZm6ZPv2h5hHDHDzDs3yBqrT_T7pj2MDZI,2301
 txt2stix/retriever.py,sha256=zU8L00RSh9N5J0NpAo3CM3IHsuZsNVjJGohRisXcMRs,5167
 txt2stix/stix.py,sha256=9nXD9a2dCY4uaatl-mlIA1k3srwQBhGW-tUSho3iYe0,30
-txt2stix/txt2stix.py,sha256=YoQ4oQcqHJOy7V94L0huM0YjoXvDQMI_4jDdlNGBRhU,16298
-txt2stix/utils.py,sha256=vCkuyTbeaCj5GH2rwm37XRuUulFoCTcmgQHlkEMiuXI,3146
+txt2stix/txt2stix.py,sha256=Vt9CUsSEO1bw5SS7vlsVxktFz1nW8M_G4-RN6idOTA0,16444
+txt2stix/utils.py,sha256=P66yq-SphsQu2S9At6BfYpavfghXsZqh4h6W13HUEoI,3256
 txt2stix/ai_extractor/__init__.py,sha256=RcXh30ZcIA3Fva2bOPH4EtWq6ffWhGE39C_II8ElAx0,417
 txt2stix/ai_extractor/anthropic.py,sha256=mdz-8CB-BSCEqnK5l35DRZURVPUf508ef2b48XMxmuk,441
-txt2stix/ai_extractor/base.py,sha256=BPqp4WCxIficLhe_Oz8iZ0tDCm_XF4WziQgzQpSDoZU,3419
+txt2stix/ai_extractor/base.py,sha256=MAtnKvWUmWZgnzwDM0i2n-WrRWq69du4KVcapNMIsEg,3523
 txt2stix/ai_extractor/deepseek.py,sha256=2XehIYbWXG6Odq68nQX4CNtl5GdmBlAmjLP_lG2eEFo,660
 txt2stix/ai_extractor/gemini.py,sha256=yJC7knYzl-TScyCBd-MTpUf-NT6znC25E7vXxNMqjLU,578
 txt2stix/ai_extractor/openai.py,sha256=DtllzeVhZw1231hj35vn1U8V2MMzm8wM7mqKLBkxazQ,489
 txt2stix/ai_extractor/openrouter.py,sha256=hAA6mTOMcpA28XYsOCvuJH7WMJqXCxfqZGJf_VrDsIk,628
-txt2stix/ai_extractor/prompts.py,sha256=85xiHmuoQsfqbeUvKEd7DEZmJkQLW_5lUvzvrXrko4Y,7707
-txt2stix/ai_extractor/utils.py,sha256=lMJGBXZ5hMqHKo1oaPIvSueSiWOy7dZCRbZJIw2mxdc,4326
+txt2stix/ai_extractor/prompts.py,sha256=3PewwmNptHEvsG1r1Yk5rs3oaEX5Jf3BRFyGaHru6r0,8137
+txt2stix/ai_extractor/utils.py,sha256=mnGIDiDa8ecwyRqDuYcKBIOiXfeQsivKxe93CfGW660,4440
 txt2stix/pattern/__init__.py,sha256=K9ofaP2AOikvzb48VSBpJZijckdqufZxSzr_kbRypLY,491
 txt2stix/pattern/extractors/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 txt2stix/pattern/extractors/base_extractor.py,sha256=ly80rp-L40g7DbhrGiCvhPWI95-ZFMtAQUEC-fH6Y-o,6130
@@ -112,8 +112,8 @@ txt2stix/includes/lookups/threat_actor.txt,sha256=QfDO9maQuqKBgW_Sdd7VGv1SHZ9Ra-
 txt2stix/includes/lookups/tld.txt,sha256=-MEgJea2NMG_KDsnc4BVvI8eRk5Dm93L-t8SGYx5wMo,8598
 txt2stix/includes/lookups/tool.txt,sha256=HGKG6JpUE26w6ezzSxOjBkp15UpSaB7N-mZ_NU_3G7A,6
 txt2stix/includes/tests/test_cases.yaml,sha256=QD1FdIunpPkOpsn6wJRqs2vil_hv8OSVaqUp4a96aZg,22247
-txt2stix-0.0.4.dist-info/METADATA,sha256=kqq-r2j9P23XerNF3P098WJ_z09D5PVzcouOelbnWcQ,12930
-txt2stix-0.0.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-txt2stix-0.0.4.dist-info/entry_points.txt,sha256=x6QPtt65hWeomw4IpJ_wQUesBl1M4WOLODbhOKyWMFg,55
-txt2stix-0.0.4.dist-info/licenses/LICENSE,sha256=BK8Ppqlc4pdgnNzIxnxde0taoQ1BgicdyqmBvMiNYgY,11364
-txt2stix-0.0.4.dist-info/RECORD,,
+txt2stix-1.0.1.dist-info/METADATA,sha256=rg6DP9idqFjH6eK2FbVmHxCgk-XnQ9mNTbFg3re-dvE,12916
+txt2stix-1.0.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+txt2stix-1.0.1.dist-info/entry_points.txt,sha256=x6QPtt65hWeomw4IpJ_wQUesBl1M4WOLODbhOKyWMFg,55
+txt2stix-1.0.1.dist-info/licenses/LICENSE,sha256=BK8Ppqlc4pdgnNzIxnxde0taoQ1BgicdyqmBvMiNYgY,11364
+txt2stix-1.0.1.dist-info/RECORD,,