Merge pull request #79 from gigkokman/wordlist_helper

Wordlist helper

VHostScan.py

@@ -7,8 +7,8 @@
 from socket import gethostbyaddr
 from lib.core.virtual_host_scanner import *
 from lib.helpers.output_helper import *
-from lib.helpers.file_helper import get_combined_word_lists
 from lib.helpers.file_helper import load_random_user_agents
+from lib.helpers.wordlist_helper import WordList
 from lib.core.__version__ import __version__
 from lib.input import cli_argument_parser
 
@@ -31,21 +31,9 @@ def main():
     parser = cli_argument_parser()
     arguments = parser.parse(sys.argv[1:])
 
-    wordlist = []
-    word_list_types = []
-
-    default_wordlist = DEFAULT_WORDLIST_FILE \
-        if sys.stdin.isatty() else None
-
-    if not sys.stdin.isatty():
-        word_list_types.append('stdin')
-        wordlist.extend(list(line for line in sys.stdin.read().splitlines()))
-
-    combined = get_combined_word_lists(arguments.wordlists or default_wordlist)
-    word_list_types.append('wordlists: {}'.format(
-        ', '.join(combined['file_paths']),
-    ))
-    wordlist.extend(combined['words'])
+    wordlist_helper = WordList()
+    wordlist, wordlist_types = wordlist_helper.get_wordlist(
+        arguments.wordlists)
 
     if len(wordlist) == 0:
         print("[!] No words found in provided wordlists, unable to scan.")
@@ -56,7 +44,7 @@ def main():
         "port {port} and {inputs}".format(
             host=arguments.target_hosts,
             port=arguments.port,
-            inputs=', '.join(word_list_types),
+            inputs=', '.join(wordlist_types),
         )
     )
 

lib/helpers/output_helper.py

@@ -35,8 +35,9 @@ def write_grepable(self, filename):
     def output_normal_likely(self):
         uniques = False
         depth = str(self.scanner.unique_depth)
-        output = "\n[+] Most likely matches with a unique count "\
-                 "of {} or less:".format(depth)
+        output = (
+            "\n[+] Most likely matches with a unique count "
+            "of {} or less:").format(depth)
 
         for p in self.scanner.likely_matches():
             output += "\n\t[>] {}".format(p)
@@ -45,8 +46,9 @@ def output_normal_likely(self):
         if(uniques):
             return output
         else:
-            return "\n[!] No matches with an" \
-                    " unique count of {} or less.".format(depth)
+            return (
+                "\n[!] No matches with an"
+                " unique count of {} or less.").format(depth)
 
     def output_json(self, filename):
         file = file_helper(filename)

lib/helpers/wordlist_helper.py

@@ -0,0 +1,41 @@
+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 __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):
+        default_wordlist_file = DEFAULT_WORDLIST_FILE
+
+        stdin_words = self.get_stdin_wordlist()
+        if stdin_words:
+            self.set_words(words_type='stdin', words=stdin_words)
+            default_wordlist_file = None
+
+        combined_files = wordlist_files or default_wordlist_file
+        combined = get_combined_word_lists(combined_files)
+        if combined:
+            words_type = 'wordlists: {}'.format(
+                ', '.join(combined['file_paths']))
+            self.set_words(words_type=words_type, words=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

@@ -0,0 +1,47 @@
+import unittest
+import pytest
+from mock import patch
+
+from lib.helpers.wordlist_helper import WordList
+from lib.helpers.wordlist_helper import DEFAULT_WORDLIST_FILE
+
+
+@pytest.fixture(scope='class')
+def user_wordlist(request, tmpdir_factory):
+    user_wordlist = ['user-word1', 'user-word2']
+    tmpdir = tmpdir_factory.mktemp('user_wordlist')
+    user_wordlist_file = tmpdir.join('user-wordlist.txt')
+    user_wordlist_file.write('\n'.join(user_wordlist))
+    request.cls.user_wordlist_file = str(user_wordlist_file)
+    request.cls.user_wordlist = user_wordlist
+
+
+@pytest.mark.usefixtures('user_wordlist')
+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_wordlist = ['keyword1', 'keyword1']
+        expected_wordlist = []
+        expected_wordlist.extend(stdin_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)
+
+    def test_get_wordlist_from_stdin_and_wordlist(self):
+        stdin_wordlist = ['keyword1', 'keyword1']
+        expected_wordlist = []
+        expected_wordlist.extend(stdin_wordlist)
+        expected_wordlist.extend(self.user_wordlist)
+        with patch('lib.helpers.wordlist_helper.WordList.get_stdin_wordlist', return_value=stdin_wordlist):
+            wordlist, wordlist_types = self.wordlist.get_wordlist(self.user_wordlist_file)
+            self.assertEqual(wordlist, expected_wordlist)
+
+    def test_using_default_wordlist(self):
+        stdin_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, self.default_wordlist)