regscale-cli 6.25.0.1__py3-none-any.whl → 6.26.0.0__py3-none-any.whl

regscale/integrations/scanner_integration.py

@@ -27,6 +27,7 @@ from regscale.integrations.commercial.durosuite.process_devices import scan_duro
 from regscale.integrations.commercial.durosuite.variables import DuroSuiteVariables
 from regscale.integrations.commercial.stig_mapper_integration.mapping_engine import StigMappingEngine as STIGMapper
 from regscale.integrations.due_date_handler import DueDateHandler
+from regscale.integrations.milestone_manager import MilestoneManager
 from regscale.integrations.public.cisa import pull_cisa_kev
 from regscale.integrations.variables import ScannerVariables
 from regscale.models import DateTimeEncoder, OpenIssueDict, Property, regscale_models
@@ -48,6 +49,31 @@ def get_thread_workers_max() -> int:
     return ScannerVariables.threadMaxWorkers
 
 
+def _create_config_override(
+    config: Optional[Dict[str, Dict]],
+    integration_name: str,
+    critical: Optional[int],
+    high: Optional[int],
+    moderate: Optional[int],
+    low: Optional[int],
+) -> Dict[str, Dict]:
+    """Create a config override for legacy parameter support."""
+    override_config = config.copy() if config else {}
+    if "issues" not in override_config:
+        override_config["issues"] = {}
+    if integration_name not in override_config["issues"]:
+        override_config["issues"][integration_name] = {}
+
+    integration_config = override_config["issues"][integration_name]
+    severity_params = {"critical": critical, "high": high, "moderate": moderate, "low": low}
+
+    for param_name, param_value in severity_params.items():
+        if param_value is not None:
+            integration_config[param_name] = param_value
+
+    return override_config
+
+
 def issue_due_date(
     severity: regscale_models.IssueSeverity,
     created_date: str,
@@ -61,6 +87,9 @@ def issue_due_date(
     """
     Calculate the due date for an issue based on its severity and creation date.
 
+    DEPRECATED: This function is kept for backward compatibility. New code should use DueDateHandler directly.
+    This function now uses DueDateHandler internally to ensure consistent behavior and proper validation.
+
     :param regscale_models.IssueSeverity severity: The severity of the issue.
     :param str created_date: The creation date of the issue.
     :param Optional[int] critical: Days until due for high severity issues.
@@ -72,40 +101,19 @@ def issue_due_date(
     :return: The due date for the issue.
     :rtype: str
     """
-    if critical is None:
-        critical = ScannerVariables.issueDueDates.get("critical", 30)
-    if high is None:
-        high = ScannerVariables.issueDueDates.get("high", 60)
-    if moderate is None:
-        moderate = ScannerVariables.issueDueDates.get("moderate", 120)
-    if low is None:
-        low = ScannerVariables.issueDueDates.get("low", 364)
-
-    if config is None:
-        config = {}
-
-    due_date_map = {
-        regscale_models.IssueSeverity.Critical: critical,
-        regscale_models.IssueSeverity.High: high,
-        regscale_models.IssueSeverity.Moderate: moderate,
-        regscale_models.IssueSeverity.Low: low,
-    }
-
-    if title and config:
-        # if title in a config key, use that key
-        issues_dict = config.get("issues", {})
-        matching_key = next((key.lower() for key in issues_dict if title.lower() in key.lower()), None)
-        if matching_key:
-            title_config = issues_dict.get(matching_key, {})
-            due_date_map = {
-                regscale_models.IssueSeverity.Critical: title_config.get("critical", critical),
-                regscale_models.IssueSeverity.High: title_config.get("high", high),
-                regscale_models.IssueSeverity.Moderate: title_config.get("moderate", moderate),
-                regscale_models.IssueSeverity.Low: title_config.get("low", low),
-            }
-
-    days = due_date_map.get(severity, low)
-    return date_str(get_day_increment(start=created_date, days=days))
+    integration_name = title or "default"
+
+    # Check if individual parameters need config override
+    if any(param is not None for param in [critical, high, moderate, low]):
+        config = _create_config_override(config, integration_name, critical, high, moderate, low)
+
+    due_date_handler = DueDateHandler(integration_name, config=config)
+    return due_date_handler.calculate_due_date(
+        severity=severity,
+        created_date=created_date,
+        cve=None,  # Legacy function doesn't have CVE parameter
+        title=title,
+    )
 
 
 class ManagedDefaultDict(Generic[K, V]):
@@ -665,6 +673,9 @@ class ScannerIntegration(ABC):
         # Initialize due date handler for this integration
         self.due_date_handler = DueDateHandler(self.title, config=self.app.config)
 
+        # Initialize milestone manager for this integration
+        self.milestone_manager = None  # Lazy initialization after scan_date is set
+
         if self.is_component:
             self.component = regscale_models.Component.get_object(self.plan_id)
             self.parent_module: str = regscale_models.Component.get_module_string()
@@ -730,6 +741,21 @@ class ScannerIntegration(ABC):
                     cls._lock_registry[key] = lock
         return lock
 
+    def get_milestone_manager(self) -> MilestoneManager:
+        """
+        Get or initialize the milestone manager.
+
+        :return: MilestoneManager instance
+        :rtype: MilestoneManager
+        """
+        if self.milestone_manager is None:
+            self.milestone_manager = MilestoneManager(
+                integration_title=self.title,
+                assessor_id=self.assessor_id,
+                scan_date=self.scan_date or get_current_datetime(),
+            )
+        return self.milestone_manager
+
     @staticmethod
     def load_stig_mapper() -> Optional[STIGMapper]:
         """
@@ -990,15 +1016,18 @@ class ScannerIntegration(ABC):
             return res[:450]
         return prefix[:450]
 
-    def get_or_create_assessment(self, control_implementation_id: int) -> regscale_models.Assessment:
+    def get_or_create_assessment(
+        self, control_implementation_id: int, status: Optional[regscale_models.AssessmentResultsStatus] = None
+    ) -> regscale_models.Assessment:
         """
-        Gets or creates a RegScale assessment
+        Gets or creates a RegScale assessment.
 
         :param int control_implementation_id: The ID of the control implementation
+        :param Optional[regscale_models.AssessmentResultsStatus] status: Optional status override (used by cci_assessment)
         :return: The assessment
         :rtype: regscale_models.Assessment
         """
-        logger.info("Getting or create assessment for control implementation %d", control_implementation_id)
+        logger.debug("Getting or create assessment for control implementation %d", control_implementation_id)
         assessment: Optional[regscale_models.Assessment] = self.assessment_map.get(control_implementation_id)
         if assessment:
             logger.debug(
@@ -1010,7 +1039,7 @@ class ScannerIntegration(ABC):
                 plannedStart=get_current_datetime(),
                 plannedFinish=get_current_datetime(),
                 status=regscale_models.AssessmentStatus.COMPLETE.value,
-                assessmentResult=regscale_models.AssessmentResultsStatus.FAIL.value,
+                assessmentResult=status.value if status else regscale_models.AssessmentResultsStatus.FAIL.value,
                 actualFinish=get_current_datetime(),
                 leadAssessorId=self.assessor_id,
                 parentId=control_implementation_id,
@@ -2106,7 +2135,9 @@ class ScannerIntegration(ABC):
 
     def _set_issue_due_date(self, issue: regscale_models.Issue, finding: IntegrationFinding) -> None:
         """Set the due date for the issue using DueDateHandler."""
+        # Always calculate or validate due date to ensure it's not in the past
         if not finding.due_date:
+            # No due date set, calculate new one
             try:
                 base_created = finding.date_created or issue.dateCreated
                 finding.due_date = self.due_date_handler.calculate_due_date(
@@ -2125,6 +2156,12 @@ class ScannerIntegration(ABC):
                     cve=finding.cve,
                     title=finding.title or self.title,
                 )
+        else:
+            # Due date already exists, but validate it's not in the past (if noPastDueDates is enabled)
+            finding.due_date = self.due_date_handler._ensure_future_due_date(
+                finding.due_date, self.due_date_handler.integration_timelines.get(finding.severity, 60)
+            )
+
         issue.dueDate = finding.due_date
 
     def _set_additional_issue_fields(
@@ -2269,87 +2306,26 @@ class ScannerIntegration(ABC):
         """
         Create milestones for an issue based on status transitions.
 
+        Delegates to MilestoneManager for cleaner separation of concerns.
+        Also ensures existing issues have creation milestones (backfills if missing).
+
         :param regscale_models.Issue issue: The issue to create milestones for
         :param IntegrationFinding finding: The finding data
         :param Optional[regscale_models.Issue] existing_issue: Existing issue for comparison
         """
-        if not (ScannerVariables.useMilestones and issue.id):
-            return
-
-        if self._should_create_reopened_milestone(existing_issue, issue):
-            self._create_milestone_safe(
-                issue, finding, "Issue reopened from", get_current_datetime(), "reopened milestone"
-            )
-        elif self._should_create_closed_milestone(existing_issue, issue):
-            self._create_milestone_safe(issue, finding, "Issue closed from", issue.dateCompleted, "closed milestone")
-        elif not existing_issue:
-            self._create_milestone_safe(issue, finding, "Issue created from", self.scan_date, "new issue milestone")
-        else:
-            logger.debug("No milestone created for issue %s from finding %s", issue.id, finding.external_id)
+        milestone_manager = self.get_milestone_manager()
 
-    def _should_create_reopened_milestone(
-        self, existing_issue: Optional[regscale_models.Issue], issue: regscale_models.Issue
-    ) -> bool:
-        """
-        Check if a reopened milestone should be created.
-
-        :param Optional[regscale_models.Issue] existing_issue: The existing issue
-        :param regscale_models.Issue issue: The current issue
-        :return: True if reopened milestone should be created
-        :rtype: bool
-        """
-        return (
-            existing_issue
-            and existing_issue.status == regscale_models.IssueStatus.Closed
-            and issue.status == regscale_models.IssueStatus.Open
-        )
-
-    def _should_create_closed_milestone(
-        self, existing_issue: Optional[regscale_models.Issue], issue: regscale_models.Issue
-    ) -> bool:
-        """
-        Check if a closed milestone should be created.
+        # For existing issues, ensure they have a creation milestone (backfill if missing)
+        if existing_issue:
+            milestone_manager.ensure_creation_milestone_exists(issue=issue, finding=finding)
 
-        :param Optional[regscale_models.Issue] existing_issue: The existing issue
-        :param regscale_models.Issue issue: The current issue
-        :return: True if closed milestone should be created
-        :rtype: bool
-        """
-        return (
-            existing_issue
-            and existing_issue.status == regscale_models.IssueStatus.Open
-            and issue.status == regscale_models.IssueStatus.Closed
+        # Handle status transition milestones
+        milestone_manager.create_milestones_for_issue(
+            issue=issue,
+            finding=finding,
+            existing_issue=existing_issue,
         )
 
-    def _create_milestone_safe(
-        self,
-        issue: regscale_models.Issue,
-        finding: IntegrationFinding,
-        title_prefix: str,
-        milestone_date: str,
-        milestone_type: str,
-    ) -> None:
-        """
-        Safely create a milestone with error handling.
-
-        :param regscale_models.Issue issue: The issue to create milestone for
-        :param IntegrationFinding finding: The finding data
-        :param str title_prefix: Prefix for milestone title
-        :param str milestone_date: Date for the milestone
-        :param str milestone_type: Description for logging purposes
-        """
-        try:
-            regscale_models.Milestone(
-                title=f"{title_prefix} {self.title} scan",
-                milestoneDate=milestone_date,
-                responsiblePersonId=self.assessor_id,
-                parentID=issue.id,
-                parentModule="issues",
-            ).create_or_update()
-            logger.debug("Added milestone for issue %s from finding %s", issue.id, finding.external_id)
-        except Exception as e:
-            logger.warning("Failed to create %s: %s", milestone_type, str(e))
-
     @staticmethod
     def extra_data_to_properties(finding: IntegrationFinding, issue_id: int) -> None:
         """
@@ -2515,6 +2491,8 @@ class ScannerIntegration(ABC):
             if found_issue.controlImplementationIds:
                 for control_id in found_issue.controlImplementationIds:
                     self.update_control_implementation_status_after_close(control_id)
+                    # Update assessment status to reflect the control implementation status
+                    self.update_assessment_status_from_control_implementation(control_id)
 
     def handle_failing_checklist(
         self,
@@ -2539,11 +2517,13 @@ class ScannerIntegration(ABC):
                 if failing_objective.name.lower().startswith("cci-"):
                     implementation_id = self.get_control_implementation_id_for_cci(failing_objective.name)
                 else:
-                    control_label = objective_to_control_dot(failing_objective.name)
-                    if control_label not in self.control_implementation_id_map:
-                        logger.warning("Control %s not found for %s", control_label, control_label)
-                        continue
-                    implementation_id = self.control_implementation_id_map[control_label]
+                    implementation_id = self._fallback_implementation_id(failing_objective)
+
+                if not implementation_id or implementation_id is None:
+                    logger.warning(
+                        "Could not map objective to a Control Implementation for objective #%i.", failing_objective.id
+                    )
+                    continue
 
                 failing_option = regscale_models.ImplementationOption(
                     name="Failed STIG",
@@ -2565,13 +2545,36 @@ class ScannerIntegration(ABC):
                 ).create_or_update()
 
                 # Create assessment and control test result
-                assessment = self.get_or_create_assessment(implementation_id)
+                assessment = self.get_or_create_assessment(
+                    implementation_id, status=regscale_models.AssessmentResultsStatus.FAIL
+                )
                 if implementation_id:
                     control_test = self.create_or_get_control_test(finding, implementation_id)
                     self.create_control_test_result(
                         finding, control_test, assessment, regscale_models.ControlTestResultStatus.FAIL
                     )
 
+    def _fallback_implementation_id(self, objective: regscale_models.ControlObjective) -> Optional[int]:
+        """
+        Fallback method to get control implementation ID from objective name if CCI mapping fails.
+
+        :param regscale_models.ControlObjective objective: The control objective
+        :return: The control implementation ID if found, None otherwise
+        :rtype: Optional[int]
+        """
+        control_label = objective_to_control_dot(objective.name)
+        if implementation_id := self.control_implementation_id_map.get(control_label):
+            return implementation_id
+
+        if control_id := self.control_id_to_implementation_map.get(objective.securityControlId):
+            if control_label := self.control_map.get(control_id):
+                implementation_id = self.control_implementation_id_map.get(control_label)
+                if not implementation_id:
+                    print("No dice.")
+                return implementation_id
+        logger.debug("Could not find fallback implementation ID for objective #%i", objective.id)
+        return None
+
     def handle_passing_checklist(
         self,
         finding: IntegrationFinding,
@@ -2595,15 +2598,12 @@ class ScannerIntegration(ABC):
                 if passing_objective.name.lower().startswith("cci-"):
                     implementation_id = self.get_control_implementation_id_for_cci(passing_objective.name)
                 else:
-                    control_label = objective_to_control_dot(passing_objective.name)
-                    if control_label not in self.control_implementation_id_map:
-                        logger.warning("Control %s not found for %s", control_label, control_label)
-                        continue
-                    implementation_id = self.control_implementation_id_map[control_label]
-
-                # Skip if we couldn't determine the implementation ID
-                if implementation_id is None:
-                    logger.warning("Could not determine implementation ID for objective %s", passing_objective.name)
+                    implementation_id = self._fallback_implementation_id(passing_objective)
+
+                if not implementation_id or implementation_id is None:
+                    logger.warning(
+                        "Could not map objective to a Control Implementation for objective #%i.", passing_objective.id
+                    )
                     continue
 
                 passing_option = regscale_models.ImplementationOption(
@@ -2626,7 +2626,9 @@ class ScannerIntegration(ABC):
                 ).create_or_update()
 
                 # Create assessment and control test result
-                assessment = self.get_or_create_assessment(implementation_id)
+                assessment = self.get_or_create_assessment(
+                    implementation_id, status=regscale_models.AssessmentResultsStatus.PASS
+                )
                 control_test = self.create_or_get_control_test(finding, implementation_id)
                 self.create_control_test_result(
                     finding, control_test, assessment, regscale_models.ControlTestResultStatus.PASS
@@ -2691,7 +2693,11 @@ class ScannerIntegration(ABC):
                 logger.error("2. Asset not found for identifier %s", finding.asset_identifier)
             return 0
 
-        tool = regscale_models.ChecklistTool.STIGs
+        tool = (
+            regscale_models.ChecklistTool.CISBenchmarks
+            if "simp.cis" in str(finding.vulnerability_number).lower()
+            else regscale_models.ChecklistTool.STIGs
+        )
         if finding.vulnerability_type == "Vulnerability Scan":
             tool = regscale_models.ChecklistTool.VulnerabilityScanner
 
@@ -3366,6 +3372,8 @@ class ScannerIntegration(ABC):
 
         for control_id in affected_control_ids:
             self.update_control_implementation_status_after_close(control_id)
+            # Update assessment status to reflect the control implementation status
+            self.update_assessment_status_from_control_implementation(control_id)
 
         (
             logger.info("Closed %d outdated issues.", closed_count)
@@ -3468,9 +3476,70 @@ class ScannerIntegration(ABC):
         if control_implementation.status != new_status:
             control_implementation.status = new_status
             self.control_implementation_map[control_id] = control_implementation.save()
-            logger.info("Updated control implementation %d status to %s", control_id, new_status)
+            logger.debug("Updated control implementation %d status to %s", control_id, new_status)
+
+    def update_assessment_status_from_control_implementation(self, control_implementation_id: int) -> None:
+        """
+        Updates the assessment status based on the control implementation status.
+        Treats the ControlImplementation status as the source of truth.
+
+        Sets assessment to PASS if ControlImplementation status is FULLY_IMPLEMENTED,
+        otherwise sets it to FAIL.
 
-    def is_issue_protected_from_auto_close(self, issue: regscale_models.Issue) -> bool:
+        This method should be called after update_control_implementation_status_after_close
+        to ensure assessments reflect the final control implementation state.
+
+        :param int control_implementation_id: The ID of the control implementation
+        :rtype: None
+        """
+        # Get the cached assessment for this control implementation
+        assessment = self.assessment_map.get(control_implementation_id)
+
+        if not assessment:
+            logger.debug(
+                "No assessment found in cache for control implementation %d, skipping assessment update",
+                control_implementation_id,
+            )
+            return
+
+        # Get the control implementation to check its status
+        control_implementation = self.control_implementation_map.get(
+            control_implementation_id
+        ) or regscale_models.ControlImplementation.get_object(object_id=control_implementation_id)
+
+        if not control_implementation:
+            logger.warning("Control implementation %d not found, cannot update assessment", control_implementation_id)
+            return
+
+        # Determine assessment result based on control implementation status
+        # Treat ControlImplementation status as the source of truth
+        new_assessment_result = (
+            regscale_models.AssessmentResultsStatus.PASS
+            if control_implementation.status == regscale_models.ImplementationStatus.FULLY_IMPLEMENTED.value
+            else regscale_models.AssessmentResultsStatus.FAIL
+        )
+
+        # Only update if the status has changed
+        if assessment.assessmentResult != new_assessment_result.value:
+            assessment.assessmentResult = new_assessment_result.value
+            assessment.save()
+            logger.debug(
+                "Updated assessment %d for control implementation %d: assessmentResult=%s (based on control status: %s)",
+                assessment.id,
+                control_implementation_id,
+                new_assessment_result.value,
+                control_implementation.status,
+            )
+        else:
+            logger.debug(
+                "Assessment %d already has correct status %s for control implementation %d",
+                assessment.id,
+                assessment.assessmentResult,
+                control_implementation_id,
+            )
+
+    @staticmethod
+    def is_issue_protected_from_auto_close(issue: regscale_models.Issue) -> bool:
         """
         Check if an issue is protected from automatic closure.