webssh

commit db3ee2b784211abd3eafa4b111a85bca6036968d
parent a51918d2ac1c28bd48f01c1069335caf2132fa90
Author: Sheng <webmaster0115@gmail.com>
Date:   Mon, 15 Oct 2018 17:49:02 +0800

Added proxies option for trusted downstream

Diffstat:
M
tests/test_settings.py
 | 
19
++++++++++++++++++-
M
webssh/settings.py
 | 
12
++++++++++++

2 files changed, 30 insertions(+), 1 deletion(-)
diff --git a/tests/test_settings.py b/tests/test_settings.py
@@ -10,7 +10,7 @@ from tests.utils import make_tests_data_path
 from webssh.policy import load_host_keys
 from webssh.settings import (
     get_host_keys_settings, get_policy_setting, base_dir, print_version,
-    get_ssl_context
+    get_ssl_context, get_trusted_downstream
 )
 from webssh.utils import UnicodeType
 from webssh._version import __version__
@@ -120,3 +120,20 @@ class TestSettings(unittest.TestCase):
         options.keyfile = make_tests_data_path('cert.key')
         ssl_ctx = get_ssl_context(options)
         self.assertIsNotNone(ssl_ctx)
+
+    def test_get_trusted_downstream(self):
+        options.proxies = ''
+        proxies = set()
+        self.assertEqual(get_trusted_downstream(options), proxies)
+
+        options.proxies = '1.1.1.1, 2.2.2.2'
+        proxies = set(['1.1.1.1', '2.2.2.2'])
+        self.assertEqual(get_trusted_downstream(options), proxies)
+
+        options.proxies = '1.1.1.1, 2.2.2.2, 2.2.2.2'
+        proxies = set(['1.1.1.1', '2.2.2.2'])
+        self.assertEqual(get_trusted_downstream(options), proxies)
+
+        options.proxies = '1.1.1.1, 2.2.2.'
+        with self.assertRaises(ValueError):
+            get_trusted_downstream(options), proxies
diff --git a/webssh/settings.py b/webssh/settings.py
@@ -7,6 +7,7 @@ from tornado.options import define
 from webssh.policy import (
     load_host_keys, get_policy_class, check_policy_setting
 )
+from webssh.utils import to_ip_address
 from webssh._version import __version__
 
 
@@ -27,6 +28,7 @@ define('policy', default='warning',
        help='Missing host key policy, reject|autoadd|warning')
 define('hostFile', default='', help='User defined host keys file')
 define('sysHostFile', default='', help='System wide host keys file')
+define('proxies', default='', help='trusted downstream, separated by comma')
 define('wpIntvl', type=int, default=0, help='Websocket ping interval')
 define('version', type=bool, help='Show version information',
        callback=print_version)
@@ -92,3 +94,13 @@ def get_ssl_context(options):
         ssl_ctx = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
         ssl_ctx.load_cert_chain(options.certfile, options.keyfile)
         return ssl_ctx
+
+
+def get_trusted_downstream(options):
+    proxies = set()
+    for ip in options.proxies.split(','):
+        ip = ip.strip()
+        if ip:
+            to_ip_address(ip)
+            proxies.add(ip)
+    return proxies