regscale-cli 6.16.0.0__py3-none-any.whl → 6.16.1.0__py3-none-any.whl

regscale/integrations/commercial/wizv2/utils.py

@@ -1,3 +1,7 @@
+"""
+Wiz V2 Utils
+"""
+
 import codecs
 import csv
 import datetime
@@ -32,10 +36,10 @@ from regscale.integrations.commercial.wizv2.constants import (
     MAX_RETRIES,
     CHECK_INTERVAL_FOR_DOWNLOAD_REPORT,
 )
+from regscale.integrations.commercial.wizv2.models import ComplianceReport, ComplianceCheckStatus
 from regscale.integrations.commercial.wizv2.variables import WizVariables
 from regscale.integrations.commercial.wizv2.wiz_auth import wiz_authenticate
 from regscale.models import File, Sbom, SecurityPlan, Catalog, ControlImplementation, Assessment, regscale_models
-from regscale.integrations.commercial.wizv2.models import ComplianceReport, ComplianceCheckStatus
 from regscale.utils import PaginatedGraphQLClient
 from regscale.utils.decorators import deprecated
 
@@ -51,41 +55,29 @@ def get_notes_from_wiz_props(wiz_entity_properties: Dict, external_id: str) -> s
     :return: Notes
     :rtype: str
     """
+    # Define property mappings with display names and keys
+    property_mappings = [
+        ("External ID", lambda: external_id, lambda x: x),
+        ("Cloud Platform", "cloudPlatform", str),
+        ("Provider Unique ID", "providerUniqueId", str),
+        ("cloudProviderURL", "cloudProviderURL", lambda x: f'<a href="{x}" target="_blank">{x}</a>'),
+        ("Vertex ID", "_vertexID", str),
+        ("Severity Name", "severity_name", str),
+        ("Severity Description", "severity_description", str),
+    ]
+
     notes = []
-    notes.append(f"External ID: {external_id}") if external_id else None
-    (
-        notes.append(f"Cloud Platform: {wiz_entity_properties.get('cloudPlatform')}")
-        if wiz_entity_properties.get("cloudPlatform")
-        else None
-    )
-    (
-        notes.append(f"Provider Unique ID: {wiz_entity_properties.get('providerUniqueId')}")
-        if wiz_entity_properties.get("providerUniqueId")
-        else None
-    )
-    (
-        notes.append(
-            f"""cloudProviderURL:<a href="{wiz_entity_properties.get("cloudProviderURL")}"
-                            target="_blank">{wiz_entity_properties.get("cloudProviderURL")}</a>"""
-        )
-        if wiz_entity_properties.get("cloudProviderURL")
-        else None
-    )
-    (
-        notes.append(f"Vertex ID: {wiz_entity_properties.get('_vertexID')}")
-        if wiz_entity_properties.get("_vertexID")
-        else None
-    )
-    (
-        notes.append(f"Severity Name: {wiz_entity_properties.get('severity_name')}")
-        if wiz_entity_properties.get("severity_name")
-        else None
-    )
-    (
-        notes.append(f"Severity Description: {wiz_entity_properties.get('severity_description')}")
-        if wiz_entity_properties.get("severity_description")
-        else None
-    )
+    for display_name, key_or_func, formatter in property_mappings:
+        # Handle external_id special case
+        if callable(key_or_func):
+            value = key_or_func()
+        else:
+            value = wiz_entity_properties.get(key_or_func)
+
+        if value:
+            formatted_value = formatter(value)
+            notes.append(f"{display_name}: {formatted_value}")
+
     return "<br>".join(notes)
 
 
@@ -108,7 +100,6 @@ def create_asset_type(asset_type: str) -> str:
     :return: Asset type
     :rtype: str
     """
-    #
     asset_type = asset_type.title().replace("_", " ")
     meta_data_list = regscale_models.Metadata.get_metadata_by_module_field(module="assets", field="Asset Type")
     if not any(meta_data.value == asset_type for meta_data in meta_data_list):
@@ -360,8 +351,9 @@ def create_report_if_needed(
         )
         logger.info(f"Wiz compliance report created with ID {wiz_report_id}")
         return [wiz_report_id]
