From 94e832b56b68642a6c0893d45f36c7ac2a2362d0 Mon Sep 17 00:00:00 2001
From: ziirish <ziirish@ziirish.info>
Date: Wed, 30 May 2018 11:04:19 +0200
Subject: [PATCH] improve wildcard marshalling performances

---
 burpui/api/custom/my_fields.py      | 72 +++++++++++++++++------------
 burpui/api/custom/my_marshalling.py | 13 +++---
 2 files changed, 49 insertions(+), 36 deletions(-)

diff --git a/burpui/api/custom/my_fields.py b/burpui/api/custom/my_fields.py
index d1fec2d1..5b42eda8 100644
--- a/burpui/api/custom/my_fields.py
+++ b/burpui/api/custom/my_fields.py
@@ -100,13 +100,11 @@ class LocalizedString(fields.String):
 
 
 class Wildcard(fields.Raw):
-    exclude = []
-    # keep a track of the last object
-    _idx = 0
+    exclude = set()
     # cache the flat object
     _flat = None
     _obj = None
-    _cache = []
+    _cache = set()
     _last = None
 
     def __init__(self, cls_or_instance, **kwargs):
@@ -122,16 +120,26 @@ class Wildcard(fields.Raw):
             self.container = cls_or_instance
 
     def _flatten(self, obj):
-        if obj == self._obj and self._flat:
+        if obj is None:
+            return None
+        if obj == self._obj and self._flat is not None:
             return self._flat
         if isinstance(obj, dict):
-            self._flat = obj.items()
+            # self._flat needs to implement pop()
+            self._flat = [x for x in obj.items()]
         else:
-            attributes = inspect.getmembers(obj, lambda a: not(inspect.isroutine(a)))
-            self._flat = [x for x in attributes if match_attributes(x)]
 
-        self._idx = 0
-        self._cache = []
+            def __match_attributes(attribute):
+                attr_name, attr_obj = attribute
+                if inspect.isroutine(attr_obj) or \
+                        (attr_name.startswith('__') and attr_name.endswith('__')):
+                    return False
+                return True
+
+            attributes = inspect.getmembers(obj)
+            self._flat = [x for x in attributes if __match_attributes(x)]
+
+        self._cache = set()
         self._obj = obj
         return self._flat
 
@@ -139,26 +147,37 @@ class Wildcard(fields.Raw):
     def key(self):
         return self._last
 
+    def reset(self):
+        self.exclude = set()
+        self._flat = None
+        self._obj = None
+        self._cache = set()
+        self._last = None
+
     def output(self, key, obj, ordered=False, **kwargs):
-        flat = self._flatten(obj)
         value = None
         reg = fnmatch.translate(key)
 
-        for idx, (objkey, val) in enumerate(flat):
-            if idx < self._idx:
-                continue
-            if objkey not in self._cache and \
-                    objkey not in self.exclude and \
-                    re.match(reg, objkey, re.IGNORECASE):
-                value = val
-                self._cache.append(objkey)
-                self._last = objkey
-                self._idx = idx
-                break
+        if self._flatten(obj):
+            while True:
+                try:
+                    # we are using pop() so that we don't
+                    # loop over the whole object every time dropping the
+                    # complexity to O(n)
+                    (objkey, val) = self._flat.pop()
+                    if objkey not in self._cache and \
+                            objkey not in self.exclude and \
+                            re.match(reg, objkey, re.IGNORECASE):
+                        value = val
+                        self._cache.add(objkey)
+                        self._last = objkey
+                        break
+                except IndexError:
+                    break
 
         if value is None:
             if self.default is not None:
-                return self.default
+                return self.container.format(self.default)
             return None
 
         return self.container.format(value)
@@ -175,10 +194,3 @@ class Wildcard(fields.Raw):
         if mask:
             model = mask.apply(model)
         return self.__class__(model, **kwargs)
-
-
-def match_attributes(attribute):
-    attr_name, _ = attribute
-    if attr_name.startswith('__') and attr_name.endswith('__'):
-        return False
-    return True
diff --git a/burpui/api/custom/my_marshalling.py b/burpui/api/custom/my_marshalling.py
index 47c96de3..c2b89c39 100644
--- a/burpui/api/custom/my_marshalling.py
+++ b/burpui/api/custom/my_marshalling.py
@@ -66,11 +66,11 @@ def marshal(data, fields, envelope=None, skip_none=False, mask=None, ordered=Fal
                 field = make(val)
                 # exclude already parsed keys from the wildcard
                 is_wildcard = isinstance(field, Wildcard)
-                if is_wildcard and keys:
-                    for tmp in keys:
-                        if tmp not in field.exclude:
-                            field.exclude.append(tmp)
-                    keys = []
+                if is_wildcard:
+                    field.reset()
+                    if keys:
+                        field.exclude |= set(keys)
+                        keys = []
                 value = field.output(dkey, data, ordered=ordered)
                 if is_wildcard:
 
@@ -83,7 +83,8 @@ def marshal(data, fields, envelope=None, skip_none=False, mask=None, ordered=Fal
                     _append(key, value)
                     while True:
                         value = field.output(dkey, data, ordered=ordered)
-                        if value is None or value == field.default:
+                        if value is None or \
+                                value == field.container.format(field.default):
                             break
                         key = field.key
                         _append(key, value)