#MAEC Malware Subject Class

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

#Compatible with MAEC v4.1
#Last updated 08/20/2014

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

import maec
import maec.bindings.maec_package as package_binding
from maec.bundle.bundle import Bundle
from maec.package.action_equivalence import ActionEquivalenceList
from maec.package.analysis import Analysis
from maec.package.malware_subject_reference import MalwareSubjectReference
from maec.package.object_equivalence import ObjectEquivalenceList

class MinorVariants(maec.EntityList):
    _contained_type = Object
    _binding_class = package_binding.MinorVariantListType
    _binding_var = "Minor_Variant"
    _namespace = maec.package._namespace
    
class Analyses(maec.EntityList):
    _contained_type = Analysis
    _binding_class = package_binding.AnalysisListType
    _binding_var = "Analysis"
    _namespace = maec.package._namespace

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

    malware_subject_reference = maec.TypedField("Malware_Subject_Reference", MalwareSubjectReference, multiple = True)
    type = maec.TypedField("type", VocabString)

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


class MalwareSubjectRelationshipList(maec.EntityList):
    _contained_type = MalwareSubjectRelationship
    _binding_class = package_binding.MalwareSubjectRelationshipListType
    _binding_var = "Relationship"
    _namespace = maec.package._namespace

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

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

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

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

    meta_analysis = maec.TypedField("Meta_Analysis", MetaAnalysis)
    bundle = maec.TypedField("Bundle", Bundle, multiple = True)
    bundle_external_reference = maec.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 = maec.package._namespace

    tools = maec.TypedField("Tools", ToolInformation)
    debugging_file = maec.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 = maec.package._namespace

    name = maec.TypedField("Name", VocabString)
    value = maec.TypedField("Value")

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


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

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

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

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

    malware_binary = maec.TypedField("Malware_Binary", MalwareBinaryConfigurationStorageDetails)
    file = maec.TypedField("File", File)
    url = maec.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 = maec.package._namespace

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

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


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

    is_encoded = maec.TypedField("is_encoded")
    is_encrypted = maec.TypedField("is_encrypted")
    algorithm_details = maec.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 = maec.package._namespace

    storage = maec.TypedField("Storage", MalwareConfigurationStorageDetails)
    obfuscation = maec.TypedField("Obfuscation", MalwareConfigurationObfuscationDetails)
    configuration_parameter = maec.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 = maec.package._namespace

    id_ = maec.TypedField("id")
    malware_instance_object_attributes = maec.TypedField("Malware_Instance_Object_Attributes", Object)
    label = maec.TypedField("Label", VocabString, multiple=True)
    configuration_details = maec.TypedField("Configuration_Details", MalwareConfigurationDetails)
    minor_variants = maec.TypedField("Minor_Variants", MinorVariants)
    development_environment = maec.TypedField("Development_Environment", MalwareDevelopmentEnvironment)
    #field_data = maec.TypedField("field_data") # TODO: support metadata:fieldDataEntry
    analyses = maec.TypedField("Analyses", Analyses)
    findings_bundles = maec.TypedField("Findings_Bundles", FindingsBundleList)
    relationships = maec.TypedField("Relationships", MalwareSubjectRelationshipList)
    compatible_platform = maec.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_ = maec.utils.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):
    _contained_type = MalwareSubject
    _binding_class = package_binding.MalwareSubjectListType
    _binding_var = "Malware_Subject"
    _namespace = maec.package._namespace