-
-    return [report["id"] for report in reports if any(frame in report["name"] for frame in frames)]
+    logger.debug(f"Returning report ids for these reports {(report['name'] for report in reports)}")
+    reports = [report["id"] for report in reports if any(frame in report["name"] for frame in frames)]
+    return reports
 
 
 def fetch_and_process_report_data(wiz_report_ids: List) -> List:
@@ -384,6 +376,95 @@ def fetch_and_process_report_data(wiz_report_ids: List) -> List:
     return report_data
 
 
+def get_or_create_report_id(
+    project_id: str,
+    frameworks: List[str],
+    wiz_frameworks: List[Dict],
+    existing_reports: List[Dict],
+    target_framework: str,
+) -> str:
+    """
+    Get an existing report ID or create a new one for the target framework.
+
+    :param project_id: Project identifier
+    :param frameworks: List of framework names
+    :param wiz_frameworks: List of framework details with IDs
+    :param existing_reports: List of existing reports
+    :param target_framework: Target framework name with underscores
+    :return: Single report ID
+    """
+    report_name = f"{target_framework}_project_{project_id}"
+
+    # Check for existing report with exact name
+    for report in existing_reports:
+        if report.get("name") == report_name:
+            logger.info(f"Found existing report '{report_name}' with ID {report['id']}")
+            return report["id"]
+
+    # Create new report if no exact match found
+    try:
+        framework_index = frameworks.index(target_framework)
+        framework_id = wiz_frameworks[framework_index].get("id")
+
+        report_id = create_compliance_report(
+            wiz_project_id=project_id, report_name=report_name, framework_id=framework_id
+        )
+        logger.info(f"Created new report '{report_name}' with ID {report_id}")
+        return report_id
+    except ValueError:
+        logger.error(f"Framework '{target_framework}' not found in frameworks list")
+        raise
+
+
+def fetch_report_data(report_id: str) -> List[Dict]:
+    """
+    Fetch and process data for a single report ID.
+
+    :param report_id: Report identifier
+    :return: List of report data rows
+    """
+    try:
+        download_url = get_report_url_and_status(report_id)
+        logger.info(f"Fetching report {report_id} from: {download_url}")
+
+        with closing(requests.get(url=download_url, stream=True, timeout=10)) as response:
+            response.raise_for_status()
+            logger.info(f"Streaming and parsing report {report_id}")
+
+            reader = csv.DictReader(codecs.iterdecode(response.iter_lines(), encoding="utf-8"), delimiter=",")
+            return list(reader)
+    except requests.RequestException as e:
+        error_and_exit(f"Failed to fetch report {report_id}: {str(e)}")
+    except csv.Error as e:
+        error_and_exit(f"Failed to parse CSV for report {report_id}: {str(e)}")
+
+
+# Usage example
+def process_single_report(
+    project_id: str,
+    frameworks: List[str],
+    wiz_frameworks: List[Dict],
+    existing_reports: List[Dict],
+    target_framework: str,
+) -> List[Dict]:
+    """Process a single report and return its data.
+    :param project_id: Project identifier
+    :param frameworks: List of framework names
+    :param wiz_frameworks: List of framework details with IDs
+    :param existing_reports: List of existing reports
+    :param target_framework: Target framework name with underscores
+    :return: List of report data rows
+    """
+    report_id = get_or_create_report_id(
+        project_id=project_id,
+        frameworks=frameworks,
+        wiz_frameworks=wiz_frameworks,
+        existing_reports=existing_reports,
+        target_framework=target_framework,
+    )
+    return fetch_report_data(report_id)
+
+
 def fetch_framework_report(wiz_project_id: str, snake_framework: str) -> List[Any]:
     """
     Fetch Framework Report from Wiz.
@@ -395,10 +476,17 @@ def fetch_framework_report(wiz_project_id: str, snake_framework: str) -> List[An
     """
     wiz_frameworks = fetch_frameworks()
     frames = get_framework_names(wiz_frameworks)
-    reports = list(query_reports())
+    reports = list(query_reports(wiz_project_id))
+
+    report_data = process_single_report(
+        project_id=wiz_project_id,
+        frameworks=frames,
+        wiz_frameworks=wiz_frameworks,
+        existing_reports=reports,
+        target_framework=snake_framework,
+    )
 
-    wiz_report_ids = create_report_if_needed(wiz_project_id, frames, wiz_frameworks, reports, snake_framework)
-    return fetch_and_process_report_data(wiz_report_ids)
+    return report_data
 
 
 def fetch_frameworks() -> list:
@@ -425,23 +513,21 @@ def fetch_frameworks() -> list:
     """
     variables = {
         "policyTypes": "CLOUD",
-        "first": 500,
     }
     resp = send_request(
         query=query,
         variables=variables,
         api_endpoint_url=WizVariables.wizUrl,
     )
-
-    if resp.ok:
-        # ["data"]["securityFrameworks"]["nodes"]
+    logger.debug(f"Response: {resp}")
+    if resp and resp.ok:
         data = resp.json()
         return data.get("data", {}).get("securityFrameworks", {}).get("nodes")
     else:
         error_and_exit(f"Wiz Error: {resp.status_code if resp else None} - {resp.text if resp else 'No response'}")
 
 
-def query_reports() -> list:
+def query_reports(wiz_project_id: str) -> list:
     """
     Query Report table from Wiz
 
@@ -450,7 +536,7 @@ def query_reports() -> list:
     """
 
     # The variables sent along with the above query
-    variables = {"first": 100, "filterBy": {}}
+    variables = {"first": 100, "filterBy": {"projectId": f"{wiz_project_id}"}}
 
     res = send_request(
         query=REPORTS_QUERY,
@@ -462,6 +548,7 @@ def query_reports() -> list:
         if "errors" in res.json().keys():
             error_and_exit(f'Wiz Error: {res.json()["errors"]}')
         json_result = res.json()
+        logger.debug("JSON Result: %s", json_result)
         result = json_result.get("data", {}).get("reports", {}).get("nodes")
     except requests.JSONDecodeError:
         error_and_exit(f"Unable to fetch reports from Wiz: {res.status_code}, {res.reason}")
@@ -583,7 +670,6 @@ def _sync_compliance(
     wiz_project_id: str,
     regscale_id: int,
     regscale_module: str,
-    include_not_implemented: bool,
     client_id: str,
     client_secret: str,
     catalog_id: int,
@@ -595,7 +681,6 @@ def _sync_compliance(
     :param str wiz_project_id: Wiz Project ID
     :param int regscale_id: RegScale ID
     :param str regscale_module: RegScale module
-    :param bool include_not_implemented: Include not implemented controls
     :param str client_id: Wiz Client ID
     :param str client_secret: Wiz Client Secret
     :param int catalog_id: Catalog ID, defaults to None
@@ -631,12 +716,12 @@ def _sync_compliance(
 
     catalog = Catalog.get_with_all_details(catalog_id=catalog_id)
     controls = catalog.get("controls") if catalog else []
-    passing_controls = dict()
-    failing_controls = dict()
-    controls_to_reports = dict()
-    existing_implementations = ControlImplementation.get_existing_control_implementations(parent_id=regscale_id)
+    passing_controls = {}
+    failing_controls = {}
+    controls_to_reports = {}
+
     compliance_job_progress.update(fetch_regscale_data_job, completed=True, advance=1)
-    logger.info(f"Analyzing ComplianceReport for framework {sync_framework} from Wiz")
+    logger.info("Analyzing ComplianceReport for framework %s from Wiz" % sync_framework)
     running_compliance_job = compliance_job_progress.add_task(
         "[#f68d1f]Building compliance posture from wiz report...",
         total=len(report_data),
@@ -654,20 +739,11 @@ def _sync_compliance(
                 )
                 report_models.append(cr)
                 compliance_job_progress.update(running_compliance_job, advance=1)
-        except ValidationError as e:
-            logger.error(f"Error creating ComplianceReport: {e}")
+        except ValidationError:
+            error_message = traceback.format_exc()
+            logger.error(f"Error creating ComplianceReport: {error_message}")
     try:
         saving_regscale_data_job = compliance_job_progress.add_task("[#f68d1f]Saving RegScale data...", total=1)
-        ControlImplementation.create_control_implementations(
-            controls=controls,
-            parent_id=regscale_id,
-            parent_module=regscale_module,
-            existing_implementation_dict=existing_implementations,
-            full_controls=passing_controls,
-            partial_controls={},
-            failing_controls=failing_controls,
-            include_not_implemented=include_not_implemented,
-        )
         create_assessment_from_compliance_report(
             controls_to_reports=controls_to_reports,
             regscale_id=regscale_id,
@@ -676,9 +752,9 @@ def _sync_compliance(
         )
         compliance_job_progress.update(saving_regscale_data_job, completed=True, advance=1)
 
-    except Exception as e:
-        logger.error(f"Error creating ControlImplementations from compliance report: {e}")
-        traceback.print_exc()
+    except Exception:
+        error_message = traceback.format_exc()
+        logger.error(f"Error creating ControlImplementations from compliance report: {error_message}")
     return report_models
 
 

regscale/integrations/public/fedramp/import_workbook.py

@@ -217,7 +217,7 @@ def map_inventory_to_asset(
             "ram": 0,
             "diskStorage": 0,
             "description": "",
-            "endOfLifeDate": date_str(mapping.get_value(inventory, "End-of-Life ", "", warnings=False)),
+            "endOfLifeDate": date_str(mapping.get_value(inventory, "End-of-Life ", "")),
             "purchaseDate": None,
             "status": "Active (On Network)",
             "wizId": "",

regscale/integrations/public/fedramp/poam/scanner.py

@@ -4,7 +4,10 @@
 
 import logging
 import re
-from typing import Iterator, List, Optional
+from typing import Iterator, List, Optional, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import numpy
 
 from openpyxl import load_workbook  # type: ignore
 from openpyxl.utils import column_index_from_string  # type: ignore
@@ -25,6 +28,14 @@ from regscale.validation.address import validate_ip_address, validate_mac_addres
 
 logger = logging.getLogger("regscale")
 
+WEAKNESS_DETECTOR_SOURCE = "Weakness Detector Source"
+ASSET_IDENTIFIER = "Asset Identifier"
+SCHEDULED_COMPLETION_DATE = "Scheduled Completion Date"
+MILESTONE_CHANGES = "Milestone Changes"
+ORIGINAL_RISK_RATING = "Original Risk Rating"
+ADJUSTED_RISK_RATING = "Adjusted Risk Rating"
+FILE_PATH_ERROR = "File path is required."
+
 
 class FedrampPoamIntegration(ScannerIntegration):
     """Integration class for FedRAMP POAM scanning."""
@@ -50,40 +61,43 @@ class FedrampPoamIntegration(ScannerIntegration):
     skipped_records: int = 0
     processed_assets: set[str] = set()  # Track processed assets across all methods
 
-    # TODO: Pair this down to only usable data
     fedramp_poam_columns = [
         "POAM ID",
         "Weakness Name",
         "Weakness Description",
-        "Weakness Detector Source",
+        WEAKNESS_DETECTOR_SOURCE,
         "Weakness Source Identifier",
-        "Asset Identifier",
+        ASSET_IDENTIFIER,
         "Point of Contact",
         "Resources Required",
         "Overall Remediation Plan",
         "Original Detection Date",
-        "Scheduled Completion Date",
+        SCHEDULED_COMPLETION_DATE,
         "Planned Milestones",
-        "Milestone Changes",
+        MILESTONE_CHANGES,
         "Status Date",
-        # "Vendor Dependency",
-        # "Last Vendor Check-in Date",
-        # "Vendor Dependent Product Name",
-        "Original Risk Rating",
-        "Adjusted Risk Rating",
+        ORIGINAL_RISK_RATING,
+        ADJUSTED_RISK_RATING,
         "Risk Adjustment",
         "False Positive",
         "Operational Requirement",
         "Deviation Rationale",
-        # "Supporting Documents",
         "Comments",
-        # "Auto-Approve",
-        # "Binding Operational Directive 22-01 tracking",
-        # "Binding Operational Directive 22-01 Due Date",
-        # "CVE",
-        # "Service Name",
     ]
 
+    """
+    Unused columns:
+    # "Vendor Dependency",
+    # "Last Vendor Check-in Date",
+    # "Vendor Dependent Product Name",
+    # "Supporting Documents",
+    # "Auto-Approve",
+    # "Binding Operational Directive 22-01 tracking",
+    # "Binding Operational Directive 22-01 Due Date",
+    # "CVE",
+    # "Service Name",
+    """
+
     def __init__(self, plan_id: int, **kwargs: dict):
         super().__init__(plan_id=plan_id)
         try:
@@ -91,7 +105,7 @@ class FedrampPoamIntegration(ScannerIntegration):
             if "file_path" in kwargs:
                 self.file_path = kwargs["file_path"]
             if not self.file_path:
-                raise ValueError("File path is required")
+                error_and_exit(FILE_PATH_ERROR)
             self.workbook = self.workbook or load_workbook(filename=self.file_path, data_only=True, read_only=True)
             self.poam_sheets = kwargs.get("poam_sheets") or [
                 sheet for sheet in self.workbook.sheetnames if re.search("POA&M Items", sheet)
@@ -130,11 +144,10 @@ class FedrampPoamIntegration(ScannerIntegration):
         """
         Fetches findings from FedRAMP POAM files.
 
-        :raises ValueError: If file path is not set
         :yield: Iterator of validated integration findings
         """
         if not self.file_path:
-            raise ValueError("File path is required")
+            error_and_exit(FILE_PATH_ERROR)
 
         findings = []
         try:
@@ -256,7 +269,7 @@ class FedrampPoamIntegration(ScannerIntegration):
             "critical": IssueSeverity.High.name,
             "low": IssueSeverity.Low.name,
         }
-        res = validator.mapping.get_value(data, "Original Risk Rating")
+        res = validator.mapping.get_value(data, ORIGINAL_RISK_RATING)
         if res.lower() not in [mem.lower() for mem in IssueSeverity.__members__]:
             res = dat_map.get(res.lower(), IssueSeverity.Low.name)
         return res
@@ -345,7 +358,7 @@ class FedrampPoamIntegration(ScannerIntegration):
                 plugin_id_int = abs(hash(poam_id)) % (10**9)
 
             # Get asset identifiers
-            asset_ids = val_mapping.get_value(data, "Asset Identifier")
+            asset_ids = val_mapping.get_value(data, ASSET_IDENTIFIER)
             if not asset_ids:
                 logger.warning(f"No asset identifier found on row {index}, sheet {sheet}. Skipping.")
                 yield from findings
@@ -374,8 +387,8 @@ class FedrampPoamIntegration(ScannerIntegration):
                         date_str(val_mapping.get_value(data, "Original Detection Date")) or get_current_datetime()
                     )
                     due_date = date_str(
-                        val_mapping.get_value(data, "Scheduled Completion Date")
-                        if val_mapping.get_value(data, "Scheduled Completion Date") != "#REF!"
+                        val_mapping.get_value(data, SCHEDULED_COMPLETION_DATE)
+                        if val_mapping.get_value(data, SCHEDULED_COMPLETION_DATE) != "#REF!"
                         else ""
                     )
                     severity: IssueSeverity = getattr(IssueSeverity, category.title(), IssueSeverity.NotAssigned)
@@ -408,18 +421,18 @@ class FedrampPoamIntegration(ScannerIntegration):
                         date_last_updated=status_date,
                         due_date=due_date,
                         cve=cve,  # Single CVE per finding
-                        plugin_name=val_mapping.get_value(data, "Weakness Detector Source") or "",
+                        plugin_name=val_mapping.get_value(data, WEAKNESS_DETECTOR_SOURCE) or "",
                         plugin_id=str(unique_plugin_id),
-                        observations=str(val_mapping.get_value(data, "Milestone Changes")) or "",
+                        observations=str(val_mapping.get_value(data, MILESTONE_CHANGES)) or "",
                         poam_comments=self.empty(val_mapping.get_value(data, "Comments")),
                         remediation=self.empty(val_mapping.get_value(data, "Overall Remediation Plan")),
                         basis_for_adjustment=str(self.get_basis_for_adjustment(val_mapping=val_mapping, data=data)),
                         vulnerability_type="FedRAMP",
-                        source_report=str(val_mapping.get_value(data, "Weakness Detector Source")),
+                        source_report=str(val_mapping.get_value(data, WEAKNESS_DETECTOR_SOURCE)),
                         point_of_contact=str(val_mapping.get_value(data, "Point of Contact")),
-                        milestone_changes=str(val_mapping.get_value(data, "Milestone Changes")),
+                        milestone_changes=str(val_mapping.get_value(data, MILESTONE_CHANGES)),
                         planned_milestone_changes=str(val_mapping.get_value(data, "Planned Milestones")),
-                        adjusted_risk_rating=val_mapping.get_value(data, "Adjusted Risk Rating"),
+                        adjusted_risk_rating=val_mapping.get_value(data, ADJUSTED_RISK_RATING),
                         risk_adjustment=self.determine_risk_adjustment(val_mapping.get_value(data, "Risk Adjustment")),
                         operational_requirements=str(val_mapping.get_value(data, "Operational Requirement")),
                         deviation_rationale=str(val_mapping.get_value(data, "Deviation Rationale")),
@@ -452,7 +465,7 @@ class FedrampPoamIntegration(ScannerIntegration):
         res = date_str(status_map.get(resolve_status), "%m-%d-%Y")
         if res:
             logger.warning(
-                f"Status Date missing on row %i, sheet %s, defaulting to %s: %s",
+                "Status Date missing on row %i, sheet %s, defaulting to %s: %s",
                 index,
                 sheet,
                 resolve_status.lower().replace("_", " "),
@@ -482,7 +495,7 @@ class FedrampPoamIntegration(ScannerIntegration):
             else:
                 logger.error("Validator mapping is None")
                 return row_assets
-            asset_ids = val_mapping.get_value(data, "Asset Identifier")
+            asset_ids = val_mapping.get_value(data, ASSET_IDENTIFIER)
             if not asset_ids:
                 return row_assets
             asset_id_list = self.gen_asset_list(asset_ids)
@@ -643,8 +656,8 @@ class FedrampPoamIntegration(ScannerIntegration):
         :rtype: Optional[str]
         """
         basis_for_adjustment = self.empty(val_mapping.get_value(data, "Comments"))  # e.g. row 23
-        risk_rating = val_mapping.get_value(data, "Original Risk Rating")
-        adjusted_risk_rating = val_mapping.get_value(data, "Adjusted Risk Rating")
+        risk_rating = val_mapping.get_value(data, ORIGINAL_RISK_RATING)
+        adjusted_risk_rating = val_mapping.get_value(data, ADJUSTED_RISK_RATING)
 
         if (adjusted_risk_rating != risk_rating) and not basis_for_adjustment:
             return "POAM Import"
@@ -745,19 +758,12 @@ class FedrampPoamIntegration(ScannerIntegration):
         """
         Fetch assets from FedRAMP POAM files.
 
-        Args:
-            *args: Variable length argument list
-            **kwargs: Arbitrary keyword arguments
-
-        Returns:
-            Iterator[IntegrationAsset]: Iterator of parsed integration assets
-
-        Raises:
-            ValueError: If file_path is not set
-            POAMProcessingError: If there's an error processing the POAM file
+        :raises POAMProcessingError: If there's an error processing the POAM file
+        :return: Iterator of parsed integration assets
+        :rtype: Iterator[IntegrationAsset]
         """
         if not self.file_path:
-            raise ValueError("File path is required")
+            error_and_exit(FILE_PATH_ERROR)
 
         assets = []
         total_processed = 0
@@ -808,6 +814,7 @@ class FedrampPoamIntegration(ScannerIntegration):
         finally:
             logger.info(f"Completed processing with {total_processed} assets and {self.error_records} errors")
 
+        self.num_assets_to_process = len(assets)
         return iter(assets)
 
     def find_max_row(self, start_row: int, ws: Worksheet) -> int:

regscale/integrations/public/fedramp/ssp_logger.py

@@ -8,23 +8,22 @@ import logging
 from regscale.core.app.logz import create_logger
 
 
-class CaptureEventsHandler:
+class CaptureEventsHandler(logging.Handler):
     def __init__(self, events, errors, infos):
-        self.handler = logging.Handler()
         self.events = events
         self.errors = errors
         self.infos = infos
+        super().__init__()
 
     def emit(self, record):
-        self.handler.emit(record)
         try:
-            log_entry = self.handler.format(record)
+            log_entry = self.format(record)
             if record.levelname == "INFO":
                 self.events.append(log_entry)
             elif record.levelname == "ERROR":
                 self.errors.append(log_entry)
         except Exception:
-            self.handler.handleError(record)
+            self.handleError(record)
 
 
 class SSPLogger:
@@ -33,7 +32,8 @@ class SSPLogger:
         self.errors = []
         self.infos = []
         self.capture_handler = CaptureEventsHandler(self.events, self.errors, self.infos)
-        self.logger = create_logger(custom_handler=self.capture_handler)
+        logger = create_logger(custom_handler=self.capture_handler)
+        self.logger = logger
 
     def create_logger(self):
         return self.logger