webssh

commit 1e4ece58c89029235d21c1c3954c7b6a6e5aa550
parent babd9bd24f651bd6a956da3bd78440b4b6b4d839
Author: Sheng <webmaster0115@gmail.com>
Date:   Fri, 10 Aug 2018 09:23:25 +0800

Added functions for validating ip and port

Diffstat:
A
tests/test_utils.py
 | 
30
++++++++++++++++++++++++++++++
M
webssh/handler.py
 | 
29
++++++++++++++++-------------
A
webssh/utils.py
 | 
29
+++++++++++++++++++++++++++++

3 files changed, 75 insertions(+), 13 deletions(-)
diff --git a/tests/test_utils.py b/tests/test_utils.py
@@ -0,0 +1,30 @@
+import unittest
+
+from webssh.utils import (is_valid_ipv4_address, is_valid_ipv6_address,
+                          is_valid_port, to_str)
+
+
+class TestUitls(unittest.TestCase):
+
+    def test_to_str(self):
+        b = b'hello'
+        u = u'hello'
+        self.assertEqual(to_str(b), u)
+        self.assertEqual(to_str(u), u)
+
+    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'))
+        self.assertTrue(is_valid_ipv4_address('127.0.0.1'))
+        self.assertTrue(is_valid_ipv4_address(b'127.0.0.1'))
+
+    def test_is_valid_ipv6_address(self):
+        self.assertFalse(is_valid_ipv6_address('abc'))
+        self.assertFalse(is_valid_ipv6_address(b'abc'))
+        self.assertTrue(is_valid_ipv6_address('::1'))
+        self.assertTrue(is_valid_ipv6_address(b'::1'))
+
+    def test_is_valid_port(self):
+        self.assertTrue(is_valid_port(80))
+        self.assertFalse(is_valid_port(0))
+        self.assertFalse(is_valid_port(65536))
diff --git a/webssh/handler.py b/webssh/handler.py
@@ -12,6 +12,8 @@ import tornado.web
 from tornado.ioloop import IOLoop
 from tornado.util import basestring_type
 from webssh.worker import Worker, recycle_worker, workers
+from webssh.utils import (is_valid_ipv4_address, is_valid_ipv6_address,
+                          is_valid_port)
 
 try:
     from concurrent.futures import Future
@@ -40,16 +42,17 @@ class MixinHandler(object):
         ip = self.request.headers.get('X-Real-Ip')
         port = self.request.headers.get('X-Real-Port')
 
-        if ip is None and port is None:
+        if ip is None and port is None:  # suppose the server doesn't use nginx
             return
 
-        try:
-            port = int(port)
-        except (TypeError, ValueError):
-            pass
-        else:
-            if ip:  # does not validate ip and port here
-                return (ip, port)
+        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)
 
         logging.warning('Bad nginx configuration.')
         return False
@@ -101,10 +104,10 @@ class IndexHandler(MixinHandler, tornado.web.RequestHandler):
         try:
             port = int(value)
         except ValueError:
-            port = 0
-
-        if 0 < port < 65536:
-            return port
+            pass
+        else:
+            if is_valid_port(port):
+                return port
 
         raise ValueError('Invalid port {}'.format(value))
 
@@ -135,7 +138,7 @@ class IndexHandler(MixinHandler, tornado.web.RequestHandler):
         except paramiko.SSHException:
             result = None
         else:
-            data = stdout.read().decode()
+            data = stdout.read().decode('utf-8')
             result = parse_encoding(data)
 
         return result if result else 'utf-8'
diff --git a/webssh/utils.py b/webssh/utils.py
@@ -0,0 +1,29 @@
+import ipaddress
+
+
+def to_str(s):
+    if isinstance(s, bytes):
+        return s.decode('utf-8')
+    return s
+
+
+def is_valid_ipv4_address(ipstr):
+    ipstr = to_str(ipstr)
+    try:
+        ipaddress.IPv4Address(ipstr)
+    except ipaddress.AddressValueError:
+        return False
+    return True
+
+
+def is_valid_ipv6_address(ipstr):
+    ipstr = to_str(ipstr)
+    try:
+        ipaddress.IPv6Address(ipstr)
+    except ipaddress.AddressValueError:
+        return False
+    return True
+
+
+def is_valid_port(port):
+    return 0 < port < 65536