txt2detection 1.0.2__tar.gz → 1.0.3__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: txt2detection
-Version: 1.0.2
+Version: 1.0.3
 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.3}/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.3}/pyproject.toml

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

{txt2detection-1.0.2 → txt2detection-1.0.3}/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.3}/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.3}/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.3}/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.3/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}")