diff --git a/lib/core/bigarray.py b/lib/core/bigarray.py index 52a9808d4..55c586b1b 100644 --- a/lib/core/bigarray.py +++ b/lib/core/bigarray.py @@ -11,12 +11,23 @@ except: import pickle import os +import sys import tempfile from lib.core.exception import SqlmapSystemException -from lib.core.settings import BIGARRAY_CHUNK_LENGTH +from lib.core.settings import BIGARRAY_CHUNK_SIZE from lib.core.settings import BIGARRAY_TEMP_PREFIX +def _size_of(object): + """ + Returns total size of a given object (in bytes) + """ + + if hasattr(object, "__iter__"): + return sum(_size_of(_) for _ in object) + else: + return sys.getsizeof(object) + class Cache(object): """ Auxiliary class used for storing cached chunks @@ -34,13 +45,20 @@ class BigArray(list): def __init__(self): self.chunks = [[]] + self.chunk_length = sys.maxint self.cache = None self.filenames = set() self._os_remove = os.remove + self._size_counter = 0 def append(self, value): self.chunks[-1].append(value) - if len(self.chunks[-1]) >= BIGARRAY_CHUNK_LENGTH: + if self.chunk_length == sys.maxint: + self._size_counter += _size_of(value) + if self._size_counter >= BIGARRAY_CHUNK_SIZE: + self.chunk_length = len(self.chunks[-1]) + self._size_counter = None + if len(self.chunks[-1]) >= self.chunk_length: filename = self._dump(self.chunks[-1]) self.chunks[-1] = filename self.chunks.append([]) @@ -54,7 +72,7 @@ class BigArray(list): self.chunks.pop() try: with open(self.chunks[-1], "rb") as fp: - self.chunks[-1] = self._load(fp) + self.chunks[-1] = pickle.load(fp) except IOError, ex: errMsg = "exception occurred while retrieving data " errMsg += "from a temporary file ('%s')" % ex @@ -67,35 +85,13 @@ class BigArray(list): return index return ValueError, "%s is not in list" % value - def _load(self, fp): - retval = [] - unpickler = pickle.Unpickler(fp) - _ = unpickler.load() - if not isinstance(_, list): - retval.append(_) - while True: - try: - retval.append(unpickler.load()) - except EOFError: - break - else: - retval = _ - return retval - def _dump(self, chunk): try: handle, filename = tempfile.mkstemp(prefix=BIGARRAY_TEMP_PREFIX) self.filenames.add(filename) os.close(handle) - try: - with open(filename, "w+b") as fp: - pickle.dump(chunk, fp, pickle.HIGHEST_PROTOCOL) - except MemoryError: - with open(filename, "w+b") as fp: - pickler = pickle.Pickler(fp, pickle.HIGHEST_PROTOCOL) - pickler.fast = True - for value in chunk: - pickler.dump(value) + with open(filename, "w+b") as fp: + pickle.dump(chunk, fp, pickle.HIGHEST_PROTOCOL) return filename except IOError, ex: errMsg = "exception occurred while storing data " @@ -109,7 +105,7 @@ class BigArray(list): if not (self.cache and self.cache.index == index): try: with open(self.chunks[index], "rb") as fp: - self.cache = Cache(index, self._load(fp), False) + self.cache = Cache(index, pickle.load(fp), False) except IOError, ex: errMsg = "exception occurred while retrieving data " errMsg += "from a temporary file ('%s')" % ex @@ -133,8 +129,8 @@ class BigArray(list): def __getitem__(self, y): if y < 0: y += len(self) - index = y / BIGARRAY_CHUNK_LENGTH - offset = y % BIGARRAY_CHUNK_LENGTH + index = y / self.chunk_length + offset = y % self.chunk_length chunk = self.chunks[index] if isinstance(chunk, list): return chunk[offset] @@ -143,8 +139,8 @@ class BigArray(list): return self.cache.data[offset] def __setitem__(self, y, value): - index = y / BIGARRAY_CHUNK_LENGTH - offset = y % BIGARRAY_CHUNK_LENGTH + index = y / self.chunk_length + offset = y % self.chunk_length chunk = self.chunks[index] if isinstance(chunk, list): chunk[offset] = value @@ -161,4 +157,4 @@ class BigArray(list): yield self[i] def __len__(self): - return len(self.chunks[-1]) if len(self.chunks) == 1 else (len(self.chunks) - 1) * BIGARRAY_CHUNK_LENGTH + len(self.chunks[-1]) + return len(self.chunks[-1]) if len(self.chunks) == 1 else (len(self.chunks) - 1) * self.chunk_length + len(self.chunks[-1]) diff --git a/lib/core/settings.py b/lib/core/settings.py index 5d99f6c44..c0a6ae523 100644 --- a/lib/core/settings.py +++ b/lib/core/settings.py @@ -443,8 +443,8 @@ WAF_ATTACK_VECTORS = ( # Used for status representation in dictionary attack phase ROTATING_CHARS = ('\\', '|', '|', '/', '-') -# Chunk length (in items) used by BigArray objects (only last chunk and cached one are held in memory) -BIGARRAY_CHUNK_LENGTH = 1024 +# Approximate chunk length (in bytes) used by BigArray objects (only last chunk and cached one are held in memory) +BIGARRAY_CHUNK_SIZE = 1024 * 1024 # Prefix used for storing dumped chunks in BigArray objects BIGARRAY_TEMP_PREFIX = "sqlmapba-%d-" % os.getpid()