text: # Copyright (c) 2017, The MITRE Corporation. All rights reserved.
  # See LICENSE.txt for complete terms.
  
  """
  STIX Extension for AIS Data Markings
  
  Unlike the other marking extensions, the AIS marking extension is not loaded
  automatically, since AIS markings are not a part of the published STIX 1.x
  specifications. They are included in python-stix because they're common enough
  that it is not worth creating a separate package.
  
  If you are writing code that needs to parse AIS markings, make sure that your
  program imports this module before beginning to parse any STIX documents:
  
  .. code-block:: python
  
      import stix.extensions.marking.ais
  
  """
  
  from mixbox import fields
  from mixbox.namespaces import Namespace
  
  import stix.bindings.extensions.marking.ais as ais_binding
  import stix.data_marking
  from stix.data_marking import MarkingStructure
  
  
  def validate_value(instance, value):
      allowed = instance._ALLOWED_VALUES
      if not value:
          return
      elif not allowed:
          return
      elif value in allowed:
          return
      else:
          error = "Value must be one of {allowed}. Received '{value}'"
          error = error.format(**locals())
          raise ValueError(error)
  
  
  class AISConsentType(stix.Entity):
      _binding = ais_binding
      _binding_class = _binding.AISConsentType
      _namespace = 'http://www.us-cert.gov/STIXMarkingStructure#AISConsentMarking-2'
      _ALLOWED_VALUES = ('EVERYONE', 'USG', 'NONE')
  
      consent = fields.TypedField("consent", preset_hook=validate_value)
  
      def __init__(self, consent=None):
          super(AISConsentType, self).__init__()
          self.consent = consent
  
  
  class TLPMarkingType(stix.Entity):
      _binding = ais_binding
      _binding_class = _binding.TLPMarkingType
      _namespace = 'http://www.us-cert.gov/STIXMarkingStructure#AISConsentMarking-2'
      _ALLOWED_VALUES = ('WHITE', 'GREEN', 'AMBER')
  
      color = fields.TypedField("color", preset_hook=validate_value)
  
      def __init__(self, color=None):
          super(TLPMarkingType, self).__init__()
          self.color = color
  
  
  class NotProprietary(stix.Entity):
      _binding = ais_binding
      _binding_class = _binding.NotProprietary
      _namespace = 'http://www.us-cert.gov/STIXMarkingStructure#AISConsentMarking-2'
  
      cisa_proprietary = fields.TypedField("CISA_Proprietary")
      ais_consent = fields.TypedField("AISConsent", AISConsentType, key_name="ais_consent")
      tlp_marking = fields.TypedField("TLPMarking", TLPMarkingType, key_name="tlp_marking")
  
      def __init__(self, cisa_proprietary='false', ais_consent=None,
                   tlp_marking=None):
          super(NotProprietary, self).__init__()
  
          self.cisa_proprietary = cisa_proprietary
          self.ais_consent = ais_consent
          self.tlp_marking = tlp_marking
  
  
  class IsProprietary(stix.Entity):
      _binding = ais_binding
      _binding_class = _binding.IsProprietary
      _namespace = 'http://www.us-cert.gov/STIXMarkingStructure#AISConsentMarking-2'
  
      cisa_proprietary = fields.TypedField("CISA_Proprietary")
      ais_consent = fields.TypedField("AISConsent", AISConsentType, key_name="ais_consent")
      tlp_marking = fields.TypedField("TLPMarking", TLPMarkingType, key_name="tlp_marking")
  
      def __init__(self, cisa_proprietary='true', ais_consent=None,
                   tlp_marking=None):
          super(IsProprietary, self).__init__()
  
          self.cisa_proprietary = cisa_proprietary
          self.ais_consent = ais_consent
          self.tlp_marking = tlp_marking
  
  
  @stix.register_extension
  class AISMarkingStructure(MarkingStructure):
      _binding = ais_binding
      _binding_class = _binding.AISMarkingStructure
      _namespace = 'http://www.us-cert.gov/STIXMarkingStructure#AISConsentMarking-2'
      _XSI_TYPE = "AIS:AISMarkingStructure"
  
      is_proprietary = fields.TypedField("Is_Proprietary", IsProprietary)
      not_proprietary = fields.TypedField("Not_Proprietary", NotProprietary)
  
      def __init__(self, is_proprietary=None, not_proprietary=None):
          super(AISMarkingStructure, self).__init__()
  
          self.is_proprietary = is_proprietary
          self.not_proprietary = not_proprietary
  
  
  NAMESPACES = [
      Namespace('http://www.us-cert.gov/STIXMarkingStructure#AISConsentMarking-2', 'AIS', 'http://www.us-cert.gov/sites/default/files/STIX_Namespace/AIS_Bundle_Marking_1.1.1_v1.0.xsd')
  ]
  
  
  def _update_namespaces():
      # Update the python-stix namespace dictionary
      from stix.utils import nsparser
      import mixbox.namespaces
  
      nsparser.STIX_NAMESPACES.add_namespace(NAMESPACES[0])
      mixbox.namespaces.register_namespace(NAMESPACES[0])
  
  
  _update_namespaces()
  
  
  # IndustryType allowed sectors
  #: Chemical Sector
  CHEMICAL_SECTOR = 'Chemical Sector'
  #: Chemical Sector
  COMMERCIAL_FACILITIES_SECTOR = 'Commercial Facilities Sector'
  #: Commercial Facilities Sector
  COMMUNICATIONS_SECTOR = 'Communications Sector'
  #: Critical Manufacturing Sector
  CRITICAL_MANUFACTURING_SECTOR = 'Critical Manufacturing Sector'
  #: Dams Sector
  DAMS_SECTOR = 'Dams Sector'
  #: Defense Industrial Base Sector
  DEFENSE_INDUSTRIAL_BASE_SECTOR = 'Defense Industrial Base Sector'
  #: Emergency Services Sector
  EMERGENCY_SERVICES_SECTOR = 'Emergency Services Sector'
  #: Energy Sector
  ENERGY_SECTOR = 'Energy Sector'
  #: Financial Services Sector
  FINANCIAL_SERVICES_SECTOR = 'Financial Services Sector'
  #: Food and Agriculture Sector
  FOOD_AND_AGRICULTURE_SECTOR = 'Food and Agriculture Sector'
  #: Government Facilities Sector
  GOVERNMENT_FACILITIES_SECTOR = 'Government Facilities Sector'
  #: Healthcare and Public Health Sector
  HEALTH_CARE_AND_PUBLIC_HEALTH_SECTOR = 'Healthcare and Public Health Sector'
  #: Information Technology Sector
  INFORMATION_TECHNOLOGY_SECTOR = 'Information Technology Sector'
  #: Nuclear Reactors, Materials, and Waste Sector
  NUCLEAR_REACTORS_MATERIALS_AND_WASTE_SECTOR = 'Nuclear Reactors, Materials, and Waste Sector'
  #: Other
  OTHER = 'Other'
  #: Transportation Systems Sector
  TRANSPORTATION_SYSTEMS_SECTOR = 'Transportation Systems Sector'
  #: Water and Wastewater Systems Sector
  WATER_AND_WASTEWATER_SYSTEMS_SECTOR = 'Water and Wastewater Systems Sector'
  
  
  def _validate_and_create_industry_type(industry_type):
      INDUSTRY_SECTORS = (CHEMICAL_SECTOR, COMMERCIAL_FACILITIES_SECTOR,
                          COMMUNICATIONS_SECTOR, CRITICAL_MANUFACTURING_SECTOR,
                          DAMS_SECTOR, DEFENSE_INDUSTRIAL_BASE_SECTOR,
                          EMERGENCY_SERVICES_SECTOR, ENERGY_SECTOR,
                          FINANCIAL_SERVICES_SECTOR, FOOD_AND_AGRICULTURE_SECTOR,
                          GOVERNMENT_FACILITIES_SECTOR,
                          HEALTH_CARE_AND_PUBLIC_HEALTH_SECTOR,
                          INFORMATION_TECHNOLOGY_SECTOR,
                          NUCLEAR_REACTORS_MATERIALS_AND_WASTE_SECTOR,
                          TRANSPORTATION_SYSTEMS_SECTOR, OTHER,
                          WATER_AND_WASTEWATER_SYSTEMS_SECTOR)
  
      lower_case_sectors = tuple(x.lower() for x in INDUSTRY_SECTORS)
      result = ""
      error = False
      val = []
  
      if isinstance(industry_type, str):
          # Pipe-delimited or single string supplied.
          val = [x.lower().strip() for x in industry_type.split("|")]
  
      elif isinstance(industry_type, (list, tuple)):
          # Create pipe-delimited string when list of strings is provided.
          val = [x.lower().strip() for x in industry_type]
  
      else:
          error = True
  
      for item in val:
          for idx, sector in enumerate(lower_case_sectors):
              if item == sector:
                  if not result:
                      result = INDUSTRY_SECTORS[idx]
                  else:
                      result = "{0}|{1}".format(result, INDUSTRY_SECTORS[idx])
                  break
          else:
              # The sectors collection was exhausted. No match found.
              error = True
              break
  
      if not error and val:
          return result
  
      msg = 'IndustryType must be one of the following: {0}. Received \'{1}\'.'
      raise ValueError(msg.format(INDUSTRY_SECTORS, industry_type))
  
  
  def add_ais_marking(stix_package, proprietary, consent, color, **kwargs):
      """
      This utility functions aids in the creation of an AIS marking and appends
      it to the provided STIX package.
  
      Args:
          stix_package: A stix.core.STIXPackage object.
          proprietary: True if marking uses IsProprietary, False for
              NotProprietary.
          consent: A string with one of the following values: "EVERYONE", "NONE"
              or "USG".
          color: A string that corresponds to TLP values: "WHITE", "GREEN" or
              "AMBER".
          **kwargs: Six required keyword arguments that are used to create a CIQ
              identity object. These are: country_name_code,
              country_name_code_type, admin_area_name_code,
              admin_area_name_code_type, organisation_name, industry_type.
  
      Raises:
          ValueError: When keyword arguments are missing. User did not supply
              correct values for: proprietary, color and consent.
  
      Note:
          The following line is required to register the AIS extension::
  
              >>> import stix.extensions.marking.ais
  
          Any Markings under STIX Header will be removed. Please follow the
          guidelines for `AIS`_.
  
          The industry_type keyword argument accepts: a list of string based on
          defined sectors, a pipe-delimited string of sectors, or a single
          sector.
  
      .. _AIS:
          https://www.us-cert.gov/ais
  
      """
      from stix.common import InformationSource
      from stix.extensions.identity.ciq_identity_3_0 import (
          CIQIdentity3_0Instance, STIXCIQIdentity3_0, PartyName, Address,
          Country, NameElement, OrganisationInfo, AdministrativeArea)
      from stix.core.stix_header import STIXHeader
      from stix.data_marking import MarkingSpecification, Marking
  
      args = ('country_name_code', 'country_name_code_type', 'industry_type',
              'admin_area_name_code', 'admin_area_name_code_type',
              'organisation_name')
  
      diff = set(args) - set(kwargs.keys())
  
      if diff:
          msg = 'All keyword arguments must be provided. Missing: {0}'
          raise ValueError(msg.format(tuple(diff)))
  
      party_name = PartyName()
      party_name.add_organisation_name(kwargs['organisation_name'])
  
      country = Country()
      country_name = NameElement()
      country_name.name_code = kwargs['country_name_code']
      country_name.name_code_type = kwargs['country_name_code_type']
      country.add_name_element(country_name)
  
      admin_area = AdministrativeArea()
      admin_area_name = NameElement()
      admin_area_name.name_code = kwargs['admin_area_name_code']
      admin_area_name.name_code_type = kwargs['admin_area_name_code_type']
      admin_area.add_name_element(admin_area_name)
  
      address = Address()
      address.country = country
      address.administrative_area = admin_area
  
      org_info = OrganisationInfo()
      org_info.industry_type = _validate_and_create_industry_type(kwargs['industry_type'])
  
      id_spec = STIXCIQIdentity3_0()
      id_spec.party_name = party_name
      id_spec.add_address(address)
      id_spec.organisation_info = org_info
  
      identity = CIQIdentity3_0Instance()
      identity.specification = id_spec
  
      if proprietary is True:
          proprietary_obj = IsProprietary()
          consent = 'EVERYONE'
      elif proprietary is False:
          proprietary_obj = NotProprietary()
      else:
          raise ValueError('proprietary expected True or False.')
  
      proprietary_obj.ais_consent = AISConsentType(consent=consent)
      proprietary_obj.tlp_marking = TLPMarkingType(color=color)
  
      ais_marking = AISMarkingStructure()
  
      if isinstance(proprietary_obj, IsProprietary):
          ais_marking.is_proprietary = proprietary_obj
      else:
          ais_marking.not_proprietary = proprietary_obj
  
      marking_spec = MarkingSpecification()
      marking_spec.controlled_structure = '//node() | //@*'
      marking_spec.marking_structures.append(ais_marking)
      marking_spec.information_source = InformationSource()
      marking_spec.information_source.identity = identity
  
      if not stix_package.stix_header:
          stix_package.stix_header = STIXHeader()
  
      # Removes any other Markings if present.
      stix_package.stix_header.handling = Marking()
      stix_package.stix_header.handling.add_marking(marking_spec)
  