minor refactoring

lib/core/wordlist.py

@@ -0,0 +1,71 @@
+#!/usr/bin/env python
+
+"""
+$Id$
+
+Copyright (c) 2006-2012 sqlmap developers (http://www.sqlmap.org/)
+See the file 'doc/COPYING' for copying permission
+"""
+
+from lib.core.common import singleTimeLogMessage
+
+class Wordlist:
+    """
+    Iterator for looping over a large dictionaries
+    """
+
+    def __init__(self, filenames):
+        self.filenames = filenames
+        self.fp = None
+        self.index = 0
+        self.iter = None
+        self.custom = []
+        self.adjust()
+        self.lock = None
+
+    def __iter__(self):
+        return self
+
+    def adjust(self):
+        self.closeFP()
+        if self.index > len(self.filenames):
+            raise StopIteration
+        elif self.index == len(self.filenames):
+            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.iter = iter(self.fp)
+
+        self.index += 1
+
+    def append(self, value):
+        self.custom.append(value)
+        
+    def closeFP(self):
+        if self.fp:
+            self.fp.close()
+            self.fp = None
+
+    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 rewind(self):
+        self.index = 0
+        self.adjust()