main.py

#!/usr/bin/env python3
"""
PHI Scanner - Salesforce PHI Field Detection Tool

Scans Salesforce org metadata and data to identify fields that may
contain Protected Health Information (PHI) and require encryption.

Usage:
    # Interactive mode - select objects from a list
    python main.py --org myorg --interactive

    # Full scan of default org
    python main.py --org myorg

    # Metadata-only scan from local SFDX project
    python main.py --mode metadata-only --source ./force-app

    # Hybrid: local metadata + API sampling
    python main.py --mode hybrid --org myorg --source ./force-app

    # Scan specific objects
    python main.py --org myorg --objects Account,Contact,Lead

    # Output as CSV
    python main.py --org myorg --format csv --output ./reports/phi-audit.csv

    # Use custom config file
    python main.py --config ./my-config.yaml
"""

import argparse
import sys
from datetime import datetime
from typing import Dict, List

from scanner.config import ScanConfig
from scanner.connection import SalesforceConnection, list_authenticated_orgs
from scanner.metadata import MetadataScanner, ObjectInfo
from scanner.categorizer import PHICategorizer, PHIAssessment
from scanner.sampler import DataSampler, SampleResult
from scanner.reporter import ReportGenerator
from scanner.interactive import interactive_object_selection, confirm_selection
from scanner.encryption import EncryptionChecker


