# MAEC Malware Subject Class

# Copyright (c) 2018, The MITRE Corporation
# All rights reserved

from mixbox import fields
from mixbox import idgen

from cybox.common import vocabs, VocabString, PlatformSpecification, ToolInformation
from cybox.objects.file_object import File
from cybox.objects.uri_object import URI
from cybox.core import Object

import maec
from . import _namespace
import maec.bindings.maec_package as package_binding
from maec.bundle import Bundle
from maec.package import (ActionEquivalenceList, Analysis,
                          MalwareSubjectReference, ObjectEquivalenceList)
from maec.vocabs.vocabs import MalwareLabel
from maec.vocabs.vocabs import MalwareConfigurationParameter as MalwareConfigParameterVocab
from maec.vocabs.vocabs import MalwareSubjectRelationship as MalwareSubjectRelationshipVocab

class MinorVariants(maec.EntityList):
    _binding_class = package_binding.MinorVariantListType
    _namespace = _namespace
    minor_variant = fields.TypedField("Minor_Variant", Object, multiple=True)

class Analyses(maec.EntityList):
    _binding_class = package_binding.AnalysisListType
    _namespace = _namespace
    analysis = fields.TypedField("Analysis", Analysis, multiple=True)

class MalwareSubjectRelationship(maec.Entity):
    _binding = package_binding
    _binding_class = package_binding.MalwareSubjectRelationshipType
    _namespace = _namespace

    malware_subject_reference = fields.TypedField("Malware_Subject_Reference", MalwareSubjectReference, multiple = True)
    type_ = vocabs.VocabField("Type", MalwareSubjectRelationshipVocab)

    def __init__(self):
        super(MalwareSubjectRelationship, self).__init__()


class MalwareSubjectRelationshipList(maec.EntityList):
    _binding_class = package_binding.MalwareSubjectRelationshipListType
    _namespace = _namespace
    relationship = fields.TypedField("Relationship", MalwareSubjectRelationship, multiple=True)

class MetaAnalysis(maec.Entity):
    _binding = package_binding
    _binding_class = package_binding.MetaAnalysisType
    _namespace = _namespace

    action_equivalences = fields.TypedField("Action_Equivalences", ActionEquivalenceList)
    object_equivalences = fields.TypedField("Object_Equivalences", ObjectEquivalenceList)

    def __init__(self):
        super(MetaAnalysis, self).__init__()

class FindingsBundleList(maec.Entity):
    _binding = package_binding
    _binding_class = package_binding.FindingsBundleListType
    _namespace = _namespace

    meta_analysis = fields.TypedField("Meta_Analysis", MetaAnalysis)
    bundle = fields.TypedField("Bundle", Bundle, multiple = True)
    bundle_external_reference = fields.TypedField("Bundle_External_Reference", multiple = True)

    def __init__(self):
        super(FindingsBundleList, self).__init__()

    def add_bundle(self, bundle):
        if not self.bundle:
            self.bundle = []
        self.bundle.append(bundle)

    def add_bundle_external_reference(self, bundle_external_reference):
        if not self.bundle_external_reference:
            self.bundle_external_reference = []
        self.bundle_external_reference.append(bundle_external_reference)

class MalwareDevelopmentEnvironment(maec.Entity):
    _binding = package_binding
    _binding_class = package_binding.MalwareDevelopmentEnvironmentType
    _namespace = _namespace

    tools = fields.TypedField("Tools", ToolInformation)
    debugging_file = fields.TypedField("Debugging_File", File, multiple = True)

    def __init__(self):
        super(MalwareDevelopmentEnvironment, self).__init__()


class MalwareConfigurationParameter(maec.Entity):
    _binding = package_binding
    _binding_class = package_binding.MalwareConfigurationParameterType
    _namespace = _namespace

    name = vocabs.VocabField("Name", MalwareConfigParameterVocab)
    value = fields.TypedField("Value")

    def __init__(self):
        super(MalwareConfigurationParameter, self).__init__()


class MalwareBinaryConfigurationStorageDetails(maec.Entity):
    _binding = package_binding
    _binding_class = package_binding.MalwareBinaryConfigurationStorageDetailsType
    _namespace = _namespace

    file_offset = fields.TypedField("File_Offset")
    section_name = fields.TypedField("Section_Name")
    section_offset = fields.TypedField("Section_Offset")

    def __init__(self):
        super(MalwareBinaryConfigurationStorageDetails, self).__init__()

class MalwareConfigurationStorageDetails(maec.Entity):
    _binding = package_binding
    _binding_class = package_binding.MalwareConfigurationStorageDetailsType
    _namespace = _namespace

    malware_binary = fields.TypedField("Malware_Binary", MalwareBinaryConfigurationStorageDetails)
    file = fields.TypedField("File", File)
    url = fields.TypedField("URL", URI, multiple = True)

    def __init__(self):
        super(MalwareConfigurationStorageDetails, self).__init__()

class MalwareConfigurationObfuscationAlgorithm(maec.Entity):
    _binding = package_binding
    _binding_class = package_binding.MalwareConfigurationObfuscationAlgorithmType
    _namespace = _namespace

    ordinal_position = fields.TypedField("ordinal_position")
    key = fields.TypedField("Key")
    algorithm_name = fields.TypedField("Algorithm_Name", VocabString)

    def __init__(self):
        super(MalwareConfigurationObfuscationAlgorithm, self).__init__()


class MalwareConfigurationObfuscationDetails(maec.Entity):
    _binding = package_binding
    _binding_class = package_binding.MalwareConfigurationObfuscationDetailsType
    _namespace = _namespace

    is_encoded = fields.TypedField("is_encoded")
    is_encrypted = fields.TypedField("is_encrypted")
    algorithm_details = fields.TypedField("Algorithm_Details", MalwareConfigurationObfuscationAlgorithm, multiple = True)

    def __init__(self):
        super(MalwareConfigurationObfuscationDetails, self).__init__()
        self.algorithm_details = []


class MalwareConfigurationDetails(maec.Entity):
    _binding = package_binding
    _binding_class = package_binding.MalwareConfigurationDetailsType
    _namespace = _namespace

    storage = fields.TypedField("Storage", MalwareConfigurationStorageDetails)
    obfuscation = fields.TypedField("Obfuscation", MalwareConfigurationObfuscationDetails)
    configuration_parameter = fields.TypedField("Configuration_Parameter", MalwareConfigurationParameter, multiple = True)

    def __init__(self):
        super(MalwareConfigurationDetails, self).__init__()

class MalwareSubject(maec.Entity):
    _binding = package_binding
    _binding_class = package_binding.MalwareSubjectType
    _namespace = _namespace

    id_ = fields.TypedField("id")
    malware_instance_object_attributes = fields.TypedField("Malware_Instance_Object_Attributes", Object)
    label = vocabs.VocabField("Label", MalwareLabel, multiple=True)
    configuration_details = fields.TypedField("Configuration_Details", MalwareConfigurationDetails)
    minor_variants = fields.TypedField("Minor_Variants", MinorVariants)
    development_environment = fields.TypedField("Development_Environment", MalwareDevelopmentEnvironment)
    #field_data = fields.TypedField("field_data") # TODO: support metadata:fieldDataEntry
    analyses = fields.TypedField("Analyses", Analyses)
    findings_bundles = fields.TypedField("Findings_Bundles", FindingsBundleList)
    relationships = fields.TypedField("Relationships", MalwareSubjectRelationshipList)
    compatible_platform = fields.TypedField("Compatible_Platform", PlatformSpecification, multiple=True)

    def __init__(self, id = None, malware_instance_object_attributes = None):
        super(MalwareSubject, self).__init__()
        if id:
            self.id_ = id
        else:
            self.id_ = idgen.create_id(prefix="malware_subject")
        #Set the Malware Instance Object Attributes (a CybOX object) if they are not none
        self.malware_instance_object_attributes = malware_instance_object_attributes

    #Public methods
    #Set the Malware_Instance_Object_Attributes with a CybOX object
    def set_malware_instance_object_attributes(self, malware_instance_object_attributes):
        self.malware_instance_object_attributes = malware_instance_object_attributes

    #Add an Analysis to the Analyses
    def add_analysis(self, analysis):
        if not self.analyses:
            self.analyses = Analyses()
        self.analyses.append(analysis)

    def get_analyses(self):
        return self.analyses

    #Get all Bundles in the Subject
    def get_all_bundles(self):
        return self.findings_bundles.bundle

    #Add a MAEC Bundle to the Findings Bundles
    def add_findings_bundle(self, bundle):
        if not self.findings_bundles:
            self.findings_bundles = FindingsBundleList()
        self.findings_bundles.add_bundle(bundle)

    def deduplicate_bundles(self):
        """DeDuplicate all Findings Bundles in the Malware Subject. For now, only handles Objects"""
        all_bundles = self.get_all_bundles()
        for bundle in all_bundles:
            bundle.deduplicate()

    def dereference_bundles(self):
        """Dereference all Findings Bundles in the Malware Subject. For now, only handles Objects"""
        all_bundles = self.get_all_bundles()
        for bundle in all_bundles:
            bundle.dereference_objects([self.malware_instance_object_attributes])

    def normalize_bundles(self):
        """Normalize all Findings Bundles in the Malware Subject. For now, only handles Objects"""
        all_bundles = self.get_all_bundles()
        for bundle in all_bundles:
            bundle.normalize_objects()

class MalwareSubjectList(maec.EntityList):
    _binding_class = package_binding.MalwareSubjectListType
    #_binding_var = "Malware_Subject"
    _namespace = _namespace

    malware_subject = fields.TypedField("Malware_Subject", MalwareSubject, multiple=True)
