webssh

commit c2c81aaeae66039d1264147da87ab1da613a63c2
parent 7e5a1703836c883b5501142a7e1e8027b9a15662
Author: Sheng <webmaster0115@gmail.com>
Date:   Wed, 17 Oct 2018 19:22:51 +0800

Use method initialize to deny forbidden acesss

Diffstat:
M
tests/test_app.py
 | 
8
++++----
M
webssh/handler.py
 | 
13
+++++++++----

2 files changed, 13 insertions(+), 8 deletions(-)
diff --git a/tests/test_app.py b/tests/test_app.py
@@ -569,19 +569,19 @@ class TestAppWithTrustedStream(OtherTestBase):
     def test_with_forbidden_get_request(self):
         response = self.fetch('/', method='GET')
         self.assertEqual(response.code, 403)
-        self.assertIn(b'403: Forbidden', response.body)
+        self.assertIn('Forbidden', response.error.message)
 
     def test_with_forbidden_post_request(self):
         response = self.fetch('/', method='POST', body=urlencode(self.body),
                               headers=self.headers)
-        self.assertEqual(response.code, 200)
-        self.assertIn(b'"status": "Forbidden"', response.body)
+        self.assertEqual(response.code, 403)
+        self.assertIn('Forbidden', response.error.message)
 
     def test_with_forbidden_put_request(self):
         response = self.fetch('/', method='PUT', body=urlencode(self.body),
                               headers=self.headers)
         self.assertEqual(response.code, 403)
-        self.assertIn(b'403: Forbidden', response.body)
+        self.assertIn('Forbidden', response.error.message)
 
 
 class TestAppNotFoundHandler(OtherTestBase):
diff --git a/webssh/handler.py b/webssh/handler.py
@@ -43,9 +43,13 @@ class MixinHandler(object):
         'Server': 'TornadoServer'
     }
 
-    def prepare(self):
+    def initialize(self):
         if self.is_forbidden():
-            raise tornado.web.HTTPError(403)
+            self.request.connection.stream.write(
+                b'%s 403 Forbidden\r\n\r\n' % to_bytes(self.request.version)
+            )
+            self.request.connection.close()
+            raise ValueError('Accesss denied')
 
     def is_forbidden(self):
         """
@@ -105,10 +109,9 @@ class MixinHandler(object):
 class NotFoundHandler(MixinHandler, tornado.web.ErrorHandler):
 
     def initialize(self):
-        pass
+        super(NotFoundHandler, self).initialize()
 
     def prepare(self):
-        super(NotFoundHandler, self).prepare()
         raise tornado.web.HTTPError(404)
 
 
@@ -122,6 +125,7 @@ class IndexHandler(MixinHandler, tornado.web.RequestHandler):
         self.privatekey_filename = None
         self.debug = self.settings.get('debug', False)
         self.result = dict(id=None, status=None, encoding=None)
+        super(IndexHandler, self).initialize()
 
     def write_error(self, status_code, **kwargs):
         if self.request.method != 'POST' or not swallow_http_errors:
@@ -322,6 +326,7 @@ class WsockHandler(MixinHandler, tornado.websocket.WebSocketHandler):
     def initialize(self, loop):
         self.loop = loop
         self.worker_ref = None
+        super(WsockHandler, self).initialize()
 
     def open(self):
         self.src_addr = self.get_client_addr()