def print_header():
    """Print tool header."""
    print("=" * 60)
    print("PHI SCANNER - Salesforce PHI Field Detection Tool")
    print("=" * 60)
    print(f"Started: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print()


def print_progress(message: str):
    """Print progress message."""
    print(f"  {message}")


def create_parser() -> argparse.ArgumentParser:
    """Create argument parser."""
    parser = argparse.ArgumentParser(
        description="Scan Salesforce orgs for PHI fields requiring encryption",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Full scan using SF CLI auth
  python main.py --org myorg

  # Metadata-only scan (no API calls)
  python main.py --mode metadata-only --source ./force-app

  # Scan specific objects only
  python main.py --org myorg --objects Account,Contact,Lead,Opportunity

  # Output as JSON
  python main.py --org myorg --format json --output ./phi-report.json

  # Use config file
  python main.py --config ./scan-config.yaml

  # List authenticated orgs
  python main.py --list-orgs
        """
    )

    # Connection options
    parser.add_argument(
        "--org", "-o",
        help="SF CLI org alias (e.g., 'myorg'). Uses default org if not specified."
    )
    parser.add_argument(
        "--list-orgs",
        action="store_true",
        help="List all authenticated SF CLI orgs and exit"
    )

    # Scan mode options
    parser.add_argument(
        "--mode", "-m",
        choices=["full", "metadata-only", "hybrid"],
        default="full",
        help="Scan mode: 'full' (API + sampling), 'metadata-only' (local files), 'hybrid' (local + sampling)"
    )
    parser.add_argument(
        "--source", "-s",
        help="Path to force-app directory for local scanning (required for metadata-only/hybrid)"
    )
    parser.add_argument(
        "--objects",
        help="Comma-separated list of objects to scan (e.g., 'Account,Contact,Lead')"
    )

    # Output options
    parser.add_argument(
        "--format", "-f",
        choices=["excel", "csv", "json"],
        default="excel",
        help="Output format (default: excel)"
    )
    parser.add_argument(
        "--output",
        help="Output file path (auto-generated if not specified)"
    )
    parser.add_argument(
        "--org-name",
        help="Organization name for report header"
    )

    # Config options
    parser.add_argument(
        "--config", "-c",
        help="Path to YAML config file"
    )
    parser.add_argument(
        "--patterns",
        help="Path to custom PHI patterns JSON file"
    )

    # Interactive mode
    parser.add_argument(
        "--interactive", "-i",
        action="store_true",
        help="Interactive mode - select objects from a list"
    )

    # Encryption check
    parser.add_argument(
        "--check-encryption",
        action="store_true",
        help="Check if PHI fields are encrypted and report gaps"
    )

    # Web UI
    parser.add_argument(
        "--web",
        action="store_true",
        help="Launch web UI dashboard"
    )
    parser.add_argument(
        "--port",
        type=int,
        default=8000,
        help="Port for web UI (default: 8000)"
    )

    # Verbosity
    parser.add_argument(
        "--quiet", "-q",
        action="store_true",
        help="Suppress progress output"
    )

    return parser


def list_orgs_and_exit():
    """List authenticated orgs and exit."""
    print("Authenticated Salesforce Orgs:")
    print("-" * 40)

    orgs = list_authenticated_orgs()
    if not orgs:
        print("  No authenticated orgs found.")
        print("  Run 'sf org login web --alias <alias>' to authenticate.")
        sys.exit(1)

    for org in orgs:
        default = " (default)" if org.get("isDefaultUsername") else ""
        alias = org.get("alias", "")
        username = org.get("username", "")
        if alias:
            print(f"  {alias}: {username}{default}")
        else:
            print(f"  {username}{default}")

    sys.exit(0)


def run_scan(config: ScanConfig, quiet: bool = False, check_encryption: bool = False) -> str:
    """
    Run the PHI scan with the given configuration.

    Args:
        config: Scan configuration
        quiet: Suppress progress output
        check_encryption: Check encryption status of PHI fields

    Returns:
        Path to generated report file
    """
    def log(msg: str):
        if not quiet:
            print_progress(msg)

    # Validate config
    errors = config.validate()
    if errors:
        print("Configuration errors:")
        for e in errors:
            print(f"  - {e}")
        sys.exit(1)

    # Initialize connection if needed
    connection = None
    if config.mode in ["full", "hybrid"]:
        log(f"Connecting to org: {config.org_alias or 'default'}...")
        try:
            connection = SalesforceConnection(
                org_alias=config.org_alias,
                api_version=config.api_version
            )
            log(f"Connected: {connection.get_org_display_name()}")

            # Use org username as default org name if not set
            if config.org_name == "Salesforce Org":
                config.org_name = connection.get_org_display_name()

        except ConnectionError as e:
            print(f"Connection failed: {e}")
            sys.exit(1)

    # Initialize metadata scanner
    metadata_scanner = MetadataScanner(
        connection=connection,
        source_path=config.source_path
    )

    # Get objects to scan
    log("Scanning metadata...")
    objects = metadata_scanner.get_objects(
        mode=config.mode,
        object_filter=config.objects if config.objects else None,
        exclude_objects=config.exclude_objects
    )

    if not objects:
        print("No objects found to scan.")
        sys.exit(1)

    total_fields = metadata_scanner.get_total_fields(objects)
    log(f"Found {len(objects)} objects with {total_fields} fields")

    # Categorize fields
    log("Categorizing fields for PHI risk...")
    categorizer = PHICategorizer(patterns_file=config.patterns_file)

    assessments: Dict[str, List[PHIAssessment]] = {}
    for obj in objects:
        obj_assessments = categorizer.categorize_fields(obj.fields)
        assessments[obj.api_name] = obj_assessments

    # Get stats
    all_assessments = [a for obj_a in assessments.values() for a in obj_a]
    stats = categorizer.get_summary_stats(all_assessments)
    log(f"PHI Classification: {stats['high']} high, {stats['medium']} medium, {stats['low']} low risk")

    # Sample data if in full/hybrid mode
    samples: Dict[str, List[SampleResult]] = {}
    if config.mode in ["full", "hybrid"] and connection:
        log("Sampling field data...")

        sampler = DataSampler(
            connection=connection,
            sample_limit=config.sample_limit,
            query_timeout=config.query_timeout,
            rate_limit_delay=config.rate_limit_delay,
            patterns_file=config.patterns_file
        )

        def object_progress(obj_name, obj_idx, total):
            log(f"Sampling {obj_name} ({obj_idx}/{total})...")

        samples = sampler.sample_all_objects(objects, progress_callback=object_progress)

        sampling_stats = sampler.get_sampling_stats(samples)
        log(f"Sampled {sampling_stats['sampled']} fields, {sampling_stats['with_data']} with data")

    # Check encryption status if requested
    encryption_status = None
    gap_analysis = None
    if check_encryption and connection:
        log("Checking field encryption status...")
        enc_checker = EncryptionChecker(connection)

        def enc_progress(obj_name, idx, total):
            log(f"Checking encryption: {obj_name} ({idx}/{total})...")

        encryption_status, gap_analysis = enc_checker.analyze_encryption_gaps(
            objects, assessments, progress_callback=enc_progress
        )

        log(f"Encryption coverage: {gap_analysis.coverage_percentage}%")
        log(f"Gaps found: {gap_analysis.high_risk_gaps} high risk, {gap_analysis.medium_risk_gaps} medium risk")

    # Generate report
    log(f"Generating {config.output_format} report...")
    reporter = ReportGenerator(
        org_name=config.org_name,
        output_format=config.output_format
    )

    output_path = reporter.generate_report(
        objects=objects,
        assessments=assessments,
        samples=samples if samples else None,
        output_path=config.output_path,
        encryption_status=encryption_status,
        gap_analysis=gap_analysis
    )

    return output_path


def launch_web_ui(port: int = 8000):
    """Launch the web UI dashboard."""
    try:
        import uvicorn
        from web.app import app
        print(f"Starting PHI Scanner Web UI on http://localhost:{port}")
        print("Press Ctrl+C to stop")
        uvicorn.run(app, host="0.0.0.0", port=port)
    except ImportError as e:
        print("Error: Web UI requires additional dependencies.")
        print("Install with: pip install fastapi uvicorn jinja2 python-multipart")
        print(f"Details: {e}")
        sys.exit(1)


def main():
    """Main entry point."""
    parser = create_parser()
    args = parser.parse_args()

    # Handle list orgs
    if args.list_orgs:
        list_orgs_and_exit()

    # Handle web UI
    if hasattr(args, 'web') and args.web:
        launch_web_ui(args.port)
        return

    # Print header
    if not args.quiet:
        print_header()

    # Build config
    config = ScanConfig.from_args(args)

    # Handle interactive mode
    if hasattr(args, 'interactive') and args.interactive:
        if config.mode not in ["full", "hybrid"]:
            print("Interactive mode requires 'full' or 'hybrid' mode (need org connection).")
            sys.exit(1)

        # Connect to org first
        print(f"Connecting to org: {config.org_alias or 'default'}...")
        try:
            connection = SalesforceConnection(
                org_alias=config.org_alias,
                api_version=config.api_version
            )
            print(f"Connected: {connection.get_org_display_name()}")

            # Update org name if not set
            if config.org_name == "Salesforce Org":
                config.org_name = connection.get_org_display_name()

        except ConnectionError as e:
            print(f"Connection failed: {e}")
            sys.exit(1)

        # Interactive object selection
        selected_objects = interactive_object_selection(connection, preselect_custom=True)

        if not selected_objects:
            print("No objects selected. Exiting.")
            sys.exit(0)

        # Confirm selection
        if not confirm_selection(selected_objects):
            print("Scan cancelled.")
            sys.exit(0)

        # Update config with selected objects
        config.objects = selected_objects
        print()

    # Run scan
    check_encryption = hasattr(args, 'check_encryption') and args.check_encryption
    output_path = run_scan(config, quiet=args.quiet, check_encryption=check_encryption)

    # Print summary
    print()
    print("=" * 60)
    print("SCAN COMPLETE")
    print("=" * 60)
    print(f"Report: {output_path}")
    print(f"Finished: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")


if __name__ == "__main__":
    main()