txt2detection 1.0.2__tar.gz → 1.0.5__tar.gz

{txt2detection-1.0.2 → txt2detection-1.0.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: txt2detection
-Version: 1.0.2
+Version: 1.0.5
 Summary: A command line tool that takes a txt file containing threat intelligence and turns it into a detection rule.
 Project-URL: Homepage, https://github.com/muchdogesec/txt2detection
 Project-URL: Issues, https://github.com/muchdogesec/txt2detection/issues
@@ -100,6 +100,30 @@ cp .env.example .env
 
 To see more information about how to set the variables, and what they do, read the `.env.markdown` file.
 
+Then test your configoration
+
+```shell
+python3 txt2detection.py \
+  check-credentials
+```
+
+It will return a response to show what API keys are working
+
+```txt
+============= Service Statuses ===============
+  ctibutler   : authorized      ✔
+  vulmatch    : authorized      ✔
+
+  LLMS:
+    openai      : authorized      ✔
+    deepseek    : unsupported     –
+    gemini      : unsupported     –
+    openrouter  : unsupported     –
+    anthropic   : unsupported     –
+```
+
+Not all services need to be configured, if you have no intention of using them.
+
 ### Run
 
 ```shell

{txt2detection-1.0.2 → txt2detection-1.0.5}/README.md

@@ -67,6 +67,30 @@ cp .env.example .env
 
 To see more information about how to set the variables, and what they do, read the `.env.markdown` file.
 
+Then test your configoration
+
+```shell
+python3 txt2detection.py \
+  check-credentials
+```
+
+It will return a response to show what API keys are working
+
+```txt
+============= Service Statuses ===============
+  ctibutler   : authorized      ✔
+  vulmatch    : authorized      ✔
+
+  LLMS:
+    openai      : authorized      ✔
+    deepseek    : unsupported     –
+    gemini      : unsupported     –
+    openrouter  : unsupported     –
+    anthropic   : unsupported     –
+```
+
+Not all services need to be configured, if you have no intention of using them.
+
 ### Run
 
 ```shell

{txt2detection-1.0.2 → txt2detection-1.0.5}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "txt2detection"
-version = "1.0.2"
+version = "1.0.5"
 authors = [
   { name = "dogesec" }
 ]

{txt2detection-1.0.2 → txt2detection-1.0.5}/tests/src/test_main.py

@@ -135,3 +135,9 @@ def test_parse_args_sigma_mode_no_ai_provider(monkeypatch):
     assert getattr(args, "ai_provider", None) is None
     assert hasattr(args, "report_id")
     assert hasattr(args, "sigma_file")
+
+
+def test_parse_args_check_credentials(monkeypatch):
+    monkeypatch.setattr(sys, "argv", ["prog", "check-credentials"])
+    with pytest.raises(SystemExit):
+        parse_args()

{txt2detection-1.0.2 → txt2detection-1.0.5}/tests/src/test_utils.py

@@ -23,7 +23,7 @@ def test_as_date_with_datetime_and_date():
     assert as_date(d_obj) == d_obj
 
 def test_validate_token_count():
-    provider = parse_model('openai')
+    provider = parse_model('anthropic')
     with patch.object(type(provider), 'count_tokens') as mock_count_tokens:
         mock_count_tokens.return_value = 1025
         with pytest.raises(Exception):

{txt2detection-1.0.2 → txt2detection-1.0.5}/txt2detection/__main__.py

@@ -13,6 +13,7 @@ import uuid
 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
@@ -91,6 +92,7 @@ def parse_args():
     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')
@@ -115,7 +117,10 @@ def parse_args():
     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()
-    print(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"
@@ -129,6 +134,8 @@ def parse_args():
     return args
 
 
+
+
 def run_txt2detection(name, identity, tlp_level, input_text: str, labels: list[str], report_id: str|uuid.UUID, ai_provider: BaseAIExtractor, **kwargs) -> Bundler:
     if sigma := kwargs.get('sigma_file'):
         detection = get_sigma_detections(sigma)

{txt2detection-1.0.2 → txt2detection-1.0.5}/txt2detection/ai_extractor/base.py

@@ -54,4 +54,14 @@ class BaseAIExtractor():
 
     @property
     def extractor_name(self):
-        return f"{self.provider}:{self.llm.model}"
+        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-1.0.2 → txt2detection-1.0.5}/txt2detection/bundler.py

@@ -244,7 +244,7 @@ class Bundler:
 
         headers = {}
         if api_key := os.environ.get('CTIBUTLER_API_KEY'):
-            headers['Authorization'] = "Bearer " + api_key
+            headers['API-KEY'] = api_key
 
         return self._get_objects(endpoint, headers)
 
@@ -255,7 +255,7 @@ class Bundler:
         endpoint = urljoin(os.environ['VULMATCH_BASE_URL'] + '/', f"v1/cve/objects/?cve_id="+','.join(cve_ids))
         headers = {}
         if api_key := os.environ.get('VULMATCH_API_KEY'):
-            headers['Authorization'] = "Bearer " + api_key
+            headers['API-KEY'] = api_key
 
         return self._get_objects(endpoint, headers)
 

txt2detection-1.0.5/txt2detection/credential_checker.py

@@ -0,0 +1,80 @@
+import argparse
+import os
+import random
+from urllib.parse import urljoin
+import requests
+
+
+
+def check_llms():
+    from txt2detection.__main__ import parse_model
+
+    auth_info = dict()
+    for model_name in ["openai", "deepseek", "gemini", "openrouter", "anthropic"]:
+        try:
+            model = parse_model(model_name)
+            auth_info[model_name] = model.check_credential()
+        except argparse.ArgumentTypeError:
+            auth_info[model_name] = "unsupported"
+        except:
+            auth_info[model_name] = "unauthorized"
+    return auth_info
+
+
+def check_ctibutler_vulmatch(service):
+    session = requests.Session()
+    if service == 'vulmatch':
+        base_url = os.getenv('VULMATCH_BASE_URL')
+        url = urljoin(base_url, 'v1/cve/objects/vulnerability--f552f6f4-39da-48dc-8717-323772c99588/')
+        session.headers['API-KEY'] = os.environ.get('VULMATCH_API_KEY')
+    elif service == 'ctibutler':
+        base_url = os.getenv('CTIBUTLER_BASE_URL')
+        url = urljoin(base_url, 'v1/location/versions/available/')
+        session.headers['API-KEY'] = os.environ.get('CTIBUTLER_API_KEY')
+
+    try:
+        resp = session.get(url)
+        match resp.status_code:
+            case 401 | 403:
+                return "unauthorized"
+            case 200:
+                return "authorized"
+            case _:
+                return "unknown"
+    except:
+        return "offline"
+
+
+def check_statuses(test_llms=False):
+    statuses = dict(
+        ctibutler=check_ctibutler_vulmatch("ctibutler"),
+        vulmatch=check_ctibutler_vulmatch("vulmatch"),
+    )
+    if test_llms:
+        statuses.update(llms=check_llms())
+    return statuses
+
+
+def format_statuses(status_dict):
+    def get_marker(status):
+        """Return a checkmark, cross, or dash based on status."""
+        match status.lower():
+            case "authorized":
+                return "✔"
+            case "unauthorized":
+                return "✖"
+            case "unknown" | "offline" | "unsupported":
+                return "–"
+            case _:
+                return "?"
+
+    print("============= Service Statuses ===============")
+    for key, value in status_dict.items():
+        if key == "llms" and isinstance(value, dict):
+            print(f"\n  {key.upper()}:")
+            for llm_name, llm_status in value.items():
+                marker = get_marker(llm_status)
+                print(f"    {llm_name:<12}: {llm_status:<15} {marker}")
+        else:
+            marker = get_marker(value)
+            print(f"  {key:<12}: {value:<15} {marker}")

{txt2detection-1.0.2 → txt2detection-1.0.5}/txt2detection/models.py

@@ -3,9 +3,10 @@ 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, List, Literal, Optional, Union
+from typing import Any, ClassVar, List, Literal, Optional, Union
 from uuid import UUID
 
 import jsonschema
@@ -194,6 +195,7 @@ class BaseDetection(BaseModel):
     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 []
@@ -249,9 +251,7 @@ class BaseDetection(BaseModel):
     def validate_rule_with_json_schema(self, rule):
         jsonschema.validate(
             rule,
-            {
-                "$ref": "https://github.com/SigmaHQ/sigma-specification/raw/refs/heads/main/json-schema/sigma-detection-rule-schema.json"
-            },
+            self.sigma_json_schema,
         )
     
     @property

{txt2detection-1.0.2 → txt2detection-1.0.5}/txt2detection/utils.py

@@ -55,11 +55,13 @@ def validate_token_count(max_tokens, input, extractor: BaseAIExtractor):
     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():
-    @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']}
     return get_licenses(datetime.now().date().isoformat())