"""
image_vul_parser
================

This module provides a class to parse the raw JSON report which is generated by Trivy.

Author:
    Pengfei liu

Date:
    2025-02-25
"""

# Standard library imports
import json
from typing import List

# Third party imports
import pandas as pd
from pandas import DataFrame

# sherlock imports
from iv_sherlock.conf.constant import (
    CVSS_V3_SCORE,
    CVSS_V3_VECTOR,
    DESCRIPTION,
    FIXED_VER,
    IMG_CREATION_TIME,
    IMG_ID,
    IMG_NAME,
    INSTALLED_VER,
    LIBRARY,
    OS_ARCHI,
    OS_FAMILY,
    OS_NAME,
    REPO_TAGS,
    SEVERITY,
    STATUS,
    TARGET_CLASS,
    TARGET_NAME,
    TARGET_TYPE,
    TARGET_VULNERABILITY_DF,
    VULNERABILITY_ID,
    VULNERABILITY_URL,
)


def print_keys_as_tree(in_dict: dict, indent=0):
    """
    Recursively prints the keys of a nested dictionary as a tree structure
    :param in_dict: The input dictionary
    :param indent: the number of spaces to indent each level
    :return:
    """
    for key, value in in_dict.items():
        print("  " * indent + str(key))  # Print the key with indentation
        if isinstance(value, dict):
            # If the value is a dictionary, recurse into it
            print_keys_as_tree(value, indent + 1)


