Skip to content

Commit 515cb06

Browse files
codingocodingo
authored
Merge pull request #64 from diogoosorio/command-argument-class
Encapsulate the argument parsing logic in a class
2 parents 7350602 + 4c13b17 commit 515cb06

File tree

5 files changed

+235
-27
lines changed

5 files changed

+235
-27
lines changed

‎.gitignore

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,6 @@ dist/
1717
downloads/
1818
eggs/
1919
.eggs/
20-
lib/
2120
lib64/
2221
parts/
2322
sdist/

‎VHostScan.py

Lines changed: 5 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
from lib.helpers.output_helper import *
1010
from lib.helpers.file_helper import get_combined_word_lists, load_random_user_agents
1111
from lib.core.__version__ import __version__
12+
from lib.input import cli_argument_parser
1213

1314

1415
def print_banner():
@@ -19,33 +20,12 @@ def print_banner():
1920

2021
def main():
2122
print_banner()
22-
parser = ArgumentParser()
23-
parser.add_argument("-t", dest="target_hosts", required=True, help="Set a target range of addresses to target. Ex 10.11.1.1-255" )
24-
parser.add_argument("-w", dest="wordlists", required=False, type=str, help="Set the wordlists to use (default ./wordlists/virtual-host-scanning.txt)", default=False)
25-
parser.add_argument("-b", dest="base_host", required=False, help="Set host to be used during substitution in wordlist (default to TARGET).", default=False)
26-
parser.add_argument("-p", dest="port", required=False, help="Set the port to use (default 80).", default=80)
27-
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)
28-
29-
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')
30-
parser.add_argument('--ignore-content-length', dest='ignore_content_length', type=int, help='Ignore content lengths of specificed amount (default 0).', default=0)
31-
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)
32-
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)
33-
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)
34-
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)
35-
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)
36-
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)
37-
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)
38-
parser.add_argument('--user-agent', dest='user_agent', type=str, help='Specify a user-agent to use for scans')
39-
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)
40-
parser.add_argument("-oN", dest="output_normal", help="Normal output printed to a file when the -oN option is specified with a filename argument." )
41-
parser.add_argument("-oJ", dest="output_json", help="JSON output printed to a file when the -oJ option is specified with a filename argument." )
42-
parser.add_argument("-", dest="stdin", action="store_true", help="By passing a blank '-' you tell VHostScan to expect input from stdin (pipe).", default=False)
43-
44-
arguments = parser.parse_args()
45-
wordlist = []
23+
24+
parser = cli_argument_parser()
25+
arguments = parser.parse(sys.argv[1:])
4626

27+
wordlist = []
4728
word_list_types = []
48-
4929
default_wordlist = "./wordlists/virtual-host-scanning.txt" if not arguments.stdin else None
5030

5131
if arguments.stdin:

‎lib/core/__version__.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,4 +2,4 @@
22
# |V|H|o|s|t|S|c|a|n| Developed by @codingo_ & @__timk
33
# +-+-+-+-+-+-+-+-+-+ https://github.com/codingo/VHostScan
44

5-
__version__ = '1.6'
5+
__version__ = '1.6.1'

‎lib/input.py

Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,111 @@
1+
from argparse import ArgumentParser
2+
3+
class cli_argument_parser(object):
4+
def __init__(self):
5+
self._parser = self.setup_parser()
6+
7+
def parse(self, argv):
8+
return self._parser.parse_args(argv)
9+
10+
@staticmethod
11+
def setup_parser():
12+
parser = ArgumentParser()
13+
14+
parser.add_argument(
15+
'-t', dest='target_hosts', required=True,
16+
help='Set a target range of addresses to target. Ex 10.11.1.1-255'
17+
),
18+
19+
parser.add_argument(
20+
'-w', dest='wordlists',
21+
help='Set the wordlists to use (default ./wordlists/virtual-host-scanning.txt)'
22+
)
23+
24+
parser.add_argument(
25+
'-b', dest='base_host', default=False,
26+
help='Set host to be used during substitution in wordlist (default to TARGET).'
27+
)
28+
29+
parser.add_argument(
30+
'-p', dest='port', default=80, type=int,
31+
help='Set the port to use (default 80).'
32+
)
33+
34+
parser.add_argument(
35+
'-r', dest='real_port', type=int, default=False,
36+
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).'
37+
)
38+
39+
parser.add_argument(
40+
'--ignore-http-codes', dest='ignore_http_codes', default='404',
41+
help='Comma separated list of http codes to ignore with virtual host scans (default 404).'
42+
)
43+
44+
parser.add_argument(
45+
'--ignore-content-length', dest='ignore_content_length', type=int, default=0,
46+
help='Ignore content lengths of specificed amount (default 0).'
47+
)
48+
49+
parser.add_argument(
50+
'--first-hit', dest='first_hit', action='store_true', default=False,
51+
help='Return first successful result. Only use in scenarios where you are sure no catch-all is configured (such as a CTF).'
52+
)
53+
54+
parser.add_argument(
55+
'--unique-depth', dest='unique_depth', type=int, default=1,
56+
help='Show likely matches of page content that is found x times (default 1).'
57+
)
58+
59+
parser.add_argument(
60+
'--ssl', dest='ssl', action='store_true', default=False,
61+
help='If set then connections will be made over HTTPS instead of HTTP (default http).'
62+
)
63+
64+
parser.add_argument(
65+
'--fuzzy-logic', dest='fuzzy_logic', action='store_true', default=False,
66+
help='If set then fuzzy match will be performed against unique hosts (default off).'
67+
)
68+
69+
parser.add_argument(
70+
'--no-lookups', dest='no_lookup', action='store_true', default=False,
71+
help='Disable reverse lookups (identifies new targets and appends to wordlist, on by default).'
72+
)
73+
74+
parser.add_argument(
75+
'--rate-limit', dest='rate_limit', type=int, default=0,
76+
help='Amount of time in seconds to delay between each scan (default 0).'
77+
)
78+
79+
parser.add_argument(
80+
'--waf', dest='add_waf_bypass_headers', action='store_true', default=False,
81+
help='If set then simple WAF bypass headers will be sent.'
82+
)
83+
84+
parser.add_argument(
85+
'-', dest='stdin', action='store_true', default=False,
86+
help="By passing a blank '-' you tell VHostScan to expect input from stdin (pipe)."
87+
)
88+
89+
output = parser.add_mutually_exclusive_group()
90+
output.add_argument(
91+
'-oN', dest='output_normal',
92+
help='Normal output printed to a file when the -oN option is specified with a filename argument.'
93+
)
94+
95+
output.add_argument(
96+
'-oJ', dest='output_json',
97+
help='JSON output printed to a file when the -oJ option is specified with a filename argument.'
98+
)
99+
100+
user_agent = parser.add_mutually_exclusive_group()
101+
user_agent.add_argument(
102+
'--random-agent', dest='random_agent', action='store_true', default=False,
103+
help='If set, then each scan will use random user-agent from predefined list.'
104+
)
105+
106+
user_agent.add_argument(
107+
'--user-agent', dest='user_agent',
108+
help='Specify a user-agent to use for scans'
109+
)
110+
111+
return parser

