Patch for an Issue #979

This commit is contained in:
Miroslav Stampar 2014-12-01 00:29:25 +01:00
parent 56b6bf72f4
commit f71a65a9a0
3 changed files with 16 additions and 14 deletions

View File

@ -15,6 +15,7 @@ import tempfile
from lib.core.exception import SqlmapSystemException from lib.core.exception import SqlmapSystemException
from lib.core.settings import BIGARRAY_CHUNK_LENGTH from lib.core.settings import BIGARRAY_CHUNK_LENGTH
from lib.core.settings import BIGARRAY_TEMP_PREFIX
class Cache(object): class Cache(object):
""" """
@ -35,14 +36,12 @@ class BigArray(list):
self.chunks = [[]] self.chunks = [[]]
self.cache = None self.cache = None
self.filenames = set() self.filenames = set()
self.protected = False
self._os_remove = os.remove self._os_remove = os.remove
def append(self, value): def append(self, value):
self.chunks[-1].append(value) self.chunks[-1].append(value)
if len(self.chunks[-1]) >= BIGARRAY_CHUNK_LENGTH: if len(self.chunks[-1]) >= BIGARRAY_CHUNK_LENGTH:
filename = self._dump(self.chunks[-1]) filename = self._dump(self.chunks[-1])
del(self.chunks[-1][:])
self.chunks[-1] = filename self.chunks[-1] = filename
self.chunks.append([]) self.chunks.append([])
@ -65,7 +64,7 @@ class BigArray(list):
def _dump(self, value): def _dump(self, value):
try: try:
handle, filename = tempfile.mkstemp(prefix="sqlmapba-") handle, filename = tempfile.mkstemp(prefix=BIGARRAY_TEMP_PREFIX)
self.filenames.add(filename) self.filenames.add(filename)
os.close(handle) os.close(handle)
with open(filename, "w+b") as fp: with open(filename, "w+b") as fp:
@ -85,12 +84,11 @@ class BigArray(list):
self.cache = Cache(index, pickle.load(fp), False) self.cache = Cache(index, pickle.load(fp), False)
def __getstate__(self): def __getstate__(self):
self.protected = True return self.chunks, self.filenames
return self.chunks, self.filenames, self.protected
def __setstate__(self, state): def __setstate__(self, state):
self.__init__() self.__init__()
self.chunks, self.filenames, self.protected = state self.chunks, self.filenames = state
def __getslice__(self, i, j): def __getslice__(self, i, j):
retval = BigArray() retval = BigArray()
@ -132,11 +130,3 @@ class BigArray(list):
def __len__(self): 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) * BIGARRAY_CHUNK_LENGTH + len(self.chunks[-1])
def __del__(self):
if not self.protected:
for filename in self.filenames:
try:
self._os_remove(filename)
except:
pass

View File

@ -441,6 +441,9 @@ ROTATING_CHARS = ('\\', '|', '|', '/', '-')
# Chunk length (in items) used by BigArray objects (only last chunk and cached one are held in memory) # 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 = 4096
# Prefix used for storing dumped chunks in BigArray objects
BIGARRAY_TEMP_PREFIX = "sqlmapba-%d-" % os.getpid()
# Only console display last n table rows # Only console display last n table rows
TRIM_STDOUT_DUMP_SIZE = 256 TRIM_STDOUT_DUMP_SIZE = 256

View File

@ -6,11 +6,13 @@ See the file 'doc/COPYING' for copying permission
""" """
import bdb import bdb
import glob
import inspect import inspect
import logging import logging
import os import os
import re import re
import sys import sys
import tempfile
import time import time
import traceback import traceback
import warnings import warnings
@ -42,6 +44,7 @@ from lib.core.exception import SqlmapUserQuitException
from lib.core.option import initOptions from lib.core.option import initOptions
from lib.core.option import init from lib.core.option import init
from lib.core.profiling import profile from lib.core.profiling import profile
from lib.core.settings import BIGARRAY_TEMP_PREFIX
from lib.core.settings import LEGAL_DISCLAIMER from lib.core.settings import LEGAL_DISCLAIMER
from lib.core.testing import smokeTest from lib.core.testing import smokeTest
from lib.core.testing import liveTest from lib.core.testing import liveTest
@ -151,6 +154,12 @@ def main():
if conf.get("showTime"): if conf.get("showTime"):
dataToStdout("\n[*] shutting down at %s\n\n" % time.strftime("%X"), forceOutput=True) dataToStdout("\n[*] shutting down at %s\n\n" % time.strftime("%X"), forceOutput=True)
for filename in glob.glob("%s*" % os.path.join(tempfile.gettempdir(), BIGARRAY_TEMP_PREFIX)):
try:
os.remove(filename)
except:
pass
kb.threadContinue = False kb.threadContinue = False
kb.threadException = True kb.threadException = True