ntermqt 0.1.4__py3-none-any.whl → 0.1.5__py3-none-any.whl

nterm/parser/tfsm_fire.py

@@ -0,0 +1,237 @@
+import sqlite3
+import textfsm
+from typing import Dict, List, Tuple, Optional
+import io
+import time
+import click
+from multiprocessing import Process, Queue
+import multiprocessing
+import sys
+import threading
+from contextlib import contextmanager
+
+
+class ThreadSafeConnection:
+    """Thread-local storage for SQLite connections"""
+
+    def __init__(self, db_path: str, verbose: bool = False):
+        self.db_path = db_path
+        self.verbose = verbose
+        self._local = threading.local()
+
+    @contextmanager
+    def get_connection(self):
+        """Get a thread-local connection"""
+        if not hasattr(self._local, 'connection'):
+            self._local.connection = sqlite3.connect(self.db_path)
+            self._local.connection.row_factory = sqlite3.Row
+            if self.verbose:
+                click.echo(f"Created new connection in thread {threading.get_ident()}")
+
+        try:
+            yield self._local.connection
+        except Exception as e:
+            if hasattr(self._local, 'connection'):
+                self._local.connection.close()
+                delattr(self._local, 'connection')
+            raise e
+
+    def close_all(self):
+        """Close connection if it exists for current thread"""
+        if hasattr(self._local, 'connection'):
+            self._local.connection.close()
+            delattr(self._local, 'connection')
+
+
+class TextFSMAutoEngine:
+    def __init__(self, db_path: str, verbose: bool = False):
+        self.db_path = db_path
+        self.verbose = verbose
+        self.connection_manager = ThreadSafeConnection(db_path, verbose)
+
+    def _calculate_template_score(
+            self,
+            parsed_data: List[Dict],
+            template: sqlite3.Row,
+            raw_output: str
+    ) -> float:
+        """
+        Score template match quality (0-100 scale).
+
+        Factors:
+        - Record count (0-30 pts): Did the template find data?
+        - Field richness (0-30 pts): How many fields per record?
+        - Population rate (0-25 pts): Are fields actually filled?
+        - Consistency (0-15 pts): Uniform data across records?
+        """
+        if not parsed_data:
+            return 0.0
+
+        num_records = len(parsed_data)
+        num_fields = len(parsed_data[0].keys()) if parsed_data else 0
+        is_version_cmd = 'version' in template['cli_command'].lower()
+
+        # === Factor 1: Record Count (0-30 points) ===
+        if is_version_cmd:
+            # Version commands: expect exactly 1 record
+            record_score = 30.0 if num_records == 1 else max(0, 15 - (num_records - 1) * 5)
+        else:
+            # Diminishing returns: log scale capped at 30
+            # 1 rec = 10, 3 rec = 20, 10+ rec = 30
+            if num_records >= 10:
+                record_score = 30.0
+            elif num_records >= 3:
+                record_score = 20.0 + (num_records - 3) * (10.0 / 7.0)
+            else:
+                record_score = num_records * 10.0
+
+        # === Factor 2: Field Richness (0-30 points) ===
+        # More fields = richer data extraction
+        # 1-2 fields = weak, 3-5 = decent, 6-10 = good, 10+ = excellent
+        if num_fields >= 10:
+            field_score = 30.0
+        elif num_fields >= 6:
+            field_score = 20.0 + (num_fields - 6) * 2.5
+        elif num_fields >= 3:
+            field_score = 10.0 + (num_fields - 3) * (10.0 / 3.0)
+        else:
+            field_score = num_fields * 5.0
+
+        # === Factor 3: Population Rate (0-25 points) ===
+        # What percentage of cells have actual data?
+        total_cells = num_records * num_fields
+        populated_cells = 0
+
+        for record in parsed_data:
+            for value in record.values():
+                if value is not None and str(value).strip():
+                    populated_cells += 1
+
+        population_rate = populated_cells / total_cells if total_cells > 0 else 0
+        population_score = population_rate * 25.0
+
+        # === Factor 4: Consistency (0-15 points) ===
+        # Are the same fields populated across all records?
+        if num_records > 1:
+            # Check which fields are populated in each record
+            field_fill_counts = {key: 0 for key in parsed_data[0].keys()}
+
+            for record in parsed_data:
+                for key, value in record.items():
+                    if value is not None and str(value).strip():
+                        field_fill_counts[key] += 1
+
+            # Consistency = fields that are either always filled or never filled
+            consistent_fields = sum(
+                1 for count in field_fill_counts.values()
+                if count == 0 or count == num_records
+            )
+            consistency_rate = consistent_fields / num_fields if num_fields > 0 else 0
+            consistency_score = consistency_rate * 15.0
+        else:
+            # Single record = perfectly consistent
+            consistency_score = 15.0
+
+        total_score = record_score + field_score + population_score + consistency_score
+
+        if self.verbose:
+            click.echo(f"    Scoring: records={record_score:.1f}, fields={field_score:.1f}, "
+                       f"population={population_score:.1f}, consistency={consistency_score:.1f} "
+                       f"-> {total_score:.1f}")
+
+        return total_score
+
+    def find_best_template(self, device_output: str, filter_string: Optional[str] = None) -> Tuple[
+        Optional[str], Optional[List[Dict]], float, List[Tuple[str, float, int]]]:
+        """Try filtered templates against the output and return the best match plus all non-zero scores."""
+        best_template = None
+        best_parsed_output = None
+        best_score = 0
+        all_scores = []  # List of (template_name, score, record_count)
+
+        with self.connection_manager.get_connection() as conn:
+            templates = self.get_filtered_templates(conn, filter_string)
+            total_templates = len(templates)
+
+            if self.verbose:
+                click.echo(f"Found {total_templates} matching templates for filter: {filter_string}")
+
+            for idx, template in enumerate(templates, 1):
+                if self.verbose:
+                    percentage = (idx / total_templates) * 100
+                    click.echo(f"\nTemplate {idx}/{total_templates} ({percentage:.1f}%): {template['cli_command']}")
+
+                try:
+                    textfsm_template = textfsm.TextFSM(io.StringIO(template['textfsm_content']))
+                    parsed = textfsm_template.ParseText(device_output)
+                    parsed_dicts = [dict(zip(textfsm_template.header, row)) for row in parsed]
+                    score = self._calculate_template_score(parsed_dicts, template, device_output)
+
+                    if self.verbose:
+                        click.echo(f" -> Score={score:.2f}, Records={len(parsed_dicts)}")
+
+                    # Track all non-zero scores
+                    if score > 0:
+                        all_scores.append((template['cli_command'], score, len(parsed_dicts)))
+
+                    if score > best_score:
+                        best_score = score
+                        best_template = template['cli_command']
+                        best_parsed_output = parsed_dicts
+                        if self.verbose:
+                            click.echo(click.style("  New best match!", fg='green'))
+
+                except Exception as e:
+                    if self.verbose:
+                        click.echo(f" -> Failed to parse: {str(e)}")
+                    continue
+
+        # Sort all_scores by score descending
+        all_scores.sort(key=lambda x: x[1], reverse=True)
+
+        return best_template, best_parsed_output, best_score, all_scores
+
+    def get_filtered_templates(self, connection: sqlite3.Connection, filter_string: Optional[str] = None):
+        """Get filtered templates from database using provided connection."""
+        cursor = connection.cursor()
+        if filter_string:
+            filter_terms = filter_string.replace('-', '_').split('_')
+            query = "SELECT * FROM templates WHERE 1=1"
+            params = []
+            for term in filter_terms:
+                if term and len(term) > 2:
+                    query += " AND cli_command LIKE ?"
+                    params.append(f"%{term}%")
+            cursor.execute(query, params)
+        else:
+            cursor.execute("SELECT * FROM templates")
+        return cursor.fetchall()
+
+    def __del__(self):
+        """Clean up connections on deletion"""
+        self.connection_manager.close_all()
+
+
+# Example usage
+if __name__ == '__main__':
+    multiprocessing.freeze_support()
+
+
+    # Example of using the engine in multiple threads
+    def worker(engine, output, filter_str):
+        result = engine.find_best_template(output, filter_str)
+        print(f"Thread {threading.get_ident()}: Found template: {result[0]}")
+
+
+    engine = TextFSMAutoEngine("./secure_cartography/tfsm_templates.db", verbose=True)
+    threads = []
+
+    # Create multiple threads
+    for i in range(3):
+        t = threading.Thread(target=worker, args=(engine, "sample output", "show version"))
+        threads.append(t)
+        t.start()
+
+    # Wait for all threads to complete
+    for t in threads:
+        t.join()