minor update

stamparm · stamparm · commit 3c31ccd16e77 · 2011-10-26T22:37:04.000Z
diff --git a/lib/core/common.py b/lib/core/common.py
@@ -153,6 +153,7 @@ def __init__(self, filenames):
         self.cursize = 0
         self.custom = []
         self.adjust()
+        self.lock = None
 
     def __iter__(self):
         return self
@@ -162,14 +163,17 @@ def adjust(self):
         if self.index > len(self.filenames):
             raise StopIteration
         elif self.index == len(self.filenames):
-            self.iter = iter(self.custom)
+            if self.custom:
+                self.iter = iter(self.custom)
+            else:
+                raise StopIteration
         else:
             current = self.filenames[self.index]
             infoMsg = "loading dictionary from '%s'" % current
             singleTimeLogMessage(infoMsg)
             self.fp = open(current, "r")
             self.cursize = os.path.getsize(current)
-            self.iter = self.fp.xreadlines()
+            self.iter = iter(self.fp)
 
         self.index += 1
 
@@ -183,11 +187,16 @@ def closeFP(self):
 
     def next(self):
         retVal = None
+        if self.lock:
+            self.lock.acquire()
         try:
             retVal = self.iter.next().rstrip()
         except StopIteration:
             self.adjust()
             retVal = self.iter.next().rstrip()
+        finally:
+            if self.lock:
+                self.lock.release()
         return retVal
 
     def percentage(self):
diff --git a/lib/utils/hash.py b/lib/utils/hash.py
@@ -541,6 +541,9 @@ def dictionaryAttack(attack_dict):
 
                     kb.wordlist = Wordlist(dictPaths)
 
+                    if _multiprocessing:
+                        kb.wordlist.lock = _multiprocessing.Lock()
+
                 except sqlmapFilePathException, msg:
                     warnMsg = "there was a problem while loading dictionaries"
                     warnMsg += " ('%s')" % msg
