regscale-cli 6.16.0.0__py3-none-any.whl

regscale/models/app_models/import_validater.py

@@ -0,0 +1,321 @@
+"""Data model that will be used to validate the input data and that it has the required headers before proceeding."""
+
+import logging
+from typing import TYPE_CHECKING, Optional, Union
+
+if TYPE_CHECKING:
+    import pandas as pd
+    from regscale.models import Mapping
+
+import json
+from pathlib import Path
+
+from regscale.exceptions.validation_exception import ValidationException
+
+logger = logging.getLogger("regscale")
+JSON = ".json"
+XLSX = ".xlsx"
+XLSM = ".xlsm"
+CSV = ".csv"
+XML = ".xml"
+NESSUS = ".nessus"
+
+
+class ImportValidater:
+    """
+    Data importer that validates the input data
+
+    :param list required_headers: List of required headers (columns)
+    :param Union[str, Path] file_path: Path to the file to import
+    :raises ValidationException: If the file type is not supported
+    """
+
+    _supported_types = [CSV, XLSX, XLSM, JSON, XML, NESSUS]
+    xml_tag = "record"
+    keys: Optional[list] = []
+    key: Optional[str] = ""
+    required_headers: list
+    parsed_headers: list
+    data: Union["pd.DataFrame", list]
+    mapping: "Mapping"
+    worksheet_name: Optional[str] = None
+    skip_rows: Optional[int] = None
+    file_path: Path
+    file_type: str
+    mapping_file_path: Path
+    disable_mapping: bool
+
+    def __init__(
+        self,
+        required_headers: list,
+        file_path: Union[str, Path],
+        mapping_file_path: Union[str, Path],
+        disable_mapping: bool = False,
+        xml_tag: Optional[str] = None,
+        key: Optional[str] = None,
+        keys: Optional[list] = None,
+        worksheet_name: Optional[str] = None,
+        skip_rows: Optional[int] = None,
+        prompt: bool = True,
+        ignore_unnamed: bool = False,
+    ):
+        self.ignore_unnamed = ignore_unnamed
+        self.prompt = prompt
+        self.required_headers = required_headers
+        file_path = self._convert_str_to_path(file_path)
+        self.mapping_file_path = self._convert_str_to_path(mapping_file_path)
+        self.file_path = file_path
+        self.file_type = file_path.suffix
+        if not self.mapping_file_path.suffix:
+            self.mapping_file_path = Path(self.mapping_file_path / f"{self.file_type[1:]}_mapping.json")
+        self.disable_mapping = disable_mapping
+        self.xml_tag = xml_tag
+        self.key = key
+        self.keys = keys
+        self.worksheet_name = worksheet_name
+        self.skip_rows = skip_rows
+        if self.file_type not in self._supported_types:
+            raise ValidationException(
+                f"Unsupported file type: {self.file_type}, supported types are: {', '.join(self._supported_types)}",
+            )
+        self._parse_headers()
+
+    @staticmethod
+    def _convert_str_to_path(file_path: Union[str, Path]) -> Path:
+        """
+        Converts a string to a Path object
+
+        :param Union[str, Path] file_path: File path as a string
+        :returns: Path object
+        :rtype: Path
+        """
+        if isinstance(file_path, Path):
+            return file_path
+        return Path(file_path)
+
+    def import_data(self) -> Union["pd.DataFrame", list, dict]:
+        """
+        Imports the data from the file and returns it
+
+        :returns: Imported data
+        :rtype: Union[pd.DataFrame, list, dict]
+        """
+        if self.file_type == CSV:
+            return self.import_csv(self.file_path)
+        elif self.file_type == XLSX or self.file_type == XLSM:
+            return self.import_xlsx(self.file_path)
+        elif self.file_type == JSON:
+            return self.import_json(self.file_path)
+        elif self.file_type in (XML, NESSUS):
+            return self.import_xml(self.file_path, self.xml_tag)
+
+    def _parse_headers(self):
+        """
+        Parses the headers from the file based on the file type and returns them
+
+        :returns: List of headers
+        :rtype: list
+        """
+        self.data = self.import_data()
+        if self.file_type in [CSV, XLSX, XLSM]:
+            self.parsed_headers = self.data.columns  # type: ignore
+        elif self.file_type == XML:
+            self.parsed_headers = list(self.data.keys())
+        elif self.file_type == JSON and not self.parsed_headers:
+            raise ValidationException(f"Unable to parse headers from JSON file: {self.file_path}")
+
+    def validate_headers(self, headers: Union[list, "pd.Index"]) -> None:
+        """
+        Validates that all required headers are present
+
+        :param Union[list, pd.Index] headers: List of headers from the file
+        """
+        import re
+
+        from regscale.models import Mapping
+
+        if not self.ignore_unnamed and any(re.search(r"unnamed", header, re.IGNORECASE) for header in headers):  # type: ignore
+            raise ValidationException(
+                f"Unable to parse headers from the file. Please ensure the headers are named in {self.file_path}"
+                f"\n{', '.join(headers)}"
+            )
+
+        if not self.prompt:
+            # Let's not prompt the user to find the header, and just raise a validation exception
+            missing_headers = [header for header in self.required_headers if header not in headers]
+            extra_headers = [header for header in headers if header not in self.required_headers]
+            if missing_headers:
+                raise ValidationException(
+                    f"{', '.join([f'`{header}`' for header in missing_headers])} header(s) not found in {self.file_path}"
+                )
+            if extra_headers:
+                logger.warning("Extra headers found in the file: %s", ", ".join(extra_headers))
+
+        if self.disable_mapping:
+            logger.debug("Mapping is disabled, using headers as is.")
+            self.mapping = Mapping(
+                mapping={header: header for header in headers},
+                expected_field_names=self.required_headers,
+                file_path_for_prompt=self.file_path,
+            )
+        else:
+            logger.debug(f"Mapping is enabled, validating headers for {self.file_path}")
+            self.mapping = Mapping.from_file(
+                file_path=self.mapping_file_path,
+                expected_field_names=self.required_headers,
+                mapping={"mapping": {header: header for header in headers}},
+                parsed_headers=[header for header in headers],  # this converts the pd.Index to a list
+                file_path_for_prompt=self.file_path,
+            )
+
+    def import_csv(self, file_path: Union[str, Path]) -> "pd.DataFrame":
+        """
+        Imports a CSV file and validates headers
+
+        :param file_path: Path to the CSV file
+        :returns: DataFrame with the imported data
+        :rtype: pd.DataFrame
+        """
+        import pandas
+
+        try:
+            if self.skip_rows:
+                df = pandas.read_csv(file_path, skiprows=self.skip_rows - 1, on_bad_lines="warn")
+            else:
+                df = pandas.read_csv(file_path, on_bad_lines="warn")
+        except pandas.errors.ParserError:
+            raise ValidationException(f"Unable to parse the {CSV} file: {file_path}")
+        self.validate_headers(df.columns)
+        df = df.fillna("")
+        return df
+
+    def import_xlsx(self, file_path: Union[str, Path]) -> "pd.DataFrame":
+        """
+        Imports an XLSX file and validates headers
+
+        :param Union[str, Path] file_path: Path to the XLSX file
+        :returns: DataFrame with the imported data
+        :rtype: pd.DataFrame
+        """
+        import pandas
+
+        try:
+            if self.skip_rows and self.worksheet_name:
+                df = pandas.read_excel(file_path, sheet_name=self.worksheet_name, skiprows=self.skip_rows - 1)
+            elif self.worksheet_name:
+                df = pandas.read_excel(file_path, sheet_name=self.worksheet_name)
+            elif self.skip_rows:
+                df = pandas.read_excel(file_path, skiprows=self.skip_rows)
+            else:
+                df = pandas.read_excel(file_path)
+        except Exception as e:
+            raise ValidationException(f"Unable to parse the {XLSX} file: {file_path}\nDetails: {e}")
+        self.validate_headers(df.columns)
+        df = df.fillna("")
+        return df
+
+    def _handle_keys(self, data: dict) -> list:
+        """
+        Handles the keys for JSON data
+
+        :param dict data: JSON data
+
+        :returns: List of keys
+        :rtype: list
+        """
+        if self.keys:
+            value = data
+            # iterate each key and see if it is in the data
+            for key in self.keys:
+                value = value.get(key, {})
+            if isinstance(value, dict):
+                return list(value.keys())
+            elif isinstance(value, list):
+                if value:
+                    return list(value[0].keys())
+                else:
+                    raise ValidationException(
+                        f"JSON file must contain a key '{self.keys[-1]}' with a list of dictionaries"
+                    )
+
+        return data.get(self.key) or data.get(self.keys[0])
+
+    def import_json(self, file_path: Union[str, Path]) -> list:
+        """
+        Imports a JSON file and validates keys (treated as headers)
+
+        :param file_path: Path to the JSON file
+        :raises ValidationException: If the JSON file is empty or not a list of dictionaries
+        :returns: List of dictionaries with the imported data
+        :rtype: list
+        """
+        with open(file_path, "r", encoding="utf-8") as f:
+            data = json.load(f)
+
+        if isinstance(data, list) and not self.key and not self.keys and data:
+            # Assuming JSON is a list of dictionaries
+            self.validate_headers(data[0].keys())
+            self.parsed_headers = list(data[0].keys())
+        elif isinstance(data, list) and self.keys and data:
+            # Assuming JSON is a list of dictionaries
+            self.validate_headers(self._handle_keys(data[0]))
+            self.parsed_headers = list(self._handle_keys(data[0]))
+        elif isinstance(data, dict) and (self.key or self.keys):
+            if self.keys:
+                self.validate_headers(self._handle_keys(data))
+                self.parsed_headers = self._handle_keys(data)
+            elif findings := data.get(self.key):
+                self.validate_headers(findings[0].keys())
+                self.parsed_headers = list(findings[0].keys())
+        elif isinstance(data, dict):
+            self.validate_headers(list(data.keys()))
+            self.parsed_headers = list(data.keys())
+        else:
+            raise ValidationException(f"Unable to parse headers from JSON file: {file_path}")
+
+        return data
+
+    def _remove_at_prefix(self, xml_data: Union[list, dict]) -> Union[list, dict]:
+        """
+        Recursively remove the '@' prefix from keys in a dictionary
+
+        :param Union[list, dict] xml_data: Parsed XML data, either a dictionary or a list of dictionaries
+        :returns: Dictionary with '@' prefix removed or list of dictionaries with '@' prefix removed
+        :rtype: Union[list, dict]
+        """
+        if isinstance(xml_data, dict):
+            return {k.lstrip("@"): self._remove_at_prefix(v) for k, v in xml_data.items()}
+        elif isinstance(xml_data, list):
+            return [self._remove_at_prefix(i) for i in xml_data]
+        else:
+            return xml_data
+
+    def import_xml(self, file_path: Union[str, Path], record_tag: Optional[str] = None) -> dict:
+        """
+        Imports an XML file and validates keys (treated as headers)
+
+        :param Union[str, Path] file_path: Path to the XML file
+        :param str record_tag: The XML tag that represents a record, defaults to "record"
+        :raises ValidationException: If the XML file contains no records
+        :returns: Dictionary with the imported data
+        :rtype: dict
+        """
+        import xmltodict
+
+        if isinstance(file_path, str):
+            file_path = Path(file_path)
+        with open(file_path, "r", encoding="utf-8") as file:
+            xml_content = file.read()
+        try:
+            if dict_content := xmltodict.parse(xml_content):
+                if record_tag:
+                    dict_content = dict_content.get(record_tag)
+                dict_content = self._remove_at_prefix(dict_content)
+                self.validate_headers(list(dict_content.keys()))
+                return dict_content
+            else:
+                raise ValidationException("XML file contains no records.")
+        except xmltodict.expat.ExpatError as e:
+            raise ValidationException(
+                f"Error parsing {self.file_type.strip('.')} file: {file_path}.\nDetails: {e.args[0]}"
+            ) from e

