[WIP] Changes to align python-stix2 with marking-prototype project.

stix2.1
Emmanuelle Vargas-Gonzalez 2017-07-13 07:57:01 -04:00
parent 7abcce7635
commit 747f0307a7
2 changed files with 32 additions and 54 deletions

View File

@ -116,9 +116,9 @@ def add_markings(obj, selectors, marking):
""" """
if selectors is None: if selectors is None:
object_markings.add_markings(obj, marking) return object_markings.add_markings(obj, marking)
else: else:
granular_markings.add_markings(obj, selectors, marking) return granular_markings.add_markings(obj, selectors, marking)
def clear_markings(obj, selectors): def clear_markings(obj, selectors):

View File

@ -1,28 +1,20 @@
import six
from stix2.markings import utils from stix2.markings import utils
def get_markings(obj): def get_markings(obj):
""" """
Get all object level markings from the given TLO object. Get all object level markings from the given SDO or SRO object.
Args: Args:
obj: A TLO object. obj: A SDO or SRO object.
Returns: Returns:
list: Marking IDs contained in the TLO. list: Marking IDs contained in the SDO or SRO. Empty list if no
markings are present in `object_marking_refs`.
""" """
object_markings = obj.get("object_marking_refs", []) return obj.get("object_marking_refs", [])
if not object_markings:
return []
elif isinstance(object_markings, six.string_types):
return [object_markings]
else:
return object_markings
def add_markings(obj, marking): def add_markings(obj, marking):
@ -30,23 +22,19 @@ def add_markings(obj, marking):
Appends an object level marking to the object_marking_refs collection. Appends an object level marking to the object_marking_refs collection.
Args: Args:
obj: A TLO object. obj: A SDO or SRO object.
marking: identifier or list of marking identifiers that apply to the marking: identifier or list of identifiers to apply SDO or SRO object.
TLO object.
Raises: Raises:
AssertionError: If `marking` fail data validation. AssertionError: If `marking` fail data validation.
""" """
marking = utils.convert_to_list(marking) marking = utils.convert_to_list(marking)
utils.validate(obj, marking=marking)
if not obj.get("object_marking_refs"): # TODO: Remove set for comparison and raise DuplicateMarkingException.
obj["object_marking_refs"] = list() object_markings = set(obj.get("object_marking_refs", []) + marking)
object_markings = set(obj.get("object_marking_refs") + marking) return obj.new_version(object_marking_refs=list(object_markings))
obj["object_marking_refs"] = list(object_markings)
def remove_markings(obj, marking): def remove_markings(obj, marking):
@ -54,13 +42,13 @@ def remove_markings(obj, marking):
Removes object level marking from the object_marking_refs collection. Removes object level marking from the object_marking_refs collection.
Args: Args:
obj: A TLO object. obj: A SDO or SRO object.
marking: identifier or list of marking identifiers that apply to the marking: identifier or list of identifiers that apply to the
TLO object. SDO or SRO object.
Raises: Raises:
AssertionError: If `marking` fail data validation. Also AssertionError: If markings to remove are not found on the provided
if markings to remove are not found on the provided TLO. SDO or SRO.
""" """
marking = utils.convert_to_list(marking) marking = utils.convert_to_list(marking)
@ -69,17 +57,14 @@ def remove_markings(obj, marking):
object_markings = obj.get("object_marking_refs", []) object_markings = obj.get("object_marking_refs", [])
if not object_markings: if not object_markings:
return [] return obj
if any(x not in obj["object_marking_refs"] for x in marking): if any(x not in obj["object_marking_refs"] for x in marking):
raise AssertionError("Unable to remove Object Level Marking(s) from " raise AssertionError("Unable to remove Object Level Marking(s) from "
"internal collection. Marking(s) not found...") "internal collection. Marking(s) not found...")
obj["object_marking_refs"] = [x for x in object_markings return obj.new_version(object_marking_refs=[x for x in object_markings
if x not in marking] if x not in marking])
if not obj.get("object_marking_refs"):
obj.pop("object_marking_refs")
def set_markings(obj, marking): def set_markings(obj, marking):
@ -88,15 +73,12 @@ def set_markings(obj, marking):
the collection. Refer to `clear_markings` and `add_markings` for details. the collection. Refer to `clear_markings` and `add_markings` for details.
Args: Args:
obj: A TLO object. obj: A SDO or SRO object.
marking: identifier or list of marking identifiers that apply to the marking: identifier or list of identifiers to apply in the
TLO object. SDO or SRO object.
""" """
utils.validate(obj, marking=marking) return add_markings(clear_markings(obj), marking)
clear_markings(obj)
add_markings(obj, marking)
def clear_markings(obj): def clear_markings(obj):
@ -104,31 +86,27 @@ def clear_markings(obj):
Removes all object level markings from the object_marking_refs collection. Removes all object level markings from the object_marking_refs collection.
Args: Args:
obj: A TLO object. obj: A SDO or SRO object.
""" """
try: return obj.new_version(object_marking_refs=None)
del obj["object_marking_refs"]
except KeyError:
raise AssertionError("Unable to clear Object Marking(s) from internal"
" collection. No Markings in object...")
def is_marked(obj, marking=None): def is_marked(obj, marking=None):
""" """
Checks if TLO is marked by any marking or by specific marking(s). Checks if SDO or SRO is marked by any marking or by specific marking(s).
Args: Args:
obj: A TLO object. obj: A SDO or SRO object.
marking: identifier or list of marking identifiers that apply to the marking: identifier or list of marking identifiers that apply to the
TLO object. SDO or SRO object.
Returns: Returns:
bool: True if TLO has object level markings. False otherwise. bool: True if SDO or SRO has object level markings. False otherwise.
Note: Note:
When a list of marking IDs is provided, if ANY of the provided marking When an identifier or list of identifiers is provided, if ANY of the
IDs matches, True is returned. provided marking refs match, True is returned.
""" """
marking = utils.convert_to_list(marking) marking = utils.convert_to_list(marking)