Using default wordlist

Author: gigkokman

lib/helpers/wordlist_helper.py

@@ -1,6 +1,24 @@
 import sys
+import os
+from lib.helpers.file_helper import get_combined_word_lists
+
+DEFAULT_WORDLIST_FILE = os.path.join(
+    os.path.dirname(os.path.abspath(__file__)),
+    '../..',
+    'wordlists',
+    'virtual-host-scanning.txt'
+)
 
 
 class WordList:
     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 = []
+        stdin_words = self.get_stdin_wordlist()
+        if stdin_words:
+            wordlist.extend(stdin_words)
+        combined = get_combined_word_lists(wordlist_files or DEFAULT_WORDLIST_FILE)
+        wordlist.extend(combined['words'])
+        return wordlist

tests/helpers/test_wordlist_helper.py

@@ -3,15 +3,24 @@
 from mock import patch
 
 from lib.helpers.wordlist_helper import WordList
+from lib.helpers.wordlist_helper import DEFAULT_WORDLIST_FILE
 
 
 class TestWordList(unittest.TestCase):
     def setUp(self):
         self.wordlist = WordList()
+        with open(DEFAULT_WORDLIST_FILE, 'r') as word_file:
+            self.default_wordlist = list(word_file.read().splitlines())
 
     def test_get_wordlist_from_stdin(self):
         stdin_list = ['keyword1', 'keyword1']
+        expected_wordlist = []
+        expected_wordlist.extend(stdin_list)
+        expected_wordlist.extend(self.default_wordlist)
         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_stdin_wordlist(), stdin_list)
+            self.assertEqual(self.wordlist.get_wordlist(), expected_wordlist)
+
+    def test_using_default_wordlist(self):
+        self.assertEqual(self.wordlist.get_wordlist(), self.default_wordlist)