webssh

commit 66ebe2ceb24382045929fd7becf657070201de6c
parent 9aebb6e4e30b92bfbaf2111e50f5b849c222bb5c
Author: Sheng <webmaster0115@gmail.com>
Date:   Sat, 13 Oct 2018 09:47:15 +0800

Use general method to validate ipaddress

Diffstat:
M
tests/test_utils.py
 | 
24
+++++++++++-------------
M
webssh/handler.py
 | 
7
+++----
M
webssh/utils.py
 | 
15
+++------------

3 files changed, 17 insertions(+), 29 deletions(-)
diff --git a/tests/test_utils.py b/tests/test_utils.py
@@ -1,8 +1,8 @@
 import unittest
 
 from webssh.utils import (
-    is_valid_ipv4_address, is_valid_ipv6_address, is_valid_port,
-    is_valid_hostname, to_str, to_bytes, to_int
+    is_valid_ip_address, is_valid_port, is_valid_hostname,
+    to_str, to_bytes, to_int
 )
 
 
@@ -26,17 +26,15 @@ class TestUitls(unittest.TestCase):
         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'))
-        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_ip_address(self):
+        self.assertFalse(is_valid_ip_address('127.0.0'))
+        self.assertFalse(is_valid_ip_address(b'127.0.0'))
+        self.assertTrue(is_valid_ip_address('127.0.0.1'))
+        self.assertTrue(is_valid_ip_address(b'127.0.0.1'))
+        self.assertFalse(is_valid_ip_address('abc'))
+        self.assertFalse(is_valid_ip_address(b'abc'))
+        self.assertTrue(is_valid_ip_address('::1'))
+        self.assertTrue(is_valid_ip_address(b'::1'))
 
     def test_is_valid_port(self):
         self.assertTrue(is_valid_port(80))
diff --git a/webssh/handler.py b/webssh/handler.py
@@ -12,8 +12,8 @@ import tornado.web
 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, to_int, UnicodeType
+    is_valid_ip_address, is_valid_port, is_valid_hostname,
+    to_bytes, to_str, to_int, UnicodeType
 )
 from webssh.worker import Worker, recycle_worker, workers
 
@@ -149,8 +149,7 @@ class IndexHandler(MixinHandler, tornado.web.RequestHandler):
 
     def get_hostname(self):
         value = self.get_value('hostname')
-        if not (is_valid_hostname(value) | is_valid_ipv4_address(value) |
-                is_valid_ipv6_address(value)):
+        if not (is_valid_hostname(value) | is_valid_ip_address(value)):
             raise InvalidValueError('Invalid hostname: {}'.format(value))
         return value
 
diff --git a/webssh/utils.py b/webssh/utils.py
@@ -30,20 +30,11 @@ def to_int(string):
         pass
 
 
-def is_valid_ipv4_address(ipstr):
+def is_valid_ip_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:
+        ipaddress.ip_address(ipstr)
+    except ValueError:
         return False
     return True