runbooks 0.7.6__py3-none-any.whl → 0.7.9__py3-none-any.whl

runbooks/security/run_script.py

@@ -16,6 +16,14 @@ import argparse
 import sys
 
 from runbooks.utils.logger import configure_logger
+from runbooks.common.rich_utils import (
+    console,
+    create_panel,
+    print_error,
+    print_info,
+    print_success,
+    print_warning,
+)
 
 from .security_baseline_tester import SecurityBaselineTester
 
@@ -65,17 +73,35 @@ def main():
     try:
         args = parse_arguments()
 
-        logger.info("Starting AWS Security Baseline Tester...")
-        logger.info(f"Using AWS profile: {args.profile}")
-        logger.info(f"Report language: {args.language}")
-        logger.info(f"Output directory: {args.output}")
+        # Display startup information with Rich formatting
+        startup_info = f"""[bold cyan]AWS Security Baseline Tester[/bold cyan]
+
+[green]Configuration:[/green]
+  [cyan]AWS Profile:[/cyan] {args.profile}
+  [cyan]Language:[/cyan] {args.language}
+  [cyan]Output Directory:[/cyan] {args.output or './results'}
+  
+[dim]Starting comprehensive security assessment...[/dim]"""
+
+        console.print(create_panel(
+            startup_info,
+            title="🔒 Security Baseline Tester",
+            border_style="cyan"
+        ))
+
+        print_info("Initializing AWS Security Baseline Tester...")
+        print_info(f"Using AWS profile: {args.profile}")
+        print_info(f"Report language: {args.language}")
+        print_info(f"Output directory: {args.output or './results'}")
 
         ## Instantiate and run the Security Baseline Tester
         tester = SecurityBaselineTester(args.profile, args.language, args.output)
         tester.run()
 
-        logger.info("AWS Security Baseline testing completed successfully.")
+        print_success("AWS Security Baseline testing completed successfully!")
+        
     except Exception as e:
+        print_error(f"An unexpected error occurred: {e}", exception=e)
         logger.error(f"An unexpected error occurred: {e}", exc_info=True)
         sys.exit(1)
 

runbooks/security/security_baseline_tester.py

@@ -5,26 +5,43 @@ import logging
 import os
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from pathlib import Path
+from typing import Any, Dict, List, Optional
 
 import boto3
 import botocore
 
+from runbooks.common.rich_utils import (
+    STATUS_INDICATORS,
+    console,
+    create_panel,
+    create_progress_bar,
+    create_table,
+    print_error,
+    print_info,
+    print_status,
+    print_success,
+    print_warning,
+)
+
 from . import (
     checklist,  # noqa: F403
     report_generator,
 )
+from .security_export import SecurityExporter
 from .utils import common, language, level_const
 
 # from .utils.language import get_translator
 
 
 class SecurityBaselineTester:
-    def __init__(self, profile, lang_code, output_dir):
+    def __init__(self, profile, lang_code, output_dir, export_formats: List[str] = None):
         self.profile = profile
         self.language = lang_code
         self.output = output_dir
+        self.export_formats = export_formats or ["json", "csv"]
         self.session = self._create_session()
         self.config = self._load_config()
+        self.exporter = SecurityExporter(output_dir)
         ## Call module 'language' and pass the string 'lang_code'
         self.translator = language.get_translator("main", lang_code)
 
@@ -50,22 +67,41 @@ class SecurityBaselineTester:
             raise
 
     def run(self):
+        """Execute the security baseline assessment with Rich CLI output."""
         try:
+            # Print security assessment header
+            console.print(create_panel(
+                "[bold cyan]AWS Security Baseline Assessment[/bold cyan]\n\n"
+                f"[dim]Profile: {self.profile} | Language: {self.language}[/dim]",
+                title="🛡️ Starting Security Assessment",
+                border_style="cyan"
+            ))
+            
             self._validate_session()
             caller_identity = self._get_caller_identity()
             self._print_auditor_info(caller_identity)
 
-            logging.info(self.translator.translate("start_test"))
+            print_info("Initiating comprehensive security baseline tests...")
 
             account_id, results = self._execute_tests()
             self._generate_report(account_id, results)
-
-            logging.info(self.translator.translate("test_completed"))
+            
+            # Export results in multiple formats
+            if self.export_formats:
+                print_info("Exporting security assessment results...")
+                self.exporter.export_security_results(
+                    account_id=account_id,
+                    results=results,
+                    language=self.language,
+                    formats=self.export_formats
+                )
+
+            print_success("Security baseline assessment completed successfully!")
+            
         except Exception as e:
-            logging.error(
-                f"An error occurred during the security baseline test: {str(e)}",
-                exc_info=True,
-            )
+            print_error(f"Security baseline test failed: {str(e)}", exception=e)
+            logging.error(f"An error occurred during the security baseline test: {str(e)}", exc_info=True)
+            raise
 
     def _validate_session(self):
         if self.session.region_name is None:
@@ -79,35 +115,122 @@ class SecurityBaselineTester:
             raise
 
     def _print_auditor_info(self, caller_identity):
-        logging.info("==================== AUDITOR INFO ====================")
-        logging.info(f"USER ID : {caller_identity['UserId']}")
-        logging.info(f"ACCOUNT : {caller_identity['Account']}")
-        logging.info(f"ARN     : {caller_identity['Arn']}")
-        logging.info("=====================================================")
+        """Display auditor information with Rich formatting."""
+        auditor_info = f"""[bold cyan]User ID:[/bold cyan] {caller_identity['UserId']}
+[bold cyan]Account:[/bold cyan] {caller_identity['Account']}
+[bold cyan]ARN:[/bold cyan] {caller_identity['Arn']}"""
+        
+        console.print(create_panel(
+            auditor_info,
+            title="🔐 Security Assessment Context",
+            border_style="cyan",
+            padding=1
+        ))
 
     def _execute_tests(self):
         iam_client = self.session.client("iam")
         sts_client = self.session.client("sts")
 
         account_id = common.get_account_id(sts_client)
