From bb122ffe139ced0825cbb3965ee1cf5166ea4b29 Mon Sep 17 00:00:00 2001
From: Andris Raugulis <moo@arthepsy.eu>
Date: Wed, 5 Apr 2017 16:02:40 +0300
Subject: [PATCH] Replace assertions with exceptions.

---
 ssh-audit.py      | 12 +++++++-----
 test/test_ssh1.py | 41 +++++++++++++++++++++++++++++------------
 2 files changed, 36 insertions(+), 17 deletions(-)

diff --git a/ssh-audit.py b/ssh-audit.py
index 9fcd04c..50a7d0b 100755
--- a/ssh-audit.py
+++ b/ssh-audit.py
@@ -570,8 +570,10 @@ class SSH1(object):
 	class PublicKeyMessage(object):
 		def __init__(self, cookie, skey, hkey, pflags, cmask, amask):
 			# type: (binary_type, Tuple[int, int, int], Tuple[int, int, int], int, int, int) -> None
-			assert len(skey) == 3
-			assert len(hkey) == 3
+			if len(skey) != 3:
+				raise ValueError('invalid server key pair: {0}'.format(skey))
+			if len(hkey) != 3:
+				raise ValueError('invalid host key pair: {0}'.format(hkey))
 			self.__cookie = cookie
 			self.__server_key = skey
 			self.__host_key = hkey
@@ -1192,9 +1194,9 @@ class SSH(object):  # pylint: disable=too-few-public-methods
 					if ssh_prefix not in result:
 						result[ssh_prefix] = [None, None, None]
 					prev, push = result[ssh_prefix][i], False
-					if ((prev is None) or
-					    (prev < ssh_version and i == 0) or
-					    (prev > ssh_version and i > 0)):
+					if (prev is None or
+					   (prev < ssh_version and i == 0) or
+					   (prev > ssh_version and i > 0)):
 						push = True
 					if push:
 						result[ssh_prefix][i] = ssh_version
diff --git a/test/test_ssh1.py b/test/test_ssh1.py
index 0f62983..f18e4be 100644
--- a/test/test_ssh1.py
+++ b/test/test_ssh1.py
@@ -66,34 +66,51 @@ class TestSSH1(object):
 		assert fp.md5 == 'MD5:9d:26:f8:39:fc:20:9d:9b:ca:cc:4a:0f:e1:93:f5:96'
 		assert fp.sha256 == 'SHA256:vZdx3mhzbvVJmn08t/ruv8WDhJ9jfKYsCTuSzot+QIs'
 	
-	def test_pkm_read(self):
-		pkm = self.ssh1.PublicKeyMessage.parse(self._pkm_payload())
-		assert pkm is not None
-		assert pkm.cookie == b'\x88\x99\xaa\xbb\xcc\xdd\xee\xff'
-		b, e, m = self._server_key()
+	def _assert_pkm_keys(self, pkm, skey, hkey):
+		b, e, m = skey
 		assert pkm.server_key_bits == b
 		assert pkm.server_key_public_exponent == e
 		assert pkm.server_key_public_modulus == m
-		b, e, m = self._host_key()
+		b, e, m = hkey
 		assert pkm.host_key_bits == b
 		assert pkm.host_key_public_exponent == e
 		assert pkm.host_key_public_modulus == m
-		fp = self.ssh.Fingerprint(pkm.host_key_fingerprint_data)
+	
+	def _assert_pkm_fields(self, pkm, skey, hkey):
+		assert pkm is not None
+		assert pkm.cookie == b'\x88\x99\xaa\xbb\xcc\xdd\xee\xff'
+		self._assert_pkm_keys(pkm, skey, hkey)
 		assert pkm.protocol_flags == 2
 		assert pkm.supported_ciphers_mask == 72
 		assert pkm.supported_ciphers == ['3des', 'blowfish']
 		assert pkm.supported_authentications_mask == 36
 		assert pkm.supported_authentications == ['rsa', 'tis']
+		fp = self.ssh.Fingerprint(pkm.host_key_fingerprint_data)
 		assert fp.md5 == 'MD5:9d:26:f8:39:fc:20:9d:9b:ca:cc:4a:0f:e1:93:f5:96'
 		assert fp.sha256 == 'SHA256:vZdx3mhzbvVJmn08t/ruv8WDhJ9jfKYsCTuSzot+QIs'
 	
+	def test_pkm_init(self):
+		cookie = b'\x88\x99\xaa\xbb\xcc\xdd\xee\xff'
+		pflags, cmask, amask = 2, 72, 36
+		skey, hkey = self._server_key(), self._host_key()
+		pkm = self.ssh1.PublicKeyMessage(cookie, skey, hkey, pflags, cmask, amask)
+		self._assert_pkm_fields(pkm, skey, hkey)
+		for skey2 in ([], [0], [0,1], [0,1,2,3]):
+			with pytest.raises(ValueError):
+				pkm = self.ssh1.PublicKeyMessage(cookie, skey2, hkey, pflags, cmask, amask)
+		for hkey2 in ([], [0], [0,1], [0,1,2,3]):
+			with pytest.raises(ValueError):
+				print(hkey2)
+				pkm = self.ssh1.PublicKeyMessage(cookie, skey, hkey2, pflags, cmask, amask)
+	
+	def test_pkm_read(self):
+		pkm = self.ssh1.PublicKeyMessage.parse(self._pkm_payload())
+		self._assert_pkm_fields(pkm, self._server_key(), self._host_key())
+	
 	def test_pkm_payload(self):
 		cookie = b'\x88\x99\xaa\xbb\xcc\xdd\xee\xff'
-		skey = self._server_key()
-		hkey = self._host_key()
-		pflags = 2
-		cmask = 72
-		amask = 36
+		skey, hkey = self._server_key(), self._host_key()
+		pflags, cmask, amask = 2, 72, 36
 		pkm1 = self.ssh1.PublicKeyMessage(cookie, skey, hkey, pflags, cmask, amask)
 		pkm2 = self.ssh1.PublicKeyMessage.parse(self._pkm_payload())
 		assert pkm1.payload == pkm2.payload