Replace assertions with exceptions.

This commit is contained in:
Andris Raugulis 2017-04-05 16:02:40 +03:00
parent 09c2e7b2d5
commit bb122ffe13
2 changed files with 36 additions and 17 deletions

View File

@ -570,8 +570,10 @@ class SSH1(object):
class PublicKeyMessage(object): class PublicKeyMessage(object):
def __init__(self, cookie, skey, hkey, pflags, cmask, amask): def __init__(self, cookie, skey, hkey, pflags, cmask, amask):
# type: (binary_type, Tuple[int, int, int], Tuple[int, int, int], int, int, int) -> None # type: (binary_type, Tuple[int, int, int], Tuple[int, int, int], int, int, int) -> None
assert len(skey) == 3 if len(skey) != 3:
assert len(hkey) == 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.__cookie = cookie
self.__server_key = skey self.__server_key = skey
self.__host_key = hkey self.__host_key = hkey
@ -1192,7 +1194,7 @@ class SSH(object): # pylint: disable=too-few-public-methods
if ssh_prefix not in result: if ssh_prefix not in result:
result[ssh_prefix] = [None, None, None] result[ssh_prefix] = [None, None, None]
prev, push = result[ssh_prefix][i], False prev, push = result[ssh_prefix][i], False
if ((prev is None) or if (prev is None or
(prev < ssh_version and i == 0) or (prev < ssh_version and i == 0) or
(prev > ssh_version and i > 0)): (prev > ssh_version and i > 0)):
push = True push = True

View File

@ -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.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' assert fp.sha256 == 'SHA256:vZdx3mhzbvVJmn08t/ruv8WDhJ9jfKYsCTuSzot+QIs'
def test_pkm_read(self): def _assert_pkm_keys(self, pkm, skey, hkey):
pkm = self.ssh1.PublicKeyMessage.parse(self._pkm_payload()) b, e, m = skey
assert pkm is not None
assert pkm.cookie == b'\x88\x99\xaa\xbb\xcc\xdd\xee\xff'
b, e, m = self._server_key()
assert pkm.server_key_bits == b assert pkm.server_key_bits == b
assert pkm.server_key_public_exponent == e assert pkm.server_key_public_exponent == e
assert pkm.server_key_public_modulus == m 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_bits == b
assert pkm.host_key_public_exponent == e assert pkm.host_key_public_exponent == e
assert pkm.host_key_public_modulus == m 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.protocol_flags == 2
assert pkm.supported_ciphers_mask == 72 assert pkm.supported_ciphers_mask == 72
assert pkm.supported_ciphers == ['3des', 'blowfish'] assert pkm.supported_ciphers == ['3des', 'blowfish']
assert pkm.supported_authentications_mask == 36 assert pkm.supported_authentications_mask == 36
assert pkm.supported_authentications == ['rsa', 'tis'] 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.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' 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): def test_pkm_payload(self):
cookie = b'\x88\x99\xaa\xbb\xcc\xdd\xee\xff' cookie = b'\x88\x99\xaa\xbb\xcc\xdd\xee\xff'
skey = self._server_key() skey, hkey = self._server_key(), self._host_key()
hkey = self._host_key() pflags, cmask, amask = 2, 72, 36
pflags = 2
cmask = 72
amask = 36
pkm1 = self.ssh1.PublicKeyMessage(cookie, skey, hkey, pflags, cmask, amask) pkm1 = self.ssh1.PublicKeyMessage(cookie, skey, hkey, pflags, cmask, amask)
pkm2 = self.ssh1.PublicKeyMessage.parse(self._pkm_payload()) pkm2 = self.ssh1.PublicKeyMessage.parse(self._pkm_payload())
assert pkm1.payload == pkm2.payload assert pkm1.payload == pkm2.payload