regscale-cli 6.21.2.0__py3-none-any.whl → 6.28.2.1__py3-none-any.whl

regscale/models/regscale_models/issue.py

@@ -84,6 +84,35 @@ class IssueStatus(str, Enum):
         return self.value
 
 
+class IssueIdentification(str, Enum):
+    """Issue Identification"""
+
+    A123Review = "A-123 Review"
+    AssessmentAuditInternal = "Assessment/Audit (Internal)"
+    AssessmentAuditExternal = "Assessment/Audit (External)"
+    CriticalControlReview = "Critical Control Review"
+    FDCCUSGCB = "FDCC/USGCB"
+    GAOAudit = "GAO Audit"
+    IGAudit = "IG Audit"
+    IncidentResponseLessonsLearned = "Incident Response Lessons Learned"
+    ITAR = "ITAR"
+    PenetrationTest = "Penetration Test"
+    RiskAssessment = "Risk Assessment"
+    SecurityAuthorization = "Security Authorization"
+    SecurityControlAssessment = "Security Control Assessment"
+    VulnerabilityAssessment = "Vulnerability Assessment"
+    Other = "Other"
+
+    def __str__(self) -> str:
+        """
+        Return the value of the Enum as a string
+
+        :return: The value of the Enum as a string
+        :rtype: str
+        """
+        return self.value
+
+
 class Issue(RegScaleModel):
     """Issue Model"""
 
@@ -455,6 +484,7 @@ class Issue(RegScaleModel):
         key: str,
         start_date: Union[str, datetime.datetime, None] = None,
         dt_format: Optional[str] = "%Y-%m-%dT%H:%M:%S",
+        default_days: Optional[int] = 60,
     ) -> str:
         """
         Function to return due date based on the severity of the issue; the values are in the init.yaml
@@ -465,6 +495,7 @@ class Issue(RegScaleModel):
         :param str key: The key to use for init.yaml from the issues section to determine the due date
         :param Union[str, datetime.datetime, None] start_date: The date to start the due date calculation from, defaults to current date
         :param Optional[str] dt_format: String of the date format to use, defaults to "%Y-%m-%dT%H:%M:%S"
+        :param Optional[int] default_days: Default number of days to return if no values match, defaults to 60
         :return: Due date for the issue
         :rtype: str
         """
@@ -479,19 +510,18 @@ class Issue(RegScaleModel):
         if isinstance(severity, IssueSeverity):
             severity = severity.value
         elif severity.lower() not in [severity.value.lower() for severity in IssueSeverity]:
-            severity = IssueSeverity.NotAssigned.value
+            severity = cls.assign_severity(severity)
 
         if severity == IssueSeverity.Critical.value:
-            days = cls._get_days_for_values(["critical"], config, key)
-            start_date = start_date + datetime.timedelta(days=days)
+            days = cls._get_days_for_values(["critical"], config, key, default_days)
         elif severity == IssueSeverity.High.value:
-            days = cls._get_days_for_values(["high"], config, key)
+            days = cls._get_days_for_values(["high"], config, key, default_days)
         elif severity == IssueSeverity.Moderate.value:
-            days = cls._get_days_for_values(["moderate", "medium"], config, key)
+            days = cls._get_days_for_values(["moderate", "medium"], config, key, default_days)
         elif severity == IssueSeverity.Low.value:
-            days = cls._get_days_for_values(["low", "minor"], config, key)
+            days = cls._get_days_for_values(["low", "minor"], config, key, default_days)
         else:
-            days = 60
+            days = default_days
         due_date = start_date + datetime.timedelta(days=days)
         return due_date.strftime(dt_format)
 
@@ -526,6 +556,7 @@ class Issue(RegScaleModel):
             "low": IssueSeverity.Low.value,
             "moderate": IssueSeverity.Moderate.value,
             "high": IssueSeverity.High.value,
+            "critical": IssueSeverity.Critical.value,
         }
         severity = IssueSeverity.NotAssigned.value
         # see if the value is an int or float
@@ -542,8 +573,10 @@ class Issue(RegScaleModel):
                 severity = severity_levels["low"]
             elif value.lower() in ["medium", "moderate", "major"]:
                 severity = severity_levels["moderate"]
-            elif value.lower() in ["high", "critical", "highest", "critical", "blocker"]:
+            elif value.lower() in ["high", "highest", "blocker"]:
                 severity = severity_levels["high"]
+            elif value.lower() in ["critical"]:
+                severity = severity_levels["critical"]
             elif value in list(severity_levels.values()):
                 severity = value
         return severity
