Return wordlist types

Author: gigkokman

lib/helpers/wordlist_helper.py

@@ -16,9 +16,14 @@ def get_stdin_wordlist(self):
 
     def get_wordlist(self, wordlist_files=None):
         wordlist = []
+        wordlist_types = []
         stdin_words = self.get_stdin_wordlist()
+        print(stdin_words)
         if stdin_words:
+            wordlist_types.append('stdin')
             wordlist.extend(stdin_words)
         combined = get_combined_word_lists(wordlist_files or DEFAULT_WORDLIST_FILE)
-        wordlist.extend(combined['words'])
-        return wordlist
+        if combined:
+            wordlist_types.append('wordlists: {}'.format(', '.join(combined['file_paths'])))
+            wordlist.extend(combined['words'])
+        return wordlist, wordlist_types

tests/helpers/test_wordlist_helper.py

@@ -20,7 +20,9 @@ def test_get_wordlist_from_stdin(self):
         with patch('sys.stdin') as mock_stdin:
             mock_stdin.read = Mock(return_value='\n'.join(stdin_list))
             mock_stdin.isatty = Mock(return_value=False)
-            self.assertEqual(self.wordlist.get_wordlist(), expected_wordlist)
+            wordlist, wordlist_types = self.wordlist.get_wordlist()
+            self.assertEqual(wordlist, expected_wordlist)
 
     def test_using_default_wordlist(self):
-        self.assertEqual(self.wordlist.get_wordlist(), self.default_wordlist)
+        wordlist, wordlist_types = self.wordlist.get_wordlist()
+        self.assertEqual(wordlist, self.default_wordlist)