Patch for an Issue #1095

This commit is contained in:
Miroslav Stampar 2015-01-07 02:04:10 +01:00
parent 5920d16cf6
commit 2030311d50
2 changed files with 28 additions and 6 deletions

View File

@ -54,7 +54,7 @@ class BigArray(list):
self.chunks.pop()
try:
with open(self.chunks[-1], "rb") as fp:
self.chunks[-1] = pickle.load(fp)
self.chunks[-1] = self._load(fp)
except IOError, ex:
errMsg = "exception occurred while retrieving data "
errMsg += "from a temporary file ('%s')" % ex
@ -67,13 +67,35 @@ class BigArray(list):
return index
return ValueError, "%s is not in list" % value
def _dump(self, 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(value, fp, pickle.HIGHEST_PROTOCOL)
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)
return filename
except IOError, ex:
errMsg = "exception occurred while storing data "
@ -87,7 +109,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, pickle.load(fp), False)
self.cache = Cache(index, self._load(fp), False)
except IOError, ex:
errMsg = "exception occurred while retrieving data "
errMsg += "from a temporary file ('%s')" % ex

View File

@ -444,7 +444,7 @@ WAF_ATTACK_VECTORS = (
ROTATING_CHARS = ('\\', '|', '|', '/', '-')
# Chunk length (in items) used by BigArray objects (only last chunk and cached one are held in memory)
BIGARRAY_CHUNK_LENGTH = 4096
BIGARRAY_CHUNK_LENGTH = 1024
# Prefix used for storing dumped chunks in BigArray objects
BIGARRAY_TEMP_PREFIX = "sqlmapba-%d-" % os.getpid()