‎tests/test_input.py

Lines changed: 118 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,118 @@
1+
import argparse
2+
import pytest
3+
4+
from lib.input import cli_argument_parser
5+
6+
def test_parse_arguments_default_value(tmpdir):
7+
words = ['word1', 'word2', 'word3']
8+
wordlist = tmpdir.mkdir('test_command').join('default')
9+
wordlist.write('\n'.join(words))
10+
11+
argv = ['-t', 'myhost']
12+
13+
arguments = cli_argument_parser().parse(argv)
14+
15+
expected_arguments = {
16+
'target_hosts': 'myhost',
17+
'wordlists': None,
18+
'base_host': False,
19+
'port': 80,
20+
'real_port': False,
21+
'ignore_http_codes': '404',
22+
'ignore_content_length': 0,
23+
'first_hit': False ,
24+
'unique_depth': 1,
25+
'fuzzy_logic': False,
26+
'no_lookup': False,
27+
'rate_limit': 0,
28+
'random_agent': False,
29+
'user_agent': None,
30+
'add_waf_bypass_headers': False,
31+
'output_normal': None,
32+
'output_json': None,
33+
'stdin': False,
34+
'ssl': False,
35+
}
36+
37+
assert vars(arguments) == expected_arguments
38+
39+
40+
def test_parse_arguments_custom_arguments(tmpdir):
41+
words = ['some', 'other', 'words']
42+
wordlist = tmpdir.mkdir('test_command').join('other_words')
43+
wordlist.write('\n'.join(words))
44+
45+
argv = [
46+
'-t', '10.11.1.1',
47+
'-w', str(wordlist),
48+
'-b', 'myhost',
49+
'-p', '8000',
50+
'-r', '8001',
51+
'--ignore-http-codes', '400,500,302',
52+
'--ignore-content-length', '100',
53+
'--unique-depth', '5',
54+
'--first-hit',
55+
'--ssl',
56+
'--fuzzy-logic',
57+
'--no-lookups',
58+
'--rate-limit', '10',
59+
'--user-agent', 'some-user-agent',
60+
'--waf',
61+
'-oN', '/tmp/on',
62+
'-',
63+
]
64+
65+
arguments = cli_argument_parser().parse(argv)
66+
67+
expected_arguments = {
68+
'target_hosts': '10.11.1.1',
69+
'wordlists': str(wordlist),
70+
'base_host': 'myhost',
71+
'port': 8000,
72+
'real_port': 8001,
73+
'ignore_http_codes': '400,500,302',
74+
'ignore_content_length': 100,
75+
'first_hit': True,
76+
'unique_depth': 5,
77+
'ssl': True,
78+
'fuzzy_logic': True,
79+
'no_lookup': True,
80+
'rate_limit': 10,
81+
'user_agent': 'some-user-agent',
82+
'random_agent': False,
83+
'add_waf_bypass_headers': True,
84+
'output_normal': '/tmp/on',
85+
'output_json': None,
86+
'stdin': True,
87+
}
88+
89+
assert vars(arguments) == expected_arguments
90+
91+
def test_parse_arguments_mutually_exclusive_user_agent():
92+
argv = [
93+
'-t', '10.11.1.1',
94+
'--user-agent', 'my-user-agent',
95+
'--random-agent',
96+
]
97+
98+
with pytest.raises(SystemExit):
99+
cli_argument_parser().parse(argv)
100+
101+
def test_parse_arguments_mutually_exclusive_output():
102+
argv = [
103+
'-t', '10.11.1.1',
104+
'-oJ',
105+
'-oN',
106+
]
107+
108+
with pytest.raises(SystemExit):
109+
cli_argument_parser().parse(argv)
110+
111+
def test_parse_arguments_unknown_argument():
112+
argv = [
113+
'-t', '10.11.1.1',
114+
'-i-do-not-exist',
115+
]
116+
117+
with pytest.raises(SystemExit):
118+
cli_argument_parser().parse(argv)

0 commit comments

Comments
 (0)