diff --git a/ssh-audit.py b/ssh-audit.py index 8cb5b06..461b953 100755 --- a/ssh-audit.py +++ b/ssh-audit.py @@ -169,14 +169,17 @@ class AuditConf(object): usage_cb() if oport is not None: host = args[0] - port = utils.parse_int(oport) else: - s = args[0].split(':') - host = s[0].strip() - if len(s) == 2: - oport, port = s[1], utils.parse_int(s[1]) + mx = re.match(r'^\[([^\]]+)\](?::(.*))?$', args[0]) + if bool(mx): + host, oport = mx.group(1), mx.group(2) else: - oport, port = '22', 22 + s = args[0].split(':') + if len(s) > 2: + host, oport = args[0], '22' + else: + host, oport = s[0], s[1] if len(s) > 1 else '22' + port = utils.parse_int(oport) if not host: usage_cb('host is empty') if port <= 0 or port > 65535: diff --git a/test/test_auditconf.py b/test/test_auditconf.py index 259a881..a901299 100644 --- a/test/test_auditconf.py +++ b/test/test_auditconf.py @@ -148,6 +148,14 @@ class TestAuditConf(object): self._test_conf(conf, host='localhost', port=2222) conf = c('-p 2222 localhost') self._test_conf(conf, host='localhost', port=2222) + conf = c('2001:4860:4860::8888') + self._test_conf(conf, host='2001:4860:4860::8888') + conf = c('[2001:4860:4860::8888]:22') + self._test_conf(conf, host='2001:4860:4860::8888') + conf = c('[2001:4860:4860::8888]:2222') + self._test_conf(conf, host='2001:4860:4860::8888', port=2222) + conf = c('-p 2222 2001:4860:4860::8888') + self._test_conf(conf, host='2001:4860:4860::8888', port=2222) with pytest.raises(SystemExit): conf = c('localhost:') with pytest.raises(SystemExit):