webssh

commit ec545ec463636457a22c1d6e54e896dc9afe83c6
parent e956c44c1e7869b8645e090ab35b7117f24c7503
Author: Sheng <webmaster0115@gmail.com>
Date:   Mon,  1 Jul 2019 22:10:49 +0800

Refactored PrivateKey

Diffstat:
M
tests/test_handler.py
 | 
33
+++++++++++++++++++++++++++++++++
M
webssh/handler.py
 | 
36
+++++++++++++++++++-----------------

2 files changed, 52 insertions(+), 17 deletions(-)
diff --git a/tests/test_handler.py b/tests/test_handler.py
@@ -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):
 
diff --git a/webssh/handler.py b/webssh/handler.py
@@ -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]
-                if tag:
-                    name = self.tag_to_name.get(tag)
-                    if name:
-                        self.name = name
-                        break
-
-        if not self.name:
+                lst = line.split(' ')
+                if len(lst) == 4:
+                    tag = lst[1]
+                    if tag:
+                        name = tag_to_name.get(tag)
+                        if name:
+                            break
+        return name, len(line_)
+
+    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: