jinql
/
webssh
mirror of https://github.com/huashengdun/webssh
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

Refactored PrivateKey

pull/75/head
Sheng 2019-07-01 22:10:49 +08:00
parent e956c44c1e
commit ec545ec463
2 changed files with 51 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

33
tests/test_handler.py
View File

@ -196,6 +196,39 @@ class TestPrivateKey(unittest.TestCase):
password = 'abc123'
self._test_with_encrypted_key(fname, password, paramiko.Ed25519Key)
def test_parse_name(self):
key = u'-----BEGIN PRIVATE KEY-----'
pk = PrivateKey(key)
name, _ = pk.parse_name(pk.iostr, pk.tag_to_name)
self.assertIsNone(name)
key = u'-----BEGIN xxx PRIVATE KEY-----'
pk = PrivateKey(key)
name, _ = pk.parse_name(pk.iostr, pk.tag_to_name)
self.assertIsNone(name)
key = u'-----BEGIN RSA PRIVATE KEY-----'
pk = PrivateKey(key)
name, _ = pk.parse_name(pk.iostr, pk.tag_to_name)
self.assertIsNone(name)
key = u'-----BEGIN RSA PRIVATE KEY-----'
pk = PrivateKey(key)
name, _ = pk.parse_name(pk.iostr, pk.tag_to_name)
self.assertIsNone(name)
key = u'-----BEGIN RSA PRIVATE KEY-----'
pk = PrivateKey(key)
name, _ = pk.parse_name(pk.iostr, pk.tag_to_name)
self.assertIsNone(name)
for tag, to_name in PrivateKey.tag_to_name.items():
key = u'-----BEGIN {} PRIVATE KEY----- \r\n'.format(tag)
pk = PrivateKey(key)
name, length = pk.parse_name(pk.iostr, pk.tag_to_name)
self.assertEqual(name, to_name)
self.assertEqual(length, len(key))
class TestWsockHandler(unittest.TestCase):

28
webssh/handler.py
View File

@ -43,7 +43,6 @@ class InvalidValueError(Exception):
class PrivateKey(object):
max_length = 16384 # rough number
name = None
tag_to_name = {
'RSA': 'RSA',
@ -63,29 +62,32 @@ class PrivateKey(object):
if len(self.privatekey) > self.max_length:
raise InvalidValueError('Invalid key length.')
def parse_name(self):
for line_orig in self.iostr:
line = line_orig.strip()
def parse_name(self, iostr, tag_to_name):
name = None
for line_ in iostr:
line = line_.strip()
if line and line.startswith('-----BEGIN ') and \
line.endswith(' PRIVATE KEY-----'):
tag = line.split(' ', 2)[1]
lst = line.split(' ')
if len(lst) == 4:
tag = lst[1]
if tag:
name = self.tag_to_name.get(tag)
name = tag_to_name.get(tag)
if name:
self.name = name
break
return name, len(line_)
if not self.name:
def get_pkey_obj(self):
name, length = self.parse_name(self.iostr, self.tag_to_name)
if not name:
raise InvalidValueError('Invalid key {}.'.format(self.filename))
offset = self.iostr.tell() - len(line_orig)
offset = self.iostr.tell() - length
self.iostr.seek(offset)
logging.debug('Reset offset to {}.'.format(offset))
def get_pkey_obj(self):
self.parse_name()
logging.info('Parsing {} key'.format(self.name))
pkeycls = getattr(paramiko, self.name+'Key')
logging.info('Parsing {} key'.format(name))
pkeycls = getattr(paramiko, name+'Key')
password = to_bytes(self.password) if self.password else None
try: