if stdin is provided, ignore default_wordlist

Author: gigkokman

lib/helpers/wordlist_helper.py

@@ -11,21 +11,30 @@
 
 
 class WordList:
+    def __init__(self):
+        self.wordlist = []
+        self.wordlist_types = []
+
     def get_stdin_wordlist(self):
         return list(line for line in sys.stdin.read().splitlines()) \
             if not sys.stdin.isatty() else []
 
     def get_wordlist(self, wordlist_files=None):
-        wordlist = []
-        wordlist_types = []
         stdin_words = self.get_stdin_wordlist()
+        default_wordlist_file = DEFAULT_WORDLIST_FILE
         if stdin_words:
-            wordlist_types.append('stdin')
-            wordlist.extend(stdin_words)
-        combined_files = wordlist_files or DEFAULT_WORDLIST_FILE
+            self.set_words('stdin', stdin_words)
+            default_wordlist_file = None
+
+        combined_files = wordlist_files or default_wordlist_file
         combined = get_combined_word_lists(combined_files)
         if combined:
-            wordlist_types.append('wordlists: {}'.format(
-                ', '.join(combined['file_paths'])))
-            wordlist.extend(combined['words'])
-        return wordlist, wordlist_types
+            words_type = 'wordlists: {}'.format(
+                ', '.join(combined['file_paths']))
+            self.set_words(words_type, combined['words'])
+
+        return self.wordlist, self.wordlist_types
+
+    def set_words(self, words_type, words):
+        self.wordlist_types.append(words_type)
+        self.wordlist.extend(words)

tests/helpers/test_wordlist_helper.py

@@ -15,7 +15,6 @@ def test_get_wordlist_from_stdin(self):
         stdin_wordlist = ['keyword1', 'keyword1']
         expected_wordlist = []
         expected_wordlist.extend(stdin_wordlist)
-        expected_wordlist.extend(self.default_wordlist)
         with patch('lib.helpers.wordlist_helper.WordList.get_stdin_wordlist', return_value=stdin_wordlist):
             wordlist, wordlist_types = self.wordlist.get_wordlist()
             self.assertEqual(wordlist, expected_wordlist)