sl5_control_catalog_extractor.py

#!/usr/bin/env python3
"""
SL5 Control Export to NIST Catalog JSON Converter

This script extracts control information from the sl5_control_export PDF and converts it
to a JSON format similar to the NIST SP 800-53 control catalog.

Key features:
- Extracts base controls and enhancements
- Converts to NIST-style JSON format with proper metadata
- Handles control text, discussions, and related controls
- Provides comprehensive summary statistics
- Filters out controls with no or empty text attributes

Usage:
    python sl5_control_catalog_extractor.py <pdf_file>
    python sl5_control_catalog_extractor.py <pdf_file> --debug-page N
"""

import fitz  # PyMuPDF
import json
import re
import sys
from typing import Dict, List, Optional

class SL5ControlCatalogExtractor:
    def __init__(self, pdf_path: str):
        self.pdf_path = pdf_path
        self.controls = []  # List of control objects in NIST format
        self.current_control = None
        self.current_attribute = None
        self.debug_mode = False
        
        # Known attributes for SL5 control export
        self.known_attributes = [
            "Control Text",
            "Assessment Procedure", 
            "Discussion",
            "Open Questions",
            "Related Controls"
        ]
        
        # Family name mapping (expand as needed)
        self.family_names = {
            "AC": "ACCESS CONTROL",
            "AT": "AWARENESS AND TRAINING",
            "AU": "AUDIT AND ACCOUNTABILITY",
            "CA": "ASSESSMENT, AUTHORIZATION, AND MONITORING",
            "CM": "CONFIGURATION MANAGEMENT",
            "CP": "CONTINGENCY PLANNING",
            "IA": "IDENTIFICATION AND AUTHENTICATION",
            "IS": "INTERFACE SECURITY",
            "IR": "INCIDENT RESPONSE",
            "MA": "MAINTENANCE",
            "MP": "MEDIA PROTECTION",
            "NS": "NETWORK SEPARATION",
            "PE": "PHYSICAL AND ENVIRONMENTAL PROTECTION",
            "PL": "PLANNING",
            "PM": "PROGRAM MANAGEMENT",
            "PS": "PERSONNEL SECURITY",
            "PT": "PERSONALLY IDENTIFIABLE INFORMATION PROCESSING AND TRANSPARENCY",
            "RA": "RISK ASSESSMENT",
            "SA": "SYSTEM AND SERVICES ACQUISITION",
            "SC": "SYSTEM AND COMMUNICATIONS PROTECTION",
            "SI": "SYSTEM AND INFORMATION INTEGRITY",
            "SR": "SUPPLY CHAIN RISK MANAGEMENT",
            "ZT": "ZERO TRUST"
        }
        
    def extract_controls(self) -> List[Dict]:
        """Extract all controls from the PDF document."""
        try:
            doc = fitz.open(self.pdf_path)
            print(f"Processing {len(doc)} pages...")
            
            for page_num in range(len(doc)):
                page = doc[page_num]
                self._process_page(page, page_num + 1)
                
            doc.close()
            print(f"Extracted {len(self.controls)} controls")
            return self.controls
            
        except Exception as e:
            print(f"Error processing PDF: {e}")
            return []
    
    def _process_page(self, page, page_num: int):
        """Process a single page extracting control information."""
        try:
            # Get formatted text with style information
            formatted_text = self._get_formatted_text(page)
            self._extract_controls_from_text(formatted_text, page_num)
        except Exception as e:
            if self.debug_mode:
                print(f"Error processing page {page_num}: {e}")
    
    def _get_formatted_text(self, page) -> List[Dict]:
        """Extract text with formatting information from the page."""
        text_dict = page.get_text("dict")
        formatted_lines = []
        
        for block in text_dict["blocks"]:
            if "lines" in block:
                for line in block["lines"]:
                    line_text = ""
                    formats = []
                    
                    for span in line["spans"]:
                        span_text = span["text"]
                        line_text += span_text
                        
                        # Track format information
                        formats.append({
                            "text": span_text,
                            "bold": "Bold" in span["font"] or span["flags"] & 2**4,
                            "italic": "Italic" in span["font"] or span["flags"] & 2**1,
                            "size": span["size"]
                        })
                    
                    if line_text.strip():
                        formatted_lines.append({
                            "text": line_text.strip(),
                            "formats": formats
                        })
        
        return formatted_lines
    
    def _extract_controls_from_text(self, formatted_text: List[Dict], page_num: int):
        """Extract control information from formatted text."""
        processing_enhancements = False
        
        for line_data in formatted_text:
            line_text = line_data["text"]
            formats = line_data["formats"]
            
            # Skip common header/footer elements
            if self._is_header_footer(line_text):
                continue
            
            # Check if any part of the line is bold (indicates potential control ID/name)
            is_bold = any(fmt.get("bold", False) for fmt in formats)
            
            # Look for control IDs - they should be bold and match the pattern
            if is_bold:
                # Check for main control pattern
                control_match = re.match(r'^([A-Z]{2}-\d{1,2}),\s*(.+)', line_text)
                if control_match:
                    control_id = control_match.group(1)
                    control_name = control_match.group(2).strip()
                    
                    # Create new control in NIST format
                    control_obj = self._create_control_object(control_id, control_name, page_num)
                    self.controls.append(control_obj)
                    self.current_control = control_obj
                    self.current_attribute = None
                    processing_enhancements = False
                    
                    if self.debug_mode:
                        print(f"Found control: {control_id} - {control_name[:50]}...")
                    continue
                
                # Check for enhancement pattern with bullet points
                if (line_text.strip().startswith('●') or line_text.strip().startswith('•')):
                    enhancement_match = re.search(r'([A-Z]{2}-\d{1,2}\(\d+\))\s+(.+)', line_text)
                    if enhancement_match:
                        enhancement_id = enhancement_match.group(1)
                        enhancement_name = enhancement_match.group(2).strip()
                        
                        # Create new enhancement in NIST format
                        enhancement_obj = self._create_enhancement_object(enhancement_id, enhancement_name, page_num)
                        self.controls.append(enhancement_obj)
                        self.current_control = enhancement_obj
                        self.current_attribute = None
                        
                        if self.debug_mode:
                            print(f"Found enhancement: {enhancement_id} - {enhancement_name[:50]}...")
                        continue
            
            # Check if we're entering a Control Enhancements section
            if self.current_control and line_text.strip() == "Control Enhancements:":
                processing_enhancements = True
                if self.debug_mode:
                    print(f"  Entering Control Enhancements section")
                continue
            
            # Look for attributes
            if self.current_control and ':' in line_text:
                attribute_result = self._extract_attribute_from_line(line_text)
                
                if attribute_result:
                    attribute_name, attribute_content = attribute_result
                    self._store_attribute(attribute_name, attribute_content)
                    self.current_attribute = attribute_name
                    
                    if self.debug_mode:
                        print(f"  Found attribute: {attribute_name}")
                    continue
            
            # Continue previous attribute content
            if self.current_control and self.current_attribute:
                self._append_to_current_attribute(line_text)
    
    def _create_control_object(self, control_id: str, control_name: str, page_num: int) -> Dict:
        """Create a control object in NIST format."""
        family = control_id.split('-')[0]
        
        return {
            "id": control_id,
            "name": control_name,
            "family": family,
            "isEnhancement": False,
            "baseControlId": control_id,
            "enhancementNumber": None,
            "text": "",
            "discussion": "",
            "openQuestions": "",
            "relatedControls": [],
            "metadata": {
                "source": "SL5 Control Export"
            },
            "_attributes": {}  # Temporary storage for raw attributes
        }
    
    def _create_enhancement_object(self, enhancement_id: str, enhancement_name: str, page_num: int) -> Dict:
        """Create an enhancement object in NIST format."""
        # Parse enhancement ID to get base control and enhancement number
        match = re.match(r'^([A-Z]{2}-\d{1,2})\((\d+)\)', enhancement_id)
        if match:
            base_control_id = match.group(1)
            enhancement_number = int(match.group(2))
        else:
            base_control_id = enhancement_id
            enhancement_number = 1
        
        family = base_control_id.split('-')[0]
        
        return {
            "id": enhancement_id,
            "name": enhancement_name,
            "family": family,
            "isEnhancement": True,
            "baseControlId": base_control_id,
            "enhancementNumber": enhancement_number,
            "text": "",
            "discussion": "",
            "openQuestions": "",
            "relatedControls": [],
            "metadata": {
                "source": "SL5 Control Export"
            },
            "_attributes": {}  # Temporary storage for raw attributes
        }
    
    def _extract_attribute_from_line(self, line_text: str) -> Optional[tuple]:
        """Extract attribute name and content from a line."""
        for attr in self.known_attributes:
            if line_text.startswith(attr + ":"):
                content = line_text[len(attr) + 1:].strip()
                return attr, content
        return None
    
    def _store_attribute(self, attribute_name: str, content: str):
        """Store attribute content in the current control."""
        if not self.current_control:
            return
        
        # Store in temporary attributes for processing
        self.current_control["_attributes"][attribute_name] = content
        
        # Map to NIST format fields
        if attribute_name == "Control Text":
            self.current_control["text"] = content
        elif attribute_name == "Discussion":
            self.current_control["discussion"] = content
        elif attribute_name == "Open Questions":
            self.current_control["openQuestions"] = content
        elif attribute_name == "Related Controls":
            # Parse related controls list
            if content and content.strip() != "TODO":
                # Simple parsing - assumes comma-separated list
                related = [ctrl.strip() for ctrl in content.split(',') if ctrl.strip()]
                self.current_control["relatedControls"] = related
    
    def _append_to_current_attribute(self, line_text: str):
        """Append text to the current attribute."""
        if not self.current_control or not self.current_attribute:
            return
        
        # Skip if line looks like a new control or attribute
        if (re.match(r'^[A-Z]{2}-\d{1,2}[,\s]', line_text) or
            any(line_text.startswith(attr + ":") for attr in self.known_attributes)):
            return
        
        # Append to current attribute
        current_content = self.current_control["_attributes"].get(self.current_attribute, "")
        if current_content:
            current_content += " " + line_text
        else:
            current_content = line_text
        
        self.current_control["_attributes"][self.current_attribute] = current_content
        
        # Update NIST format fields
        if self.current_attribute == "Control Text":
            self.current_control["text"] = current_content
        elif self.current_attribute == "Discussion":
            self.current_control["discussion"] = current_content
        elif self.current_attribute == "Open Questions":
            self.current_control["openQuestions"] = current_content
        elif self.current_attribute == "Related Controls":
            if current_content and current_content.strip() != "TODO":
                related = [ctrl.strip() for ctrl in current_content.split(',') if ctrl.strip()]
                self.current_control["relatedControls"] = related
    
    def _is_header_footer(self, text: str) -> bool:
        """Check if text is likely a header or footer."""
        text = text.strip().lower()
        return (
            text.startswith('page ') or
            text.endswith(' of ') or
            len(text) < 3 or
            text.isdigit() or
            'sl5' in text or
            'security level' in text
        )
    
    def _has_meaningful_text_content(self, control: Dict) -> bool:
        """Check if a control has meaningful text content in its text attribute."""
        text = control.get("text", "").strip()
        
        # Consider control as having meaningful content if:
        # 1. It has non-empty text content
        # 2. The text is not just placeholder values like "TODO" or empty strings
        if not text:
            return False
        
        # Check for common placeholder values that should be considered "empty"
        placeholder_values = ["TODO", "TBD", "N/A", "", " "]
        if text in placeholder_values:
            return False
        
        # Must have at least some meaningful content (more than just whitespace)
        if len(text.strip()) < 5:  # Arbitrary minimum length for meaningful text
            return False
        
        return True
    
    def create_nist_catalog_json(self) -> Dict:
        """Create the full NIST-style catalog JSON structure."""
        # Clean up controls - remove temporary attributes and filter out controls with no text
        cleaned_controls = []
        filtered_count = 0
        
        for control in self.controls:
            # Check if control has meaningful text content
            if self._has_meaningful_text_content(control):
                # Remove temporary attributes and add to final list
                cleaned_control = {k: v for k, v in control.items() if k != "_attributes"}
                cleaned_controls.append(cleaned_control)
            else:
                filtered_count += 1
                if self.debug_mode:
                    print(f"Filtered out control {control.get('id', 'unknown')} - no meaningful text content")
        
        # Print filtering summary
        if filtered_count > 0:
            print(f"Filtered out {filtered_count} controls with no or empty text attributes")
            print(f"Remaining controls: {len(cleaned_controls)}")
        
        # Calculate summary statistics
        summary = self._calculate_summary(cleaned_controls)
        
        catalog = {
            "metadata": {
                "title": "SL5 Security Controls",
                "version": "SL5 Export",
                "source": "SL5 Control Export PDF"
            },
            "summary": summary,
            "controls": cleaned_controls
        }
        
        return catalog
    
    def _calculate_summary(self, controls: List[Dict]) -> Dict:
        """Calculate summary statistics for the catalog."""
        total_controls = len(controls)
        base_controls = sum(1 for c in controls if not c['isEnhancement'])
        enhancements = sum(1 for c in controls if c['isEnhancement'])
        
        # Group by family
        families = {}
        for control in controls:
            family = control['family']
            if family not in families:
                families[family] = {
                    "name": self.family_names.get(family, family),
                    "total": 0,
                    "baseControls": 0,
                    "enhancements": 0
                }
            
            families[family]["total"] += 1
            if control['isEnhancement']:
                families[family]["enhancements"] += 1
            else:
                families[family]["baseControls"] += 1
        
        return {
            "totalControls": total_controls,
            "baseControls": base_controls,
            "enhancements": enhancements,
            "families": families
        }
    
    def save_to_json(self, output_file: str):
        """Save extracted controls to JSON file in NIST catalog format."""
        try:
            catalog = self.create_nist_catalog_json()
            
            with open(output_file, 'w', encoding='utf-8') as f:
                json.dump(catalog, f, indent=2, ensure_ascii=False)
            
            print(f"\nSuccessfully saved catalog with {len(catalog['controls'])} controls to {output_file}")
        except Exception as e:
            print(f"Error saving to JSON: {e}")
    
    def print_summary(self):
        """Print a summary of extracted controls."""
        print(f"\n=== SL5 CONTROL CATALOG EXTRACTION SUMMARY ===")
        print(f"Total controls extracted: {len(self.controls)}")
        
        if not self.controls:
            print("No controls found!")
            return
        
        # Count by type
        base_controls = sum(1 for c in self.controls if not c['isEnhancement'])
        enhancements = sum(1 for c in self.controls if c['isEnhancement'])
        
        print(f"Base controls: {base_controls}")
        print(f"Control enhancements: {enhancements}")
        
        # Count by family
        families = {}
        for control in self.controls:
            family = control['family']
            families[family] = families.get(family, 0) + 1
        
        print(f"\nControls by family:")
        for family, count in sorted(families.items()):
            family_name = self.family_names.get(family, family)
            print(f"  {family} ({family_name}): {count} controls")
        
        # Show sample controls
        print(f"\nSample controls:")
        for i, control in enumerate(self.controls[:5]):
            control_type = "Enhancement" if control['isEnhancement'] else "Base Control"
            print(f"  {control['id']} ({control_type}): {control['name'][:60]}...")
    
    def debug_page(self, page_num: int):
        """Debug a specific page to understand formatting."""
        try:
            doc = fitz.open(self.pdf_path)
            if page_num < 1 or page_num > len(doc):
                print(f"Invalid page number. Document has {len(doc)} pages.")
                return
            
            page = doc[page_num - 1]  # Convert to 0-based index
            print(f"\n=== DEBUG PAGE {page_num} ===")
            
            formatted_text = self._get_formatted_text(page)
            
            print(f"Found {len(formatted_text)} text lines:")
            for i, line_data in enumerate(formatted_text):
                line_text = line_data["text"]
                formats = line_data["formats"]
                is_bold = any(fmt.get("bold", False) for fmt in formats)
                
                print(f"{i+1:3d}: {'[B]' if is_bold else '   '} {line_text[:80]}")
                
                # Show detailed format info for first few lines
                if i < 5:
                    for fmt in formats:
                        if fmt["text"].strip():
                            print(f"     Format: '{fmt['text']}' - Bold: {fmt['bold']}, Size: {fmt['size']}")
            
            doc.close()
            
        except Exception as e:
            print(f"Error debugging page: {e}")

