Encapsulate the argument parsing logic in a class

.gitignore

@@ -17,7 +17,6 @@ dist/
 downloads/
 eggs/
 .eggs/
-lib/
 lib64/
 parts/
 sdist/

VHostScan.py

@@ -9,6 +9,7 @@
 from lib.helpers.output_helper import *
 from lib.helpers.file_helper import get_combined_word_lists, load_random_user_agents
 from lib.core.__version__ import __version__
+from lib.input import cli_argument_parser
 
 
 def print_banner():
@@ -19,33 +20,12 @@ def print_banner():
 
 def main():
     print_banner()
-    parser = ArgumentParser()
-    parser.add_argument("-t",   dest="target_hosts", required=True, help="Set a target range of addresses to target. Ex 10.11.1.1-255" )
-    parser.add_argument("-w",   dest="wordlists", required=False, type=str, help="Set the wordlists to use (default ./wordlists/virtual-host-scanning.txt)", default=False)
-    parser.add_argument("-b",   dest="base_host", required=False, help="Set host to be used during substitution in wordlist (default to TARGET).", default=False)
-    parser.add_argument("-p",   dest="port", required=False, help="Set the port to use (default 80).", default=80)
-    parser.add_argument("-r",   dest="real_port", required=False, help="The real port of the webserver to use in headers when not 80 (see RFC2616 14.23), useful when pivoting through ssh/nc etc (default to PORT).", default=False)
-
-    parser.add_argument('--ignore-http-codes', dest='ignore_http_codes', type=str, help='Comma separated list of http codes to ignore with virtual host scans (default 404).', default='404')
-    parser.add_argument('--ignore-content-length', dest='ignore_content_length', type=int, help='Ignore content lengths of specificed amount (default 0).', default=0)
-    parser.add_argument('--first-hit', dest='first_hit', action='store_true', help='Return first successful result. Only use in scenarios where you are sure no catch-all is configured (such as a CTF).', default=False)
-    parser.add_argument('--unique-depth', dest='unique_depth', type=int, help='Show likely matches of page content that is found x times (default 1).', default=1)
-    parser.add_argument("--ssl", dest="ssl",   action="store_true", help="If set then connections will be made over HTTPS instead of HTTP (default http).", default=False)
-    parser.add_argument("--fuzzy-logic", dest="fuzzy_logic", action="store_true", help="If set then fuzzy match will be performed against unique hosts (default off).", default=False)
-    parser.add_argument("--no-lookups", dest="no_lookup", action="store_true", help="Disable reverse lookups (identifies new targets and appends to wordlist, on by default).", default=False)
-    parser.add_argument("--rate-limit", dest="rate_limit", type=int, help='Amount of time in seconds to delay between each scan (default 0).', default=0)
-    parser.add_argument('--random-agent', dest='random_agent', action='store_true', help='If set, then each scan will use random user-agent from predefined list.', default=False)
-    parser.add_argument('--user-agent', dest='user_agent', type=str, help='Specify a user-agent to use for scans')
-    parser.add_argument("--waf", dest="add_waf_bypass_headers",   action="store_true", help="If set then simple WAF bypass headers will be sent.", default=False)
-    parser.add_argument("-oN",   dest="output_normal", help="Normal output printed to a file when the -oN option is specified with a filename argument." )
-    parser.add_argument("-oJ", dest="output_json", help="JSON output printed to a file when the -oJ option is specified with a filename argument." )
-    parser.add_argument("-", dest="stdin", action="store_true", help="By passing a blank '-' you tell VHostScan to expect input from stdin (pipe).", default=False)
-
-    arguments = parser.parse_args()
-    wordlist = []
+
+    parser = cli_argument_parser()
+    arguments = parser.parse(sys.argv[1:])
 
+    wordlist = []
     word_list_types = []
-
     default_wordlist = "./wordlists/virtual-host-scanning.txt" if not arguments.stdin else None
 
     if arguments.stdin:

lib/core/__version__.py

@@ -2,4 +2,4 @@
 # |V|H|o|s|t|S|c|a|n|  Developed by @codingo_ & @__timk
 # +-+-+-+-+-+-+-+-+-+  https://github.com/codingo/VHostScan
 
-__version__ = '1.6'
+__version__ = '1.6.1'

lib/input.py

@@ -0,0 +1,111 @@
+from argparse import ArgumentParser
+
+class cli_argument_parser(object):
+    def __init__(self):
+        self._parser = self.setup_parser()
+
+    def parse(self, argv):
+        return self._parser.parse_args(argv)
+
+    @staticmethod
+    def setup_parser():
+        parser = ArgumentParser()
+
+        parser.add_argument(
+            '-t', dest='target_hosts', required=True,
+            help='Set a target range of addresses to target. Ex 10.11.1.1-255'
+        ),
+
+        parser.add_argument(
+            '-w', dest='wordlists',
+            help='Set the wordlists to use (default ./wordlists/virtual-host-scanning.txt)'
+        )
+
+        parser.add_argument(
+            '-b', dest='base_host', default=False,
+            help='Set host to be used during substitution in wordlist (default to TARGET).'
+        )
+
+        parser.add_argument(
+            '-p', dest='port', default=80, type=int,
+            help='Set the port to use (default 80).'
+        )
+
+        parser.add_argument(
+            '-r', dest='real_port', type=int, default=False,
+            help='The real port of the webserver to use in headers when not 80 (see RFC2616 14.23), useful when pivoting through ssh/nc etc (default to PORT).'
+        )
+
+        parser.add_argument(
+            '--ignore-http-codes', dest='ignore_http_codes', default='404',
+            help='Comma separated list of http codes to ignore with virtual host scans (default 404).'
+        )
+
+        parser.add_argument(
+            '--ignore-content-length', dest='ignore_content_length', type=int, default=0,
+            help='Ignore content lengths of specificed amount (default 0).'
+        )
+
+        parser.add_argument(
+            '--first-hit', dest='first_hit', action='store_true', default=False,
+            help='Return first successful result. Only use in scenarios where you are sure no catch-all is configured (such as a CTF).'
+        )
+
+        parser.add_argument(
+            '--unique-depth', dest='unique_depth', type=int, default=1,
+            help='Show likely matches of page content that is found x times (default 1).'
+        )
+
+        parser.add_argument(
+            '--ssl', dest='ssl', action='store_true', default=False,
+            help='If set then connections will be made over HTTPS instead of HTTP (default http).'
+        )
+
+        parser.add_argument(
+            '--fuzzy-logic', dest='fuzzy_logic', action='store_true', default=False,
+            help='If set then fuzzy match will be performed against unique hosts (default off).'
+        )
+
+        parser.add_argument(
+            '--no-lookups', dest='no_lookup', action='store_true', default=False,
+            help='Disable reverse lookups (identifies new targets and appends to wordlist, on by default).'
+        )
+
+        parser.add_argument(
+            '--rate-limit', dest='rate_limit', type=int, default=0,
+            help='Amount of time in seconds to delay between each scan (default 0).'
+        )
+
+        parser.add_argument(
+            '--waf', dest='add_waf_bypass_headers', action='store_true', default=False,
+            help='If set then simple WAF bypass headers will be sent.'
+        )
+
+        parser.add_argument(
+            '-', dest='stdin', action='store_true', default=False,
+            help="By passing a blank '-' you tell VHostScan to expect input from stdin (pipe)."
+        )
+
+        output = parser.add_mutually_exclusive_group()
+        output.add_argument(
+            '-oN', dest='output_normal',
+            help='Normal output printed to a file when the -oN option is specified with a filename argument.'
+        )
+
+        output.add_argument(
+            '-oJ', dest='output_json',
+            help='JSON output printed to a file when the -oJ option is specified with a filename argument.'
+        )
+
+        user_agent = parser.add_mutually_exclusive_group()
+        user_agent.add_argument(
+            '--random-agent', dest='random_agent', action='store_true', default=False,
+            help='If set, then each scan will use random user-agent from predefined list.'
+        )
+
+        user_agent.add_argument(
+            '--user-agent', dest='user_agent',
+            help='Specify a user-agent to use for scans'
+        )
+
+        return parser

tests/test_input.py

@@ -0,0 +1,118 @@
+import argparse
+import pytest
+
+from lib.input import cli_argument_parser
+
+def test_parse_arguments_default_value(tmpdir):
+    words = ['word1', 'word2', 'word3']
+    wordlist = tmpdir.mkdir('test_command').join('default')
+    wordlist.write('\n'.join(words))
+
+    argv = ['-t', 'myhost']
+
+    arguments = cli_argument_parser().parse(argv)
+
+    expected_arguments = {
+        'target_hosts': 'myhost',
+        'wordlists': None,
+        'base_host': False,
+        'port': 80,
+        'real_port': False,
+        'ignore_http_codes': '404',
+        'ignore_content_length': 0,
+        'first_hit': False ,
+        'unique_depth': 1,
+        'fuzzy_logic': False,
+        'no_lookup': False,
+        'rate_limit': 0,
+        'random_agent': False,
+        'user_agent': None,
+        'add_waf_bypass_headers': False,
+        'output_normal': None,
+        'output_json': None,
+        'stdin': False,
+        'ssl': False,
+    }
+
+    assert vars(arguments) == expected_arguments
+
+
+def test_parse_arguments_custom_arguments(tmpdir):
+    words = ['some', 'other', 'words']
+    wordlist = tmpdir.mkdir('test_command').join('other_words')
+    wordlist.write('\n'.join(words))
+
+    argv = [
+        '-t', '10.11.1.1',
+        '-w', str(wordlist),
+        '-b', 'myhost',
+        '-p', '8000',
+        '-r', '8001',
+        '--ignore-http-codes', '400,500,302',
+        '--ignore-content-length', '100',
+        '--unique-depth', '5',
+        '--first-hit',
+        '--ssl',
+        '--fuzzy-logic',
+        '--no-lookups',
+        '--rate-limit', '10',
+        '--user-agent', 'some-user-agent',
+        '--waf',
+        '-oN', '/tmp/on',
+        '-',
+    ]
+
+    arguments = cli_argument_parser().parse(argv)
+
+    expected_arguments = {
+        'target_hosts': '10.11.1.1',
+        'wordlists': str(wordlist),
+        'base_host': 'myhost',
+        'port': 8000,
+        'real_port': 8001,
+        'ignore_http_codes': '400,500,302',
+        'ignore_content_length': 100,
+        'first_hit': True,
+        'unique_depth': 5,
+        'ssl': True,
+        'fuzzy_logic': True,
+        'no_lookup': True,
+        'rate_limit': 10,
+        'user_agent': 'some-user-agent',
+        'random_agent': False,
+        'add_waf_bypass_headers': True,
+        'output_normal': '/tmp/on',
+        'output_json': None,
+        'stdin': True,
+    }
+
+    assert vars(arguments) == expected_arguments
+
+def test_parse_arguments_mutually_exclusive_user_agent():
+    argv = [
+        '-t', '10.11.1.1',
+        '--user-agent', 'my-user-agent',
+        '--random-agent',
+    ]
+
+    with pytest.raises(SystemExit):
+        cli_argument_parser().parse(argv)
+
+def test_parse_arguments_mutually_exclusive_output():
+    argv = [
+        '-t', '10.11.1.1',
+        '-oJ',
+        '-oN',
+    ]
+
+    with pytest.raises(SystemExit):
+        cli_argument_parser().parse(argv)
+
+def test_parse_arguments_unknown_argument():
+    argv = [
+        '-t', '10.11.1.1',
+        '-i-do-not-exist',
+    ]
+
+    with pytest.raises(SystemExit):
+        cli_argument_parser().parse(argv)