# Copyright (C) 2024 Anthony Harrison
# SPDX-License-Identifier: Apache-2.0


class Vulnerability:
    def __init__(self, validation=None):
        self.vulnerability = {}
        if validation is not None and validation in [
            "openvex",
            "csaf",
            "cyclonedx",
            "spdx",
        ]:
            self.validation = validation
        else:
            self.validation = None

    def initialise(self):
        self.vulnerability = {}

    def set_name(self, name):
        self.vulnerability["product"] = name

    def set_release(self, release):
        self.vulnerability["release"] = release

    def set_id(self, id):
        self.vulnerability["id"] = id

    def set_description(self, description):
        self.vulnerability["description"] = description

    def set_comment(self, comment):
        self.vulnerability["comment"] = comment

    def validate_status(self, status):
        if self.validation in ["openvex", "spdx"] and status in [
            "not_affected",
            "affected",
            "fixed",
            "under_investigation",
        ]:
            return True
        elif self.validation == "cyclonedx" and status in [
            "resolved",
            "resolved_with_pedigree",
            "exploitable",
            "in_triage",
            "false_positive",
            "not_affected",
        ]:
            return True
        elif self.validation == "csaf" and status in [
            "first_affected",
            "first_fixed",
            "fixed",
            "known_affected",
            "known_not_affected",
            "last_affected",
            "recommended",
            "under_investigation",
        ]:
            return True
        return False

    def set_status(self, status):
        if self.validation is None:
            # No validation performed
            self.vulnerability["status"] = status.lower()
        elif self.validate_status(status.lower()):
            # Valid field according to specification
            self.vulnerability["status"] = status.lower()
        else:
            # Invalid field detected
            # Set to default state
            if self.validation == "cyclonedx":
                self.vulnerability["status"] = "in_triage"
            else:
                self.vulnerability["status"] = "under_investigation"

    def validate_justification(self, justification):
        if self.validation == "openvex" and justification in [
            "component_not_present",
            "inline_mitigations_already_exist",
            "vulnerable_code_cannot_be_controlled_by_adversary",
            "vulnerable_code_not_in_execute_path",
            "vulnerable_code_not_present",
        ]:
            return True
        elif self.validation == "cyclonedx" and justification in [
            "code_not_present",
            "code_not_reachable",
            "requires_configuration",
            "requires_dependency",
            "requires_environment",
            "protected_by_compiler",
            "protected_at_runtime",
            "protected_at_perimeter",
            "protected_by_mitigating_control",
        ]:
            return True
        elif self.validation == "csaf" and justification in [
            "component_not_present",
            "inline_mitigations_already_exist",
            "vulnerable_code_cannot_be_controlled_by_adversary",
            "vulnerable_code_not_in_execute_path",
            "vulnerable_code_not_present",
        ]:
            return True
        elif self.validation == "spdx" and justification in [
            "componentNotPresent",
            "inlineMitigationsAlreadyExist",
            "vulnerableCodeCannotBeControlledByAdversary",
            "vulnerableCodeNotInExecutePath",
            "vulnerableCodeNotPresent",
        ]:
            return True
        return False

    def set_justification(self, justification):
        if self.validation is None:
            # No validation performed
            self.vulnerability["justification"] = justification.lower()
        elif self.validate_justification(justification.lower()):
            # Valid field according to specification
            self.vulnerability["justification"] = justification.lower()
        elif self.validate_justification(justification):
            # Valid field according to specification
            self.vulnerability["justification"] = justification
        else:
            # Invalid field detected
            self.vulnerability["justification"] = None

    def validate_remediation(self, remediation):
        if self.validation == "openvex":
            # No remediation action to be validated
            return True
        elif self.validation == "cyclonedx" and remediation in [
            "can_not_fix",
            "will_not_fix",
            "update",
            "rollback",
            "workaround_available",
        ]:
            return True
        elif self.validation == "csaf" and remediation in [
            "mitigation",
            "no_fix_planned",
            "none_available",
            "vendor_fix",
            "workaround",
        ]:
            return True
        elif self.validation == "spdx":
            # No remediation action to be validated
            return True
        return False

    def set_remediation(self, remediation):
        if self.validation is None:
            # No validation performed
            self.vulnerability["remediation"] = remediation.lower()
        elif self.validate_remediation(remediation.lower()):
            # Valid field according to specification
            self.vulnerability["remediation"] = remediation.lower()
        else:
            # Invalid field detected
            self.vulnerability["remediation"] = None

    def set_action(self, value):
        self.vulnerability["action"] = value

    def set_value(self, key, value):
        self.vulnerability[key] = value

    def get_vulnerability(self):
        return self.vulnerability

    def get_value(self, attribute):
        return self.vulnerability.get(attribute, None)

    def debug_vulnerability(self):
        print("OUTPUT:", self.vulnerability)

    def show_vulnerability(self):
        for key in self.vulnerability:
            print(f"{key}    : {self.vulnerability[key]}")

    def copy_vulnerability(self, vulnerability_info):
        for key in vulnerability_info:
            self.set_value(key, vulnerability_info[key])

    def get_name(self):
        return self.get_value("name")
