webssh

commit 88405eddac6f3dca0be0a628dd941a1c3fc016a7
parent de0fda1ae168d5e1bb860e34ab86eba4604494a9
Author: Sheng <webmaster0115@gmail.com>
Date:   Wed, 10 Oct 2018 08:49:59 +0800

Added to_int to utils

Diffstat:
M
tests/test_utils.py
 | 
12
++++++++++--
M
webssh/handler.py
 | 
23
+++++++----------------
M
webssh/utils.py
 | 
7
+++++++

3 files changed, 24 insertions(+), 18 deletions(-)
diff --git a/tests/test_utils.py b/tests/test_utils.py
@@ -1,7 +1,9 @@
 import unittest
 
-from webssh.utils import (is_valid_ipv4_address, is_valid_ipv6_address,
-                          is_valid_port, is_valid_hostname, to_str, to_bytes)
+from webssh.utils import (
+    is_valid_ipv4_address, is_valid_ipv6_address, is_valid_port,
+    is_valid_hostname, to_str, to_bytes, to_int
+)
 
 
 class TestUitls(unittest.TestCase):
@@ -18,6 +20,12 @@ class TestUitls(unittest.TestCase):
         self.assertEqual(to_bytes(b), b)
         self.assertEqual(to_bytes(u), b)
 
+    def test_to_int(self):
+        self.assertEqual(to_int(''), None)
+        self.assertEqual(to_int(None), None)
+        self.assertEqual(to_int('22'), 22)
+        self.assertEqual(to_int(' 22 '), 22)
+
     def test_is_valid_ipv4_address(self):
         self.assertFalse(is_valid_ipv4_address('127.0.0'))
         self.assertFalse(is_valid_ipv4_address(b'127.0.0'))
diff --git a/webssh/handler.py b/webssh/handler.py
@@ -13,7 +13,7 @@ from tornado.ioloop import IOLoop
 from webssh.settings import swallow_http_errors
 from webssh.utils import (
     is_valid_ipv4_address, is_valid_ipv6_address, is_valid_port,
-    is_valid_hostname, to_bytes, to_str, UnicodeType
+    is_valid_hostname, to_bytes, to_str, to_int, UnicodeType
 )
 from webssh.worker import Worker, recycle_worker, workers
 
@@ -52,13 +52,9 @@ class MixinHandler(object):
             return  # suppose this app doesn't run after an nginx server
 
         if is_valid_ipv4_address(ip) or is_valid_ipv6_address(ip):
-            try:
-                port = int(port)
-            except (TypeError, ValueError):
-                pass
-            else:
-                if is_valid_port(port):
-                    return (ip, port)
+            port = to_int(port)
+            if port and is_valid_port(port):
+                return (ip, port)
 
         logging.warning('Bad nginx configuration.')
         return False
@@ -154,14 +150,9 @@ class IndexHandler(MixinHandler, tornado.web.RequestHandler):
 
     def get_port(self):
         value = self.get_value('port')
-        try:
-            port = int(value)
-        except ValueError:
-            pass
-        else:
-            if is_valid_port(port):
-                return port
-
+        port = to_int(value)
+        if port and is_valid_port(port):
+            return port
         raise InvalidValueError('Invalid port: {}'.format(value))
 
     def lookup_hostname(self, hostname, port):
diff --git a/webssh/utils.py b/webssh/utils.py
@@ -23,6 +23,13 @@ def to_bytes(ustr, encoding='utf-8'):
     return ustr
 
 
+def to_int(string):
+    try:
+        return int(string)
+    except (TypeError, ValueError):
+        pass
+
+
 def is_valid_ipv4_address(ipstr):
     ipstr = to_str(ipstr)
     try: