mirror of https://github.com/sgoudham/Enso-Bot.git
You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
118 lines
4.1 KiB
Python
118 lines
4.1 KiB
Python
import contextlib
|
|
import unittest
|
|
|
|
from .exceptions import InvalidHandshake
|
|
from .handshake import *
|
|
from .handshake import accept # private API
|
|
|
|
|
|
class HandshakeTests(unittest.TestCase):
|
|
|
|
def test_accept(self):
|
|
# Test vector from RFC 6455
|
|
key = "dGhlIHNhbXBsZSBub25jZQ=="
|
|
acc = "s3pPLMBiTxaQ9kYGzzhZRbK+xOo="
|
|
self.assertEqual(accept(key), acc)
|
|
|
|
def test_round_trip(self):
|
|
request_headers = {}
|
|
request_key = build_request(request_headers.__setitem__)
|
|
response_key = check_request(request_headers.__getitem__)
|
|
self.assertEqual(request_key, response_key)
|
|
response_headers = {}
|
|
build_response(response_headers.__setitem__, response_key)
|
|
check_response(response_headers.__getitem__, request_key)
|
|
|
|
@contextlib.contextmanager
|
|
def assert_invalid_request_headers(self):
|
|
"""
|
|
Provide request headers for corruption.
|
|
|
|
Assert that the transformation made them invalid.
|
|
|
|
"""
|
|
headers = {}
|
|
build_request(headers.__setitem__)
|
|
yield headers
|
|
with self.assertRaises(InvalidHandshake):
|
|
check_request(headers.__getitem__)
|
|
|
|
def test_request_invalid_upgrade(self):
|
|
with self.assert_invalid_request_headers() as headers:
|
|
headers['Upgrade'] = 'socketweb'
|
|
|
|
def test_request_missing_upgrade(self):
|
|
with self.assert_invalid_request_headers() as headers:
|
|
del headers['Upgrade']
|
|
|
|
def test_request_invalid_connection(self):
|
|
with self.assert_invalid_request_headers() as headers:
|
|
headers['Connection'] = 'Downgrade'
|
|
|
|
def test_request_missing_connection(self):
|
|
with self.assert_invalid_request_headers() as headers:
|
|
del headers['Connection']
|
|
|
|
def test_request_invalid_key_not_base64(self):
|
|
with self.assert_invalid_request_headers() as headers:
|
|
headers['Sec-WebSocket-Key'] = "!@#$%^&*()"
|
|
|
|
def test_request_invalid_key_not_well_padded(self):
|
|
with self.assert_invalid_request_headers() as headers:
|
|
headers['Sec-WebSocket-Key'] = "CSIRmL8dWYxeAdr/XpEHRw"
|
|
|
|
def test_request_invalid_key_not_16_bytes_long(self):
|
|
with self.assert_invalid_request_headers() as headers:
|
|
headers['Sec-WebSocket-Key'] = "ZLpprpvK4PE="
|
|
|
|
def test_request_missing_key(self):
|
|
with self.assert_invalid_request_headers() as headers:
|
|
del headers['Sec-WebSocket-Key']
|
|
|
|
def test_request_invalid_version(self):
|
|
with self.assert_invalid_request_headers() as headers:
|
|
headers['Sec-WebSocket-Version'] = '42'
|
|
|
|
def test_request_missing_version(self):
|
|
with self.assert_invalid_request_headers() as headers:
|
|
del headers['Sec-WebSocket-Version']
|
|
|
|
@contextlib.contextmanager
|
|
def assert_invalid_response_headers(self, key='CSIRmL8dWYxeAdr/XpEHRw=='):
|
|
"""
|
|
Provide response headers for corruption.
|
|
|
|
Assert that the transformation made them invalid.
|
|
|
|
"""
|
|
headers = {}
|
|
build_response(headers.__setitem__, key)
|
|
yield headers
|
|
with self.assertRaises(InvalidHandshake):
|
|
check_response(headers.__getitem__, key)
|
|
|
|
def test_response_invalid_upgrade(self):
|
|
with self.assert_invalid_response_headers() as headers:
|
|
headers['Upgrade'] = 'socketweb'
|
|
|
|
def test_response_missing_upgrade(self):
|
|
with self.assert_invalid_response_headers() as headers:
|
|
del headers['Upgrade']
|
|
|
|
def test_response_invalid_connection(self):
|
|
with self.assert_invalid_response_headers() as headers:
|
|
headers['Connection'] = 'Downgrade'
|
|
|
|
def test_response_missing_connection(self):
|
|
with self.assert_invalid_response_headers() as headers:
|
|
del headers['Connection']
|
|
|
|
def test_response_invalid_accept(self):
|
|
with self.assert_invalid_response_headers() as headers:
|
|
other_key = "1Eq4UDEFQYg3YspNgqxv5g=="
|
|
headers['Sec-WebSocket-Accept'] = accept(other_key)
|
|
|
|
def test_response_missing_accept(self):
|
|
with self.assert_invalid_response_headers() as headers:
|
|
del headers['Sec-WebSocket-Accept']
|