def main():
    if len(sys.argv) < 2:
        print("Usage: python sl5_control_catalog_extractor.py <pdf_file> [--debug-page N]")
        print("Example: python sl5_control_catalog_extractor.py sl5_control_export.pdf")
        print("Example: python sl5_control_catalog_extractor.py sl5_control_export.pdf --debug-page 10")
        sys.exit(1)
    
    pdf_file = sys.argv[1]
    
    # Check for debug mode
    if len(sys.argv) == 4 and sys.argv[2] == "--debug-page":
        try:
            debug_page_num = int(sys.argv[3])
            extractor = SL5ControlCatalogExtractor(pdf_file)
            extractor.debug_mode = True
            extractor.debug_page(debug_page_num)
            return
        except ValueError:
            print("Invalid page number for debug mode")
            sys.exit(1)
    
    output_file = "sl5_control_catalog.json"
    
    print("SL5 Control Catalog Extractor")
    print("=" * 40)
    
    extractor = SL5ControlCatalogExtractor(pdf_file)
    extractor.debug_mode = False
    controls = extractor.extract_controls()
    
    if controls:
        extractor.save_to_json(output_file)
        extractor.print_summary()
    else:
        print("No controls were extracted.")
        print("Try using --debug-page N to see the formatting details.")

if __name__ == "__main__":
    main()