Correct IPv6 parsing in command-line. Fixes #26.

arthepsy · May 5, 2017 · d8eb46d · d8eb46d
1 parent 96d442e
commit d8eb46d
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 6 deletions.
diff --git a/ssh-audit.py b/ssh-audit.py
@@ -169,14 +169,17 @@ def from_cmdline(cls, args, usage_cb):
 			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
@@ -148,6 +148,14 @@ def test_audit_conf_cmdline(self):
 		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):