-        logging.info(self.translator.translate("request_credential_report"))
+        print_info(f"Generating credential report for account {account_id}")
         credential_report = common.generate_credential_report(iam_client)
 
-        with ThreadPoolExecutor(max_workers=self.config.get("max_workers", 5)) as executor:
-            futures = {
-                executor.submit(self._run_check, check_name, credential_report): check_name
-                for check_name in self.config.get("checks", [])
-            }
-
-            results = {
-                level: [] for level in ["Success", "Warning", "Danger", "Error", "Info"] if isinstance(level, str)
-            }
-            for future in as_completed(futures):
-                result = future.result()
-                results[result.level].append(result)
-
+        # Create progress bar for security checks
+        checks = self.config.get("checks", [])
+        total_checks = len(checks)
+        
+        with create_progress_bar(description="Security Assessment") as progress:
+            task = progress.add_task("Running security checks...", total=total_checks)
+            
+            with ThreadPoolExecutor(max_workers=self.config.get("max_workers", 5)) as executor:
+                futures = {
+                    executor.submit(self._run_check, check_name, credential_report): check_name
+                    for check_name in checks
+                }
+
+                results = {
+                    level: [] for level in ["Success", "Warning", "Danger", "Error", "Info"] if isinstance(level, str)
+                }
+                completed_checks = 0
+                
+                for future in as_completed(futures):
+                    result = future.result()
+                    results[result.level].append(result)
+                    completed_checks += 1
+                    progress.update(task, completed=completed_checks)
+
+        # Display security assessment summary
+        self._display_security_summary(results, total_checks)
         return account_id, results
 
+    def _display_security_summary(self, results: Dict[str, List], total_checks: int):
+        """Display security assessment summary with Rich formatting."""
+        # Create summary table
+        summary_table = create_table(
+            title="🛡️ Security Assessment Summary",
+            columns=[
+                {"name": "Status", "style": "bold", "justify": "left"},
+                {"name": "Count", "style": "bold", "justify": "center"},
+                {"name": "Percentage", "style": "dim", "justify": "right"},
+            ]
+        )
+        
+        # Calculate statistics
+        for level in ["Danger", "Warning", "Success", "Info", "Error"]:
+            count = len(results.get(level, []))
+            if total_checks > 0:
+                percentage = (count / total_checks) * 100
+                percentage_str = f"{percentage:.1f}%"
+            else:
+                percentage_str = "0%"
+                
+            # Style based on level
+            if level == "Danger":
+                status_text = f"🔴 {level}"
+                style = "error"
+            elif level == "Warning":
+                status_text = f"🟡 {level}"
+                style = "warning"
+            elif level == "Success":
+                status_text = f"🟢 {level}"
+                style = "success"
+            elif level == "Info":
+                status_text = f"🔵 {level}"
+                style = "info"
+            else:  # Error
+                status_text = f"❌ {level}"
+                style = "critical"
+                
+            summary_table.add_row(status_text, str(count), percentage_str, style=style)
+        
+        console.print(summary_table)
+        
+        # Calculate overall security score
+        total_issues = len(results.get("Danger", [])) + len(results.get("Warning", []))
+        total_success = len(results.get("Success", []))
+        
+        if total_checks > 0:
+            security_score = (total_success / total_checks) * 100
+            if security_score >= 90:
+                score_style = "success"
+                score_icon = "🛡️"
+            elif security_score >= 75:
+                score_style = "warning" 
+                score_icon = "⚠️"
+            else:
+                score_style = "error"
+                score_icon = "🚨"
+                
+            score_summary = f"""[bold {score_style}]{score_icon} Overall Security Score: {security_score:.1f}%[/bold {score_style}]
+
+[dim]Total Checks: {total_checks} | Issues Found: {total_issues} | Successful: {total_success}[/dim]"""
+        
+            console.print(create_panel(
+                score_summary,
+                title="Security Posture Assessment", 
+                border_style=score_style
+            ))
+
     def _run_check(self, check_name, credential_report):
         # check_module = __import__(f"checklist.{check_name}", fromlist=[check_name])
         check_module = importlib.import_module(f"runbooks.security.checklist.{check_name}")
@@ -164,9 +287,9 @@ class SecurityBaselineTester:
 
         if not results_dir.exists():
             results_dir.mkdir(parents=True, exist_ok=True)
-            logging.info(self.translator.translate("results_folder_created"))
+            print_info(f"Created results directory: {results_dir}")
         else:
-            logging.info(self.translator.translate("results_folder_already_exists"))
+            print_info(f"Using existing results directory: {results_dir}")
 
         return results_dir
 
@@ -198,4 +321,20 @@ class SecurityBaselineTester:
         with report_path.open("w", encoding="utf-8") as file:
             file.write(html_report)
 
-        logging.info(self.translator.translate("generate_result_report"))
+        # Display report generation success with Rich formatting
+        report_success = f"""[bold green]Security Report Generated Successfully[/bold green]
+
+[cyan]Report Location:[/cyan] {report_path}
+[cyan]Account ID:[/cyan] {account_id}
+[cyan]Language:[/cyan] {self.language}
+[cyan]Report Time:[/cyan] {current_time}
+
+[dim]Open the HTML report in your browser to view detailed findings.[/dim]"""
+
+        console.print(create_panel(
+            report_success,
+            title="📊 Report Generation Complete",
+            border_style="green"
+        ))
+        
+        print_success(f"HTML report saved to: {report_path}")