minor update

This commit is contained in:
Miroslav Stampar 2011-10-26 22:37:04 +00:00
parent 9d31230d5e
commit 3c31ccd16e
2 changed files with 14 additions and 2 deletions

View File

@ -153,6 +153,7 @@ class Wordlist:
self.cursize = 0 self.cursize = 0
self.custom = [] self.custom = []
self.adjust() self.adjust()
self.lock = None
def __iter__(self): def __iter__(self):
return self return self
@ -162,14 +163,17 @@ class Wordlist:
if self.index > len(self.filenames): if self.index > len(self.filenames):
raise StopIteration raise StopIteration
elif self.index == len(self.filenames): elif self.index == len(self.filenames):
if self.custom:
self.iter = iter(self.custom) self.iter = iter(self.custom)
else:
raise StopIteration
else: else:
current = self.filenames[self.index] current = self.filenames[self.index]
infoMsg = "loading dictionary from '%s'" % current infoMsg = "loading dictionary from '%s'" % current
singleTimeLogMessage(infoMsg) singleTimeLogMessage(infoMsg)
self.fp = open(current, "r") self.fp = open(current, "r")
self.cursize = os.path.getsize(current) self.cursize = os.path.getsize(current)
self.iter = self.fp.xreadlines() self.iter = iter(self.fp)
self.index += 1 self.index += 1
@ -183,11 +187,16 @@ class Wordlist:
def next(self): def next(self):
retVal = None retVal = None
if self.lock:
self.lock.acquire()
try: try:
retVal = self.iter.next().rstrip() retVal = self.iter.next().rstrip()
except StopIteration: except StopIteration:
self.adjust() self.adjust()
retVal = self.iter.next().rstrip() retVal = self.iter.next().rstrip()
finally:
if self.lock:
self.lock.release()
return retVal return retVal
def percentage(self): def percentage(self):

View File

@ -541,6 +541,9 @@ def dictionaryAttack(attack_dict):
kb.wordlist = Wordlist(dictPaths) kb.wordlist = Wordlist(dictPaths)
if _multiprocessing:
kb.wordlist.lock = _multiprocessing.Lock()
except sqlmapFilePathException, msg: except sqlmapFilePathException, msg:
warnMsg = "there was a problem while loading dictionaries" warnMsg = "there was a problem while loading dictionaries"
warnMsg += " ('%s')" % msg warnMsg += " ('%s')" % msg