jinql
/
webssh
mirror of https://github.com/huashengdun/webssh
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

Added attribute custom_headers to MixinHandler

pull/38/head
Sheng 2018-10-16 15:14:34 +08:00
parent af60cd1cd5
commit a68eff592f
2 changed files with 17 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
tests/test_app.py
View File

@ -586,22 +586,30 @@ class TestAppWithTrustedStream(OtherTestBase):
class TestAppNotFoundHandler(OtherTestBase): class TestAppNotFoundHandler(OtherTestBase):
custom_headers = handler.MixinHandler.custom_headers
def test_with_not_found_get_request(self): def test_with_not_found_get_request(self):
response = self.fetch('/pathnotfound', method='GET') response = self.fetch('/pathnotfound', method='GET')
self.assertEqual(response.code, 404) self.assertEqual(response.code, 404)
self.assertEqual(response.headers['Server'], 'TornadoServer') self.assertEqual(
response.headers['Server'], self.custom_headers['Server']
)
self.assertIn(b'404: Not Found', response.body) self.assertIn(b'404: Not Found', response.body)
def test_with_not_found_post_request(self): def test_with_not_found_post_request(self):
response = self.fetch('/pathnotfound', method='POST', response = self.fetch('/pathnotfound', method='POST',
body=urlencode(self.body), headers=self.headers) body=urlencode(self.body), headers=self.headers)
self.assertEqual(response.code, 404) self.assertEqual(response.code, 404)
self.assertEqual(response.headers['Server'], 'TornadoServer') self.assertEqual(
response.headers['Server'], self.custom_headers['Server']
)
self.assertIn(b'404: Not Found', response.body) self.assertIn(b'404: Not Found', response.body)
def test_with_not_found_put_request(self): def test_with_not_found_put_request(self):
response = self.fetch('/pathnotfound', method='PUT', response = self.fetch('/pathnotfound', method='PUT',
body=urlencode(self.body), headers=self.headers) body=urlencode(self.body), headers=self.headers)
self.assertEqual(response.code, 404) self.assertEqual(response.code, 404)
self.assertEqual(response.headers['Server'], 'TornadoServer') self.assertEqual(
response.headers['Server'], self.custom_headers['Server']
)
self.assertIn(b'404: Not Found', response.body) self.assertIn(b'404: Not Found', response.body)

7
webssh/handler.py
View File

@ -39,6 +39,10 @@ class InvalidValueError(Exception):
class MixinHandler(object): class MixinHandler(object):
custom_headers = {
'Server': 'TornadoServer'
}
def prepare(self): def prepare(self):
if self.is_forbidden(): if self.is_forbidden():
raise tornado.web.HTTPError(403) raise tornado.web.HTTPError(403)
@ -66,7 +70,8 @@ class MixinHandler(object):
return True return True
def set_default_headers(self): def set_default_headers(self):
self.set_header('Server', 'TornadoServer') for header in self.custom_headers.items():
self.set_header(*header)
def get_value(self, name): def get_value(self, name):
value = self.get_argument(name) value = self.get_argument(name)