@@ -914,103 +947,285 @@ class Issue(RegScaleModel):
         """
         import logging
 
-        cache_disabled = cls._is_cache_disabled()
-        use_cache: bool = not cache_disabled
-
         logger = logging.getLogger("regscale")
+
         # Check cache first
-        if use_cache:
-            cached_data = cls._get_from_cache(plan_id)
-            if cached_data is not None:
-                logger.info(f"Using cached open issues data for security plan {plan_id}")
-                return cached_data
-
-        # Performance optimization: Use larger batch size and optimize query
-        take = 50  # Increased from 50 to reduce API roundtrips
-        skip = 0
-        control_issues: Dict[int, List[OpenIssueDict]] = defaultdict(list)
+        cached_data = cls._check_cache(plan_id, logger)
+        if cached_data is not None:
+            return cached_data
 
-        start_time = time.time()
-        logger.info(f"Fetching open issues for controls and for security plan {plan_id}...")
+        # Fetch open issues from API
+        control_issues = cls._fetch_open_issues_from_api(plan_id, is_component, logger)
+
+        # Cache the results if caching is enabled
+        if not cls._is_cache_disabled():
+            cls._cache_data(plan_id, control_issues)
+
+        return control_issues
+
+    @classmethod
+    def _check_cache(cls, plan_id: int, logger) -> Optional[Dict[int, List[OpenIssueDict]]]:
+        """
+        Check cache for open issues data
+
+        :param int plan_id: The ID of the parent
+        :param logger: Logger instance
+        :return: Cached data if available and valid, None otherwise
+        :rtype: Optional[Dict[int, List[OpenIssueDict]]]
+        """
+        if cls._is_cache_disabled():
+            return None
+
+        cached_data = cls._get_from_cache(plan_id)
+        if cached_data is not None:
+            logger.info(f"Using cached open issues data for security plan {plan_id}")
+        return cached_data
+
+    @classmethod
+    def _fetch_open_issues_from_api(cls, plan_id: int, is_component: bool, logger) -> Dict[int, List[OpenIssueDict]]:
+        """
+        Fetch open issues from API with pagination
 
+        :param int plan_id: The ID of the parent
+        :param bool is_component: Whether parent is a component
+        :param logger: Logger instance
+        :return: Dictionary of control IDs to open issues
+        :rtype: Dict[int, List[OpenIssueDict]]
+        """
+        start_time = time.time()
         module_str = "component" if is_component else "security plan"
-        logger.info(
-            f"Fetching open issues for controls and for {module_str} {plan_id}...",
-        )
-        supports_multiple_controls: bool = cls.is_multiple_controls_supported()
+        logger.info(f"Fetching open issues for controls and for {module_str} {plan_id}...")
+
+        control_issues: Dict[int, List[OpenIssueDict]] = defaultdict(list)
+
+        try:
+            total_fetched = cls._paginate_and_process_issues(plan_id, is_component, control_issues, logger)
+            cls._log_completion(plan_id, total_fetched, len(control_issues), start_time, logger)
+        except Exception as e:
+            logger.error(f"Error fetching open issues for security plan {plan_id}: {e}")
+            return defaultdict(list)
+
+        return control_issues
+
+    @classmethod
+    def _paginate_and_process_issues(
+        cls,
+        plan_id: int,
+        is_component: bool,
+        control_issues: Dict[int, List[OpenIssueDict]],
+        logger,
+    ) -> int:
+        """
+        Paginate through API results and process issues using concurrent requests
+
+        :param int plan_id: The ID of the parent
+        :param bool is_component: Whether parent is a component
+        :param Dict[int, List[OpenIssueDict]] control_issues: Dictionary to populate with results
+        :param logger: Logger instance
+        :return: Total number of items fetched
+        :rtype: int
+        """
+        take = 50
+        supports_multiple_controls = cls.is_multiple_controls_supported()
+        fields = cls._get_query_fields(supports_multiple_controls)
+
+        # First fetch to get total count
+        query = cls._build_query(plan_id, is_component, 0, take, fields)
+        response = cls._get_api_handler().graph(query)
+
+        items = response.get(cls.get_module_string(), {}).get("items", [])
+        total_count = response.get(cls.get_module_string(), {}).get("totalCount", 0)
+
+        # Process first page
+        cls._process_issue_items(items, supports_multiple_controls, control_issues)
+        logger.info("Fetched first page with %d items, total: %d", len(items), total_count)
+
+        # Use concurrent pagination for remaining pages if there are more
+        if total_count > take:
+            logger.info("Fetching remaining pages concurrently...")
+            from regscale.core.utils.async_graphql_client import run_async_paginated_query
+
+            # Get API credentials
+            api_handler = cls._get_api_handler()
+            domain = api_handler.config.get("domain", "")
+            token = api_handler.config.get("token", "")
+            endpoint = f"{domain}/graphql"
+            headers = {"Authorization": f"{token}"}
+
+            # Create query builder
+            def query_builder(skip: int, page_size: int) -> str:
+                return cls._build_query(plan_id, is_component, skip, page_size, fields)
+
+            # Define token refresh callback
+            def token_refresh() -> str:
+                # Re-read token from config in case it was refreshed
+                return api_handler.config.get("token", "")
+
+            # Fetch remaining pages concurrently
+            remaining_items = run_async_paginated_query(
+                endpoint=endpoint,
+                headers=headers,
+                query_builder=query_builder,
+                topic_key=cls.get_module_string(),
+                total_count=total_count - take,  # Subtract first page
+                page_size=take,
+                starting_skip=take,  # Start from second page
+                max_concurrent=5,
+                timeout=60,
+                task_name="Open Issues",
+                token_refresh_callback=token_refresh,
+            )
+
+            # Process remaining items
+            for batch_items in [remaining_items[i : i + 100] for i in range(0, len(remaining_items), 100)]:
+                cls._process_issue_items(batch_items, supports_multiple_controls, control_issues)
+            logger.info("Fetched %d items from remaining pages", len(remaining_items))
+
+        return total_count
+
+    @classmethod
+    def _get_query_fields(cls, supports_multiple_controls: bool) -> str:
+        """
+        Get GraphQL query fields based on control support
 
