sl5_overlay_extractor.py

#!/usr/bin/env python3
"""
SL5 Overlay Control Extractor

This script extracts control information from the sl5_overlay PDF.
Key differences from insider threat overlay:
- Attributes are plain text (not underlined)
- Controls are "selected" if they have "Justification to Select" rather than just "Open Questions"
- Different attribute patterns
- Enhanced to properly extract Open Questions for control enhancements
"""

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

class SL5ControlExtractor:
    def __init__(self, pdf_path: str):
        self.pdf_path = pdf_path
        self.controls = {}
        self.current_control = None
        self.current_attribute = None  
        self.debug_mode = False
        
        # Known attributes for SL5 overlay (case-sensitive)
        self.known_attributes = [
            "Justification to Select",
            "Open Questions", 
            "Inherit From",
            "Parameter Value(s)",
            "Test Method",
            "Control Enhancements",
            "Supplemental Guidance",
            "Regulatory/Statutory Reference(s)"
        ]
        
    def extract_controls(self) -> 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()
            
            # Post-process to determine selected status
            self._determine_selected_status()
            
            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."""
        try:
            text_dict = page.get_text("dict")
            formatted_text = self._extract_formatted_text(text_dict)
            self._find_controls_and_attributes(formatted_text, page_num)
        except Exception as e:
            if self.debug_mode:
                print(f"Error processing page {page_num}: {e}")
    
    def _extract_formatted_text(self, text_dict: dict) -> List[dict]:
        """Extract text with formatting information."""
        formatted_text = []
        
        for block in text_dict.get("blocks", []):
            if block.get("type") == 0:  # Text block
                for line in block.get("lines", []):
                    line_text = ""
                    line_formats = []
                    
                    for span in line.get("spans", []):
                        text = span.get("text", "")
                        flags = span.get("flags", 0)
                        
                        # Extract formatting
                        bold = bool(flags & 16)
                        italic = bool(flags & 2)
                        underline = bool(flags & 4)
                        
                        line_text += text
                        line_formats.append({
                            "text": text,
                            "bold": bold,
                            "italic": italic,
                            "underline": underline
                        })
                    
                    if line_text.strip():
                        formatted_text.append({
                            "text": line_text,
                            "formats": line_formats
                        })
        
        return formatted_text
    
    def _find_controls_and_attributes(self, formatted_text: List[dict], page_num: int):
        """Find controls and their attributes in the formatted text."""
        processing_enhancements = False
        parent_control = None
        
        for text_info in formatted_text:
            line_text = text_info["text"]
            formats = text_info["formats"]
            
            # Look for control IDs (typically bold)
            is_bold = any(fmt.get("bold", False) for fmt in formats)
            if is_bold:
                control_match = re.match(r'^([A-Z]{2}-\d{1,2}(?:\(\d+\))?)', line_text)
                if control_match:
                    control_id = control_match.group(1)
                    # Extract control name (everything after the ID and comma/space)
                    remaining_text = line_text[len(control_id):].strip()
                    if remaining_text.startswith(','):
                        remaining_text = remaining_text[1:].strip()
                    
                    control_name = remaining_text
                    
                    # Initialize new control
                    self.controls[control_id] = {
                        "name": control_name,
                        "attributes": {},
                        "selected": False  # Will be determined later
                    }
                    self.current_control = control_id
                    self.current_attribute = None
                    processing_enhancements = False
                    parent_control = None
                    
                    print(f"Found control: {control_id} - {control_name[:50]}...")
                    continue
            
            # Check if we're entering a Control Enhancements section
            if self.current_control and line_text.strip() == "Control Enhancements:":
                # Don't create an empty attribute, just set the processing flag
                processing_enhancements = True
                parent_control = self.current_control  # Remember the parent control
                if self.debug_mode:
                    print(f"  Entering Control Enhancements section for {self.current_control}")
                continue
            
            # Look for control enhancements (indicated by bullet points and enhancement format)
            if processing_enhancements and (line_text.strip().startswith('●') or line_text.strip().startswith('•')):
                # Extract enhancement pattern
                enhancement_match = re.search(r'([A-Z]{2}-\d{1,2}\(\d+\))\s+(.+?)\s*\|\s*(.+)', line_text)
                if enhancement_match:
                    enhancement_id = enhancement_match.group(1)
                    base_control_name = enhancement_match.group(2).strip()
                    enhancement_name = enhancement_match.group(3).strip()
                    
                    # Create full enhancement name combining base control and enhancement
                    full_enhancement_name = f"{base_control_name} | {enhancement_name}"
                    
                    # Initialize new enhancement as a separate control
                    self.controls[enhancement_id] = {
                        "name": full_enhancement_name,
                        "attributes": {},
                        "selected": False  # Will be determined later
                    }
                    self.current_control = enhancement_id
                    self.current_attribute = None
                    
                    print(f"Found enhancement: {enhancement_id} - {full_enhancement_name[:50]}...")
                    continue
            
            # Look for attributes (not necessarily bold in SL5 overlay)
            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
                    
                    # Store the attribute content - NO LONGER redirect Open Questions from enhancements
                    self.controls[self.current_control]["attributes"][attribute_name] = attribute_content.strip()
                    self.current_attribute = attribute_name
                    
                    if self.debug_mode:
                        print(f"  Added attribute to {self.current_control}: {attribute_name}")
                    
                    # Special handling: if we encounter Open Questions that are not part of an enhancement
                    # and we were processing enhancements, we may need to stop processing enhancements
                    # This handles cases where Open Questions appear after all enhancements
                    if (attribute_name == "Open Questions" and 
                        processing_enhancements and 
                        parent_control and 
                        not '(' in self.current_control):
                        # This is an Open Questions for the parent control, not an enhancement
                        processing_enhancements = False
                    
                    continue
            
            # If we have a current attribute and this line doesn't start a new attribute or control,
            # append this line to the current attribute's content
            elif (self.current_control and self.current_attribute and 
                  line_text.strip() and 
                  not self._is_new_control_line(line_text, formats) and
                  not self._starts_with_known_attribute(line_text) and
                  not self._is_enhancement_line(line_text)):
                
                # This is continuation content for the current attribute
                current_content = self.controls[self.current_control]["attributes"][self.current_attribute]
                self.controls[self.current_control]["attributes"][self.current_attribute] = current_content + " " + line_text.strip()
                
                if self.debug_mode:
                    print(f"    Appended to {self.current_attribute}: {line_text.strip()[:50]}...")
    
    def _extract_attribute_from_line(self, line_text: str):
        """Extract attribute name and content if line starts with known attribute pattern."""
        for attr in self.known_attributes:
            if line_text.startswith(attr + ":"):
                # Extract content after the colon
                content = line_text[len(attr + ":"):].strip()
                return attr, content
        
        return None
    
    def _is_new_control_line(self, line_text: str, formats: List[dict]) -> bool:
        """Check if this line starts a new control."""
        is_bold = any(fmt.get("bold", False) for fmt in formats)
        if is_bold and re.match(r'^([A-Z]{2}-\d{1,2}(?:\(\d+\))?)', line_text):
            return True
        return False
    
    def _starts_with_known_attribute(self, line_text: str) -> bool:
        """Check if line starts with a known attribute."""
        for attr in self.known_attributes:
            if line_text.startswith(attr + ":"):
                return True
        return False
    
    def _is_enhancement_line(self, line_text: str) -> bool:
        """Check if this line is a control enhancement line."""
        # 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+\))', line_text)
            return enhancement_match is not None
        return False
    
    def _determine_selected_status(self):
        """Determine if each control is selected based on presence of 'Justification to Select'."""
        for control_id, control_info in self.controls.items():
            # A control is selected if it has "Justification to Select"
            control_info["selected"] = "Justification to Select" in control_info["attributes"]
    
    def save_to_json(self, filename: str):
        """Save extracted controls to JSON file."""
        try:
            with open(filename, 'w', encoding='utf-8') as f:
                json.dump(self.controls, f, indent=2, ensure_ascii=False)
            print(f"Controls saved to {filename}")
        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 OVERLAY EXTRACTION SUMMARY ===")
        print(f"Total controls extracted: {len(self.controls)}")
        
        if not self.controls:
            print("No controls found!")
            return
        
        # Count selected vs not selected
        selected_count = sum(1 for control in self.controls.values() if control["selected"])
        not_selected_count = len(self.controls) - selected_count
        
        print(f"Selected controls (with Justification): {selected_count}")
        print(f"Not selected controls (Open Questions only): {not_selected_count}")
        
        # Count controls by family
        families = {}
        enhancements = 0
        base_controls = 0
        enhancements_with_open_questions = 0
        
        for control_id in self.controls.keys():
            family = control_id.split('-')[0]
            families[family] = families.get(family, 0) + 1
            
            if '(' in control_id:
                enhancements += 1
                # Check if this enhancement has Open Questions
                if "Open Questions" in self.controls[control_id]["attributes"]:
                    enhancements_with_open_questions += 1
            else:
                base_controls += 1
        
        print(f"Base controls: {base_controls}")
        print(f"Control enhancements: {enhancements}")
        print(f"Control enhancements with Open Questions: {enhancements_with_open_questions}")
        
        print("\nControls by family:")
        for family, count in sorted(families.items()):
            print(f"  {family}: {count} controls")
        
        # Show sample of selected and not selected controls
        print("\nSample selected controls:")
        selected_controls = [(cid, cdata) for cid, cdata in self.controls.items() if cdata["selected"]]
        for i, (control_id, control_info) in enumerate(selected_controls[:5]):
            control_type = "Enhancement" if '(' in control_id else "Base Control"
            print(f"  {control_id} ({control_type}): {control_info['name'][:50]}...")
        
        print("\nSample not selected controls (Open Questions):")
        not_selected_controls = [(cid, cdata) for cid, cdata in self.controls.items() if not cdata["selected"]]
        for i, (control_id, control_info) in enumerate(not_selected_controls[:5]):
            control_type = "Enhancement" if '(' in control_id else "Base Control"
            print(f"  {control_id} ({control_type}): {control_info['name'][:50]}...")
            
        # Show sample of enhancements with Open Questions
        print("\nSample control enhancements with Open Questions:")
        enhancement_open_questions = [(cid, cdata) for cid, cdata in self.controls.items() 
                                    if '(' in cid and "Open Questions" in cdata["attributes"]]
        for i, (control_id, control_info) in enumerate(enhancement_open_questions[:5]):
            open_questions = control_info["attributes"]["Open Questions"]
            print(f"  {control_id}: {open_questions[:60]}...")
    
    def debug_page(self, page_num: int):
        """Debug a specific page."""
        try:
            doc = fitz.open(self.pdf_path)
            if page_num <= len(doc):
                page = doc[page_num - 1]
                
                print(f"\n=== DEBUG PAGE {page_num} ===")
                
                text_dict = page.get_text("dict")
                
                for block_num, block in enumerate(text_dict.get("blocks", [])):
                    if block.get("type") == 0:
                        print(f"Block {block_num}:")
                        for line_num, line in enumerate(block.get("lines", [])):
                            line_text = ""
                            for span in line.get("spans", []):
                                line_text += span.get("text", "")
                            
                            if line_text.strip():
                                print(f"  Line: '{line_text.strip()}'")
                                
                                # Check if potential attribute
                                for attr in self.known_attributes:
                                    if line_text.strip().startswith(attr + ":"):
                                        print(f"    *** FOUND ATTRIBUTE: {attr} ***")
                                
                                # Check if potential control ID  
                                if re.match(r'^([A-Z]{2}-\d{1,2}(?:\(\d+\))?)', line_text.strip()):
                                    print(f"    *** POTENTIAL CONTROL ID ***")
                                
                                # Check if potential enhancement
                                if (line_text.strip().startswith('●') or line_text.strip().startswith('•')):
                                    enhancement_match = re.search(r'([A-Z]{2}-\d{1,2}\(\d+\))', line_text)
                                    if enhancement_match:
                                        print(f"    *** POTENTIAL ENHANCEMENT: {enhancement_match.group(1)} ***")
                                
                                for span in line.get("spans", []):
                                    text = span.get("text", "")
                                    flags = span.get("flags", 0)
                                    font = span.get("font", "")
                                    if text.strip():
                                        bold = bool(flags & 16)
                                        underline = bool(flags & 2)
                                        print(f"    Span: '{text}' | Font: {font} | Bold: {bold} | Underline: {underline}")
            
            doc.close()
            
        except Exception as e:
            print(f"Error debugging page: {e}")

def main():
    if len(sys.argv) < 2:
        print("Usage: python sl5_overlay_extractor.py <pdf_file> [--debug-page N]")
        print("Example: python sl5_overlay_extractor.py sl5_overlay.pdf")
        print("Example: python sl5_overlay_extractor.py sl5_overlay.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 = SL5ControlExtractor(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 = "extracted_sl5_overlay.json"
    
    print("SL5 Overlay Control Extractor")
    print("=" * 40)
    
    extractor = SL5ControlExtractor(pdf_file)
    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()