regscale/models/app_models/mapping.py

@@ -0,0 +1,260 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""Pydantic class for custom container scan mappings"""
+import json
+import logging
+from json import JSONDecodeError
+from logging import Logger
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Type, Union
+
+from pydantic import BaseModel, ConfigDict, field_validator
+
+from regscale.core.app import create_logger
+from regscale.core.app.utils.app_utils import check_file_path, error_and_exit, save_data_to
+from regscale.exceptions import ValidationException
+
+
+class Mapping(BaseModel):
+    """
+    Pydantic class for custom container scan mappings
+    """
+
+    model_config = ConfigDict(populate_by_name=True, json_encoders={bytes: lambda v: v.decode()})
+
+    mapping: Dict[str, str]
+    # expected field names for validation
+    expected_field_names: List[str] = []
+    _logger: Optional[Logger] = logging.getLogger("regscale")
+    file_path_for_prompt: Optional[Path] = None
+
+    @classmethod
+    def _prompt_for_field(
+        cls,
+        field: str,
+        mapping: Union[dict, "Mapping"],
+        parsed_headers: Optional[List[str]] = None,
+        file_path_for_prompt: Optional[Path] = None,
+    ) -> str:
+        """
+        Prompt for a field value
+
+        :param str field: Field to prompt for
+        :param Union[dict, "Mapping"] mapping: Mapping object or dictionary of the mapping object
+        :param Optional[List[str]] parsed_headers: Parsed headers from file, defaults to None
+        :return: Field value
+        :rtype: str
+        """
+        info_msg = f"Field '{field}' not found in mapping, please enter the mapping for this field, enter 'exit' to exit, or enter 'skip' to skip"
+        if file_path_for_prompt:
+            info_msg += f"\nFile path for headers: {file_path_for_prompt.name}"
+        if parsed_headers:
+            info_msg += f"\nParsed headers: {parsed_headers}"
+        cls._logger.default.info(info_msg)
+        custom_mapping = input(f"Enter the mapping for field '{field}':")
+        cls._handle_user_input(custom_mapping, field, mapping, parsed_headers, file_path_for_prompt)
+
+    @classmethod
+    def _handle_user_input(
+        cls,
+        custom_mapping: str,
+        field: str,
+        mapping: Union[dict, "Mapping"],
+        parsed_headers: Optional[List[str]],
+        file_path_for_prompt: Optional[Path],
+    ) -> None:
+        if custom_mapping.lower() in ["exit", "skip"]:
+            cls._handle_exit_or_skip(custom_mapping, file_path_for_prompt)
+        else:
+            cls._confirm_and_update_mapping(custom_mapping, field, mapping, parsed_headers)
+
+    @classmethod
+    def _handle_exit_or_skip(cls, custom_mapping: str, file_path_for_prompt: Optional[Path]) -> None:
+        if custom_mapping.lower() == "exit":
+            error_and_exit("Exiting...")
+        elif custom_mapping.lower() == "skip":
+            raise ValidationException(f"Skipping file {file_path_for_prompt.name if file_path_for_prompt else ''}...")
+
+    @classmethod
+    def _confirm_and_update_mapping(
+        cls, custom_mapping: str, field: str, mapping: Union[dict, "Mapping"], parsed_headers: Optional[List[str]]
+    ) -> None:
+        confirm = input(f"Confirm the mapping for field '{field}' is '{custom_mapping}' (y/n)")
+        if confirm.lower() == "y":
+            cls._validate_and_update_mapping(custom_mapping, field, mapping, parsed_headers)
+        elif confirm.lower() in ["exit", "skip"]:
+            cls._handle_exit_or_skip(confirm, None)
+        else:
+            cls._prompt_for_field(field=field, mapping=mapping, parsed_headers=parsed_headers)
+
+    @classmethod
+    def _validate_and_update_mapping(
+        cls, custom_mapping: str, field: str, mapping: Union[dict, "Mapping"], parsed_headers: Optional[List[str]]
+    ) -> None:
+        if parsed_headers and custom_mapping not in parsed_headers:
+            cls._logger.default.warning(
+                f"Mapping {custom_mapping} is not found in the headers. "
+                f"Please select one of the following headers: {parsed_headers}."
+            )
+            cls._prompt_for_field(field=field, mapping=mapping, parsed_headers=parsed_headers)
+        elif custom_mapping not in mapping.values() and not parsed_headers:
+            cls._logger.default.warning(
+                f"Mapping {custom_mapping} is not found in the headers. "
+                f"Please select one of the following headers: {list(mapping.values())}."
+            )
+            cls._prompt_for_field(field=field, mapping=mapping, parsed_headers=parsed_headers)
+        else:
+            if isinstance(mapping, dict):
+                mapping[field] = custom_mapping
+            else:
+                mapping.mapping[field] = custom_mapping
+
+    @field_validator("expected_field_names")
+    def validate_mapping(cls: Type["Mapping"], expected_field_names: List[str], values: Dict[str, Any]) -> List[str]:
+        """
+        Validate the expected field names
+
+        :param List[str] expected_field_names: Expected field names
+        :param Dict[str, Any] values: Values
+        :return: Expected field names
+        """
+        mapping = values.data.get("mapping")
+        if mapping is not None and expected_field_names is not None:
+            if missing_fields := [field for field in expected_field_names if field not in mapping]:
+                for field in missing_fields:
+                    cls._prompt_for_field(field, mapping, cls.file_path_for_prompt)
+        return expected_field_names
+
+    @field_validator("expected_field_names")
+    def validate_expected_field_names(cls: Type["Mapping"], expected_field_names: Any) -> List[str]:
+        """
+        Validate the expected field names and types
+
+        :param Any expected_field_names: Expected field names
+        :raises ValidationError: If expected_field_names is not a list or if any element in the list is not a string
+        :rtype: List[str]
+        :return: Expected field names
+        """
+        if not isinstance(expected_field_names, list):
+            raise ValidationException("expected_field_names must be a list")
+        if not all(isinstance(field_name, str) for field_name in expected_field_names):
+            raise ValidationException("All elements in expected_field_names must be strings")
+        return expected_field_names
+
+    # Add a from file method to load the mapping from a JSON file
+    @classmethod
+    def from_file(cls, file_path: Path, expected_field_names: List[str], **kwargs) -> "Mapping":
+        """
+        Load the mapping from a JSON file
+
+        :param Path file_path: Path to the JSON file
+        :param List[str] expected_field_names: Expected field names
+        :return: Validated Mapping ob
+        :rtype: Mapping
+        """
+        file_path_for_prompt = kwargs.get("file_path_for_prompt")
+        if not file_path.exists() and kwargs.get("mapping"):
+            check_file_path(file_path.parent)
+            dat = kwargs.get("mapping")
+        else:
+            with open(file_path, "r") as file:
+                try:
+                    dat = json.load(file)
+                    # if mapping is not found in the JSON file, check the kwargs for the provided mapping and use that
+                    if not dat.get("mapping"):
+                        dat["mapping"] = kwargs.get("mapping")
+                        if not dat.get("mapping"):
+                            error_and_exit("Mapping not found in JSON file")
+                except JSONDecodeError as jex:
+                    cls._logger.default.debug(jex)
+                    error_and_exit("JSON file is badly formatted, please check the file")
+                except (ValueError, SyntaxError) as exc:
+                    error_and_exit(f"Error parsing JSON file: {exc}")
+        if parsed_headers := kwargs.get("parsed_headers"):
+            cls._verify_parsed_headers(parsed_headers, expected_field_names, dat, file_path_for_prompt)
+        mapping = cls(
+            mapping=dat["mapping"],
+            expected_field_names=expected_field_names,
+            file_path_for_prompt=file_path_for_prompt,
+        )
+        if not mapping:
+            error_and_exit("Error loading mapping from file. Exiting...")
+        mapping.save_mapping(file_path)
+        return mapping
+
+    @classmethod
+    def _verify_parsed_headers(
+        cls, parsed_headers: List[str], expected_field_names: List[str], dat: dict, file_path_for_prompt: Optional[Path]
+    ) -> None:
+        """
+        Verify the parsed headers and prompt for missing headers
+
+        :param List[str] parsed_headers: Parsed headers
+        :param List[str] expected_field_names: Expected field names
+        :param Dict[str, Any] dat: Data dictionary
+        :param Optional[Path] file_path_for_prompt: File path for prompt
+        :rtype: None
+        """
+        if missing_headers := [
+            header
+            for header in expected_field_names
+            if header not in parsed_headers and dat["mapping"].get(header) not in parsed_headers
+        ]:
+            for header in missing_headers:
+                cls._logger.default.info(
+                    f"Header '{header}' not found in parsed headers in {file_path_for_prompt.name if file_path_for_prompt else 'file'} Please enter the mapping."
+                )
+                cls._prompt_for_field(field=header, mapping=dat["mapping"], parsed_headers=parsed_headers)
+
+    def save_mapping(self, file_path: Path) -> None:
+        """
+        Save the mapping to a JSON file
+
+        :param Path file_path: Path to save the mapping JSON file
+        :rtype: None
+        """
+        check_file_path(file_path.parent)
+        save_data_to(file_path, {"mapping": self.mapping})
+
+    def get_header(self, key: str) -> str:
+        """
+        Get the header for a key
+
+        :param str key: Key to get the header for
+        :return: Header for the key
+        :rtype: str
+        """
+        return self.mapping.get(key)
+
+    def get_value(self, dat: Optional[dict], key: str, default_val: Optional[Any] = "", warnings: bool = True) -> Any:
+        """
+        Get the value from a dictionary by mapped key
+
+        :param Optional[dict] dat: Data dictionary, defaults to None
+        :param str key: Key to get the value for
+        :param Optional[Any] default_val: Default value to return, defaults to empty string
+        :param bool warnings: Whether to log warnings, defaults to False
+        :return: Value for the key
+        :rtype: Any
+        """
+        # check mapping
+        if key == "None" or key is None:
+            return default_val
+        mapped_key = self.mapping.get(key)
+        if not mapped_key and warnings:
+            self._logger.warning(f"Value for key '{key}' not found in mapping.")
+        if dat and mapped_key:
+            val = dat.get(mapped_key)
+            if isinstance(val, str):
+                return val.strip()
+            return val or default_val
+        return default_val
+
+    def to_header(self) -> list[str]:
+        """
+        Convert the mapping to a header
+        :return: Mapping as a header
+        :rtype: list[str]
+        """
+        # convert mapping to a list of strings
+        return [f"{value}" for key, value in self.mapping.items()]

regscale/models/app_models/pipeline.py

@@ -0,0 +1,37 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""Dataclass for a user's pipeline"""
+
+# standard python imports
+from dataclasses import dataclass
+from typing import Any
+
+
+@dataclass
+class Pipeline:
+    """Pipeline Model"""
+
+    email: str  # Required
+    fullName: str = None
+    pipelines: Any = None
+    totalTasks: int = None
+    analyzed: bool = False
+    emailed: bool = False
+
+    def __getitem__(self, key: Any) -> Any:
+        """
+        Get attribute from Pipeline
+        :param Any key:
+        :return: value of provided key
+        :rtype: Any
+        """
+        return getattr(self, key)
+
+    def __setitem__(self, key: Any, value: Any) -> None:
+        """
+        Set attribute in Pipeline with provided key
+        :param Any key: Key to change to provided value
+        :param Any value: New value for provided Key
+        :rtype: None
+        """
+        return setattr(self, key, value)