class ImageVulParser:
    """
    This class represents an image_vulnerability_parser that parses the raw image vulnerability
    report that is generated by trivy in JSON format.

    :param raw_report_path:  The file path which contains the raw image vulnerability report.
    :type raw_report_path: Str

    :ivar raw_data: The dict which corresponds to the json file.
    :vartype raw_data: dict
    """

    def __init__(self, raw_report_path: str):
        """
        Initialize the image vulnerability parser
        :param raw_report_path: The JSON file path of the raw image vulnerability report.
        """
        with open(raw_report_path, "r", encoding="utf-8") as file:
            self.raw_data:dict = json.load(file)

    def get_raw_data(self) -> dict:
        """
        This method returns the raw dictionary of the trivy report.
        :return:
        """
        return self.raw_data

    def get_available_attributes(self) -> list:
        """
        This method returns the available attributes of the raw trivy report.
        :return: A list of available attributes.
        """
        return list(self.raw_data.keys())

    def build_image_vul_report(self) -> tuple[dict, list]:
        """
        This function builds an image vulnerability analysis summary report based on the raw data. It
        contains the basic image information and the list of image CVE
        :return:
        """
        img_basic_info = self.get_image_basic_info()
        cve_list = self.parser_cve_results()
        return img_basic_info, cve_list

    def parser_cve_results(self) -> List[dict]:
        """
        This function parses the Results section of the raw report. The results section contains a list of
        `Target`. A Target has the below attributes:
        - target_name: for example, redis (debian 12.7)
        - target_class: os-pkgs
        - target_type: debian
        - vulnerabilities: a list of dict which contains the details of each vulnerability
        For each target, we generate a parsed_target which contains:
        - target_name: for example, redis (debian 12.7)
        - target_class: os-pkgs
        - target_type: debian
        - target_vul_df: a pandas data frame to represent vulnerabilities(each row is a CVE)
        :return:
        """
        parsed_targets = []
        targets = self.raw_data.get("Results")
        # if the target is not empty, parse each target
        if targets:
            for target in targets:
                parsed_target = {}
                target_name = target.get("Target")
                target_class = target.get("Class")
                target_type = target.get("Type")
                vulnerabilities = target.get("Vulnerabilities")
                # build a pandas data frame to represent vulnerabilities(each row is a CVE)
                vul_pdf = self.parser_target_vulnerabilities(vulnerabilities)
                parsed_target[TARGET_NAME] = target_name
                parsed_target[TARGET_CLASS] = target_class
                parsed_target[TARGET_TYPE] = target_type
                parsed_target[TARGET_VULNERABILITY_DF] = vul_pdf
                parsed_targets.append(parsed_target)
        return parsed_targets

    @staticmethod
    def parser_target_vulnerabilities(vulnerabilities: List[dict]) -> DataFrame | None:
        """
        This function takes a list of vulnerabilities (raw.Results.Targets[i].Vulnerabilities), then parses them
        into a dataframe. It contains below columns:
        - library: name of the package which contains vulnerabilities
        - vulnerability_id: The id of the vulnerability(cve)
        - severity: The severity of the vulnerability(cve)
        - status: The status of the vulnerability(cve)
        - installed_version: The installed version of the package.
        - fixed_version: If a fix exists, in which version the fix will be applied
        - description: Description of the vulnerability(cve)
        - cvss_v3_score: The cvss score of the vulnerability(cve). Various organizations provide this score,
                 in this implementation, we took the first available, for some cve, this column maybe null.
        - cvss_v3_vector: The cvss vector of the vulnerability(cve)
        :param vulnerabilities:
        :return:
        """
        vul_list: list = []
        if vulnerabilities:
            for vul in vulnerabilities:
                vul_attrs = {}
                library = vul.get("PkgName")
                vul_id = vul.get("VulnerabilityID")
                vul_url = vul.get("PrimaryURL")
                severity = vul.get("Severity")
                status = vul.get("Status")
                installed_version = vul.get("InstalledVersion")
                fixed_version = vul.get("FixedVersion")
                description = vul.get("Title")
                # cvss_v3_score = vul.get("CVSS").get("nvd").get("V3Score")
                # cvss_v3_vector = vul.get("CVSS").get("nvd").get("V3Vector")
                cvss_v3_score, cvss_v3_vector = ImageVulParser.parse_cvss(
                    vul.get("CVSS")
                )
                vul_attrs[LIBRARY] = library
                vul_attrs[VULNERABILITY_ID] = vul_id
                vul_attrs[VULNERABILITY_URL] = vul_url
                vul_attrs[SEVERITY] = severity
                vul_attrs[STATUS] = status
                vul_attrs[INSTALLED_VER] = installed_version
                vul_attrs[FIXED_VER] = fixed_version
                vul_attrs[CVSS_V3_SCORE] = cvss_v3_score
                vul_attrs[CVSS_V3_VECTOR] = cvss_v3_vector
                vul_attrs[DESCRIPTION] = description
                vul_list.append(vul_attrs)
            pdf = pd.DataFrame(vul_list)
            return pdf
        else:
            return None

    @staticmethod
    def parse_cvss(cvss:dict)->(float, str):
        """
        This function takes the cvss section (raw.Results.Targets[i].Vulnerabilities[i].CVSS), and get the first
        CVSS available record, parse them and return score and vector. If no cvss record exists, return none, none.
        :param cvss: A dict which contains CVSS record issued by different organizations
        :return: a tuple of cvss score and cvss vector, if nothing found return None, None
        """
        if cvss:
            cvss_org = next(
                iter(cvss), None
            )  # Get the first organization name, if available
            cvss_v3_score = cvss.get(cvss_org, {}).get("V3Score")
            cvss_v3_vector = cvss.get(cvss_org, {}).get("V3Vector")
        else:
            cvss_v3_score = None
            cvss_v3_vector = None
        return cvss_v3_score, cvss_v3_vector

    def get_image_basic_info(self) -> dict:
        """
        This function return a dictionary of the basic info of the image, such as the below example:
        - image_name: redis
        - repo_tag: redis:latest
        - creation_time: 2024-09-13T09:13:14.153929471+02:00
        - os_family: debian
        - os_name: 12.7
        - os_architecture: amd64
        - image_id: sha256:590b81f2fea1af9798db4580a6299dafba020c2f5dc7d8d734663e7fa5299ca0
        :return:
        """
        # get some basic info from the metadata
        basic_info = self.parse_image_metadata()

        # add image name to dict
        basic_info[IMG_NAME] = self.raw_data.get("ArtifactName")

        # add creation time to dict
        basic_info[IMG_CREATION_TIME] = self.raw_data.get("CreatedAt")
        return basic_info

    def parse_image_metadata(self) -> dict:
        """
        This function parse the metadata of the image and return a dictionary of the basic infomation of the image,
        such as the below example:
        - os_family: debian
        - os_name: 12.7
        - os_architecture: amd64
        - repo_tag: redis:latest
        - image_id: sha256:590b81f2fea1af9798db4580a6299dafba020c2f5dc7d8d734663e7fa5299ca0

        :return:
        """
        result: dict = {}
        metadata = self.raw_data.get("Metadata")
        os_family: str = metadata.get("OS").get("Family")
        os_name: str = metadata.get("OS").get("Name")
        os_archi: str = metadata.get("ImageConfig").get("architecture")
        repo_tags: List = metadata.get("RepoTags")
        image_id: str = metadata.get("ImageID")
        result[OS_FAMILY] = os_family
        result[OS_NAME] = os_name
        result[OS_ARCHI] = os_archi
        result[REPO_TAGS] = repo_tags
        result[IMG_ID] = image_id
        return result
