Minor update for an Issue #105

This commit is contained in:
Miroslav Stampar 2012-07-18 14:09:04 +02:00
parent 2656b8fc51
commit 4fc462c4d9

View File

@ -29,10 +29,7 @@ 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 not self.proc_id: 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]
self.fp = open(current, "r") self.fp = open(current, "r")
@ -48,15 +45,13 @@ class Wordlist:
def next(self): def next(self):
retVal = None retVal = None
while True: while True:
self.counter += 1
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()
if not self.proc_count: if not self.proc_count or self.counter % self.proc_count == self.proc_id:
break
self.counter += 1
if self.counter % self.proc_count == self.proc_id:
break break
return retVal return retVal