From 5a9f627669243d4e77b73a457a8ced019b29ad9f Mon Sep 17 00:00:00 2001 From: Michael Chisholm Date: Wed, 6 Jun 2018 15:30:45 -0400 Subject: [PATCH] Pickle-proof stix objects --- stix2/base.py | 8 +++++++- stix2/test/test_pickle.py | 17 +++++++++++++++++ 2 files changed, 24 insertions(+), 1 deletion(-) create mode 100644 stix2/test/test_pickle.py diff --git a/stix2/base.py b/stix2/base.py index 2afba16..9b5308f 100644 --- a/stix2/base.py +++ b/stix2/base.py @@ -185,7 +185,13 @@ class _STIXBase(collections.Mapping): # Handle attribute access just like key access def __getattr__(self, name): - if name in self: + # Pickle-proofing: pickle invokes this on uninitialized instances (i.e. + # __init__ has not run). So no "self" attributes are set yet. The + # usual behavior of this method reads an __init__-assigned attribute, + # which would cause infinite recursion. So this check disables all + # attribute reads until the instance has been properly initialized. + unpickling = "_inner" not in self.__dict__ + if not unpickling and name in self: return self.__getitem__(name) raise AttributeError("'%s' object has no attribute '%s'" % (self.__class__.__name__, name)) diff --git a/stix2/test/test_pickle.py b/stix2/test/test_pickle.py new file mode 100644 index 0000000..9e2cc9a --- /dev/null +++ b/stix2/test/test_pickle.py @@ -0,0 +1,17 @@ +import pickle + +import stix2 + + +def test_pickling(): + """ + Ensure a pickle/unpickle cycle works okay. + """ + identity = stix2.Identity( + id="identity--d66cb89d-5228-4983-958c-fa84ef75c88c", + name="alice", + description="this is a pickle test", + identity_class="some_class" + ) + + pickle.loads(pickle.dumps(identity))