From d8eb46d7669aa92ad08ad4e714cdab88ff14496f Mon Sep 17 00:00:00 2001
From: Andris Raugulis <moo@arthepsy.eu>
Date: Fri, 5 May 2017 14:12:45 +0300
Subject: [PATCH] Correct IPv6 parsing in command-line. Fixes #26.

---
 ssh-audit.py           | 15 +++++++++------
 test/test_auditconf.py |  8 ++++++++
 2 files changed, 17 insertions(+), 6 deletions(-)

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):