webssh

commit e94c84637917ae5cec9d27cb3e604450413d392f
parent 499f3b6dcd475e2e3c70d617735f46d993407c9b
Author: Sheng <webmaster0115@gmail.com>
Date:   Tue, 16 Oct 2018 13:56:44 +0800

Added default handler NotFoundHandler

Diffstat:
M
tests/test_app.py
 | 
9
+++++++--
M
webssh/handler.py
 | 
28
++++++++++++++++++++++------
M
webssh/main.py
 | 
3
++-

3 files changed, 31 insertions(+), 9 deletions(-)
diff --git a/tests/test_app.py b/tests/test_app.py
@@ -497,8 +497,13 @@ class TestAppInDebug(OtherTestBase):
             self.assertEqual(response.code, 500)
             self.assertIn(b'Uncaught exception', response.body)
 
-    def test_server_error(self):
-        response = self.fetch('/?error=generate', method='GET')
+    def test_server_error_for_post_method(self):
+        response = self.fetch(
+            '/',
+            method='POST',
+            body=urlencode(dict(self.body, error='raise')),
+            headers=self.headers
+        )
         self.assert_response(b'"status": "Internal Server Error"', response)
 
     def test_html(self):
diff --git a/webssh/handler.py b/webssh/handler.py
@@ -54,11 +54,15 @@ class MixinHandler(object):
         lst = context.trusted_downstream
 
         if lst and ip not in lst:
+            logging.info(
+                'IP {!r} not found in trusted downstream {!r}'.format(ip, lst)
+            )
             return True
 
         if context._orig_protocol == 'http':
             ipaddr = to_ip_address(ip)
             if not ipaddr.is_private:
+                logging.info('Public non-https request is forbidden.')
                 return True
 
     def set_default_headers(self):
@@ -93,6 +97,16 @@ class MixinHandler(object):
         return (ip, port)
 
 
+class NotFoundHandler(MixinHandler, tornado.web.ErrorHandler):
+
+    def initialize(self):
+        pass
+
+    def prepare(self):
+        super(NotFoundHandler, self).prepare()
+        raise tornado.web.HTTPError(404)
+
+
 class IndexHandler(MixinHandler, tornado.web.RequestHandler):
 
     def initialize(self, loop, policy, host_keys_settings):
@@ -101,11 +115,12 @@ class IndexHandler(MixinHandler, tornado.web.RequestHandler):
         self.host_keys_settings = host_keys_settings
         self.ssh_client = self.get_ssh_client()
         self.privatekey_filename = None
+        self.debug = self.settings.get('debug', False)
         self.result = dict(id=None, status=None, encoding=None)
 
     def write_error(self, status_code, **kwargs):
-        if not swallow_http_errors:
-            super(MixinHandler, self).write_error(status_code, **kwargs)
+        if self.request.method != 'POST' or not swallow_http_errors:
+            super(IndexHandler, self).write_error(status_code, **kwargs)
         else:
             exc_info = kwargs.get('exc_info')
             if exc_info:
@@ -269,13 +284,14 @@ class IndexHandler(MixinHandler, tornado.web.RequestHandler):
             future.set_result(worker)
 
     def get(self):
-        debug = self.settings.get('debug', False)
-        if debug and self.get_argument('error', u''):
-            raise ValueError('Uncaught exception')
-        self.render('index.html', debug=debug)
+        self.render('index.html', debug=self.debug)
 
     @tornado.gen.coroutine
     def post(self):
+        if self.debug and self.get_argument('error', u''):
+            # for testing purpose only
+            raise ValueError('Uncaught exception')
+
         future = Future()
         t = threading.Thread(target=self.ssh_connect_wrapped, args=(future,))
         t.setDaemon(True)
diff --git a/webssh/main.py b/webssh/main.py
@@ -3,7 +3,7 @@ import tornado.web
 import tornado.ioloop
 
 from tornado.options import options
-from webssh.handler import IndexHandler, WsockHandler
+from webssh.handler import IndexHandler, WsockHandler, NotFoundHandler
 from webssh.settings import (
     get_app_settings,  get_host_keys_settings, get_policy_setting,
     get_ssl_context, get_server_settings
@@ -23,6 +23,7 @@ def make_handlers(loop, options):
 
 
 def make_app(handlers, settings):
+    settings.update(default_handler_class=NotFoundHandler)
     return tornado.web.Application(handlers, **settings)