-        # Optimize field selection - only get what we need
+        :param bool supports_multiple_controls: Whether multiple controls are supported
+        :return: GraphQL field selection string
+        :rtype: str
+        """
         if supports_multiple_controls:
-            fields = "id, otherIdentifier, integrationFindingId, controlImplementations { id }"
-        else:
-            fields = "id, controlId, otherIdentifier, integrationFindingId"
+            return "id, otherIdentifier, integrationFindingId, controlImplementations { id }"
+        return "id, controlId, otherIdentifier, integrationFindingId"
 
-        total_fetched = 0
+    @classmethod
+    def _build_query(cls, plan_id: int, is_component: bool, skip: int, take: int, fields: str) -> str:
+        """
+        Build GraphQL query for fetching open issues
 
-        try:
-            while True:
-                query = f"""
-                    query GetOpenIssuesByPlanOrComponent {{
-                        {cls.get_module_string()}(
-                            skip: {skip},
-                            take: {take},
-                            where: {{
-                                {"componentId" if is_component else "securityPlanId"}: {{eq: {plan_id}}},
-                                status: {{eq: "Open"}}
-                            }}
-                        ) {{
-                            items {{ {fields} }}
-                            pageInfo {{ hasNextPage }}
-                            totalCount
-                        }}
+        :param int plan_id: The ID of the parent
+        :param bool is_component: Whether parent is a component
+        :param int skip: Number of items to skip
+        :param int take: Number of items to take
+        :param str fields: GraphQL fields to select
+        :return: GraphQL query string
+        :rtype: str
+        """
+        parent_field = "componentId" if is_component else "securityPlanId"
+        return f"""
+            query GetOpenIssuesByPlanOrComponent {{
+                {cls.get_module_string()}(
+                    skip: {skip},
+                    take: {take},
+                    where: {{
+                        {parent_field}: {{eq: {plan_id}}},
+                        status: {{eq: "Open"}}
                     }}
-                """
+                ) {{
+                    items {{ {fields} }}
+                    pageInfo {{ hasNextPage }}
+                    totalCount
+                }}
+            }}
+        """
 
-                response = cls._get_api_handler().graph(query)
-                items = response.get(cls.get_module_string(), {}).get("items", [])
-                total_count = response.get(cls.get_module_string(), {}).get("totalCount", 0)
+    @classmethod
+    def _log_progress(cls, skip: int, take: int, items_count: int, total_count: int, logger) -> None:
+        """
+        Log progress for large datasets
 
-                # Log progress for large datasets
-                if total_count > 1000:
-                    logger.info(
-                        f"Processing batch {skip // take + 1} - fetched {len(items)} items ({skip + len(items)}/{total_count})"
-                    )
+        :param int skip: Number of items skipped
+        :param int take: Batch size
+        :param int items_count: Number of items in current batch
+        :param int total_count: Total count of items
+        :param logger: Logger instance
+        :rtype: None
+        """
+        if total_count > 1000:
+            logger.info(
+                f"Processing batch {skip // take + 1} - fetched {items_count} items ({skip + items_count}/{total_count})"
+            )
 
-                for item in items:
-                    issue_dict = OpenIssueDict(
-                        id=item["id"],
-                        otherIdentifier=item.get("otherIdentifier", ""),
-                        integrationFindingId=item.get("integrationFindingId", ""),
-                    )
+    @classmethod
+    def _process_issue_items(
+        cls,
+        items: List[Dict[str, Any]],
+        supports_multiple_controls: bool,
+        control_issues: Dict[int, List[OpenIssueDict]],
+    ) -> None:
+        """
+        Process issue items and populate control_issues dictionary
 
-                    if supports_multiple_controls and item.get("controlImplementations"):
-                        for control in item.get("controlImplementations", []):
-                            control_issues[control["id"]].append(issue_dict)
-                    elif item.get("controlId"):
-                        control_issues[item["controlId"]].append(issue_dict)
+        :param List[Dict[str, Any]] items: List of issue items from API
+        :param bool supports_multiple_controls: Whether multiple controls are supported
+        :param Dict[int, List[OpenIssueDict]] control_issues: Dictionary to populate
+        :rtype: None
+        """
+        for item in items:
+            issue_dict = OpenIssueDict(
+                id=item["id"],
+                otherIdentifier=item.get("otherIdentifier", ""),
+                integrationFindingId=item.get("integrationFindingId", ""),
+            )
 
-                total_fetched += len(items)
-                if not response.get(cls.get_module_string(), {}).get("pageInfo", {}).get("hasNextPage", False):
-                    break
-                skip += take
+            if supports_multiple_controls:
+                cls._add_issue_to_multiple_controls(item, issue_dict, control_issues)
+            else:
+                cls._add_issue_to_single_control(item, issue_dict, control_issues)
 
-        except Exception as e:
-            logger.error(f"Error fetching open issues for security plan {plan_id}: {e}")
-            # Return empty dict on error to prevent breaking the calling code
-            return defaultdict(list)
+    @classmethod
+    def _add_issue_to_multiple_controls(
+        cls,
+        item: Dict[str, Any],
+        issue_dict: OpenIssueDict,
+        control_issues: Dict[int, List[OpenIssueDict]],
+    ) -> None:
+        """
+        Add issue to multiple control implementations
 
+        :param Dict[str, Any] item: Issue item from API
+        :param OpenIssueDict issue_dict: Issue dictionary
+        :param Dict[int, List[OpenIssueDict]] control_issues: Dictionary to populate
+        :rtype: None
+        """
+        if item.get("controlImplementations"):
+            for control in item.get("controlImplementations", []):
+                control_issues[control["id"]].append(issue_dict)
+
+    @classmethod
+    def _add_issue_to_single_control(
+        cls,
+        item: Dict[str, Any],
+        issue_dict: OpenIssueDict,
+        control_issues: Dict[int, List[OpenIssueDict]],
+    ) -> None:
+        """
+        Add issue to single control
+
+        :param Dict[str, Any] item: Issue item from API
+        :param OpenIssueDict issue_dict: Issue dictionary
+        :param Dict[int, List[OpenIssueDict]] control_issues: Dictionary to populate
+        :rtype: None
+        """
+        if item.get("controlId"):
+            control_issues[item["controlId"]].append(issue_dict)
+
+    @classmethod
+    def _log_completion(cls, plan_id: int, total_fetched: int, control_count: int, start_time: float, logger) -> None:
+        """
+        Log completion statistics
+
+        :param int plan_id: The ID of the parent
+        :param int total_fetched: Total number of items fetched
+        :param int control_count: Number of controls with issues
+        :param float start_time: Start time of the operation
+        :param logger: Logger instance
+        :rtype: None
+        """
         elapsed_time = time.time() - start_time
         logger.info(
-            f"Finished fetching {total_fetched} open issue(s) for {len(control_issues)} control(s) "
+            f"Finished fetching {total_fetched} open issue(s) for {control_count} control(s) "
             f"in security plan {plan_id} - took {elapsed_time:.2f} seconds"
         )
 
-        # Cache the results
-        if use_cache:
-            cls._cache_data(plan_id, control_issues)
-
-        return control_issues
-
     @classmethod
     def get_sort_position_dict(cls) -> Dict[str, int]:
         """
@@ -1100,36 +1315,20 @@ class Issue(RegScaleModel):
         }
 
     @classmethod
-    def get_enum_values(cls, field_name: str) -> List[Union[IssueSeverity, IssueStatus, str]]:
+    def get_enum_values(cls, field_name: str) -> List[Union[IssueSeverity, IssueStatus, IssueIdentification, str]]:
         """
         Overrides the base method.
 
         :param str field_name: The property name to provide enum values for
         :return: List of enum values or strings
-        :rtype: List[Union[IssueSeverity, IssueStatus, str]]
+        :rtype: List[Union[IssueSeverity, IssueStatus, IssueIdentification, str]]
         """
         if field_name == "severityLevel":
             return [severity.__str__() for severity in IssueSeverity]
         if field_name == "status":
             return [status.__str__() for status in IssueStatus]
         if field_name == "identification":
-            return [
-                "A-123 Review",
-                "Assessment/Audit (External)",
-                "Assessment/Audit (Internal)",
-                "Critical Control Review",
-                "FDCC/USGCB",
-                "GAO Audit",
-                "IG Audit",
-                "Incidnet Response Lessons Learned",
-                "ITAR",
-                "Other",
-                "Penetration Test",
-                "Risk Assessment",
-                "Security Authorization",
-                "Security Control Assessment",
-                "Vulnerability Assessment",
-            ]
+            return [identification.__str__() for identification in IssueIdentification]
         return cls.get_bool_enums(field_name)
 
     @classmethod

regscale/models/regscale_models/milestone.py

@@ -3,7 +3,8 @@
 """Class for milestone model in RegScale platform"""
 
 from typing import Optional
-from pydantic import ConfigDict, Field
+
+from pydantic import Field, field_validator, model_validator
 
 from regscale.core.app.utils.app_utils import get_current_datetime
 from regscale.models.regscale_models.regscale_model import RegScaleModel
@@ -14,27 +15,45 @@ class Milestone(RegScaleModel):
 
     _module_slug = "milestones"
     _module_string = "milestones"
+    _unique_fields = ["title", "parentModule", "parentID"]
+    _parent_id_field = "parentID"
 
     title: str
     id: int = 0
     isPublic: Optional[bool] = True
     milestoneDate: Optional[str] = Field(default_factory=get_current_datetime)
     responsiblePersonId: Optional[str] = None
-    predecessorStepId: Optional[int] = 0
+    predecessorStepId: Optional[int] = None
     completed: Optional[bool] = False
-    dateCompleted: Optional[str] = Field(default_factory=get_current_datetime)
+    dateCompleted: Optional[str] = None
     notes: Optional[str] = ""
     parentID: Optional[int] = None
     parentModule: str = ""
 
-    @staticmethod
-    def _get_additional_endpoints() -> ConfigDict:
-        """
-        Get additional endpoints for the Milestone model
-
-        :return: A dictionary of additional endpoints
-        :rtype: ConfigDict
-        """
-        return ConfigDict(
-            get_all_by_parent="/api/{model_slug}/getAllByParent/{intParentID}",
-        )
+    @field_validator("milestoneDate")
+    @classmethod
+    def validate_milestone_date(cls, v: Optional[str]) -> str:
+        """Ensure milestoneDate is never empty."""
+        if not v or v == "":
+            return get_current_datetime()
+        return v
+
+    @field_validator("dateCompleted")
+    @classmethod
+    def set_date_completed(cls, v: Optional[str], info) -> Optional[str]:
+        """Set dateCompleted based on completed field."""
+        completed = info.data.get("completed", False)
+        if completed and (v is None or v == ""):
+            return get_current_datetime()
+        if not completed:
+            return None
+        return v
+
+    @model_validator(mode="after")
+    def validate_completion(self):
+        """Ensure dateCompleted is set when completed is True."""
+        if self.completed and (self.dateCompleted is None or self.dateCompleted == ""):
+            self.dateCompleted = get_current_datetime()
+        elif not self.completed:
+            self.dateCompleted = None
+        return self

regscale/models/regscale_models/organization.py

@@ -15,7 +15,10 @@ class Organization(RegScaleModel):
     id: int = 0
     name: Optional[str] = ""
     description: Optional[str] = ""
+    orgCode: Optional[str] = ""
+    orgUrl: Optional[str] = ""
     status: Optional[str] = "Active"
+    externalId: Optional[str] = ""
 
     @staticmethod
     def _get_additional_endpoints() -> dict: