import io import unittest import paramiko from tornado.httputil import HTTPServerRequest from tornado.options import options from tests.utils import read_file, make_tests_data_path from webssh import handler from webssh import worker from webssh.handler import ( IndexHandler, MixinHandler, WsockHandler, PrivateKey, InvalidValueError, SSHClient, LoginRateLimiter, EncodingCache ) try: from unittest.mock import Mock except ImportError: from mock import Mock class TestMixinHandler(unittest.TestCase): def test_is_forbidden(self): mhandler = MixinHandler() handler.redirecting = True options.fbidhttp = True context = Mock( address=('8.8.8.8', 8888), trusted_downstream=['127.0.0.1'], _orig_protocol='http' ) hostname = '4.4.4.4' self.assertTrue(mhandler.is_forbidden(context, hostname)) context = Mock( address=('8.8.8.8', 8888), trusted_downstream=[], _orig_protocol='http' ) hostname = 'www.google.com' self.assertEqual(mhandler.is_forbidden(context, hostname), False) context = Mock( address=('8.8.8.8', 8888), trusted_downstream=[], _orig_protocol='http' ) hostname = '4.4.4.4' self.assertTrue(mhandler.is_forbidden(context, hostname)) context = Mock( address=('192.168.1.1', 8888), trusted_downstream=[], _orig_protocol='http' ) hostname = 'www.google.com' self.assertIsNone(mhandler.is_forbidden(context, hostname)) options.fbidhttp = False self.assertIsNone(mhandler.is_forbidden(context, hostname)) hostname = '4.4.4.4' self.assertIsNone(mhandler.is_forbidden(context, hostname)) handler.redirecting = False self.assertIsNone(mhandler.is_forbidden(context, hostname)) context._orig_protocol = 'https' self.assertIsNone(mhandler.is_forbidden(context, hostname)) def test_get_redirect_url(self): mhandler = MixinHandler() hostname = 'www.example.com' uri = '/' port = 443 self.assertEqual( mhandler.get_redirect_url(hostname, port, uri=uri), 'https://www.example.com/' ) port = 4433 self.assertEqual( mhandler.get_redirect_url(hostname, port, uri), 'https://www.example.com:4433/' ) def test_get_client_addr(self): mhandler = MixinHandler() client_addr = ('8.8.8.8', 8888) context_addr = ('127.0.0.1', 1234) options.xheaders = True mhandler.context = Mock(address=context_addr) mhandler.get_real_client_addr = lambda: None self.assertEqual(mhandler.get_client_addr(), context_addr) mhandler.context = Mock(address=context_addr) mhandler.get_real_client_addr = lambda: client_addr self.assertEqual(mhandler.get_client_addr(), client_addr) options.xheaders = False mhandler.context = Mock(address=context_addr) mhandler.get_real_client_addr = lambda: client_addr self.assertEqual(mhandler.get_client_addr(), context_addr) def test_get_real_client_addr(self): x_forwarded_for = '1.1.1.1' x_forwarded_port = 1111 x_real_ip = '2.2.2.2' x_real_port = 2222 fake_port = 65535 mhandler = MixinHandler() mhandler.request = HTTPServerRequest(uri='/') mhandler.request.remote_ip = x_forwarded_for self.assertIsNone(mhandler.get_real_client_addr()) mhandler.request.headers.add('X-Forwarded-For', x_forwarded_for) self.assertEqual(mhandler.get_real_client_addr(), (x_forwarded_for, fake_port)) mhandler.request.headers.add('X-Forwarded-Port', str(fake_port + 1)) self.assertEqual(mhandler.get_real_client_addr(), (x_forwarded_for, fake_port)) mhandler.request.headers['X-Forwarded-Port'] = x_forwarded_port self.assertEqual(mhandler.get_real_client_addr(), (x_forwarded_for, x_forwarded_port)) mhandler.request.remote_ip = x_real_ip mhandler.request.headers.add('X-Real-Ip', x_real_ip) self.assertEqual(mhandler.get_real_client_addr(), (x_real_ip, fake_port)) mhandler.request.headers.add('X-Real-Port', str(fake_port + 1)) self.assertEqual(mhandler.get_real_client_addr(), (x_real_ip, fake_port)) mhandler.request.headers['X-Real-Port'] = x_real_port self.assertEqual(mhandler.get_real_client_addr(), (x_real_ip, x_real_port)) class TestPrivateKey(unittest.TestCase): def get_pk_obj(self, fname, password=None): key = read_file(make_tests_data_path(fname)) return PrivateKey(key, password=password, filename=fname) def _test_with_encrypted_key(self, fname, password, klass): pk = self.get_pk_obj(fname, password='') with self.assertRaises(InvalidValueError) as ctx: pk.get_pkey_obj() self.assertIn('Need a passphrase', str(ctx.exception)) pk = self.get_pk_obj(fname, password='wrongpass') with self.assertRaises(InvalidValueError) as ctx: pk.get_pkey_obj() self.assertIn('wrong passphrase', str(ctx.exception)) pk = self.get_pk_obj(fname, password=password) self.assertIsInstance(pk.get_pkey_obj(), klass) def test_class_with_invalid_key_length(self): key = u'a' * (PrivateKey.max_length + 1) with self.assertRaises(InvalidValueError) as ctx: PrivateKey(key) self.assertIn('Invalid key length', str(ctx.exception)) def test_get_pkey_obj_with_invalid_key(self): key = u'a b c' fname = 'abc' pk = PrivateKey(key, filename=fname) with self.assertRaises(InvalidValueError) as ctx: pk.get_pkey_obj() self.assertIn('Invalid key {}'.format(fname), str(ctx.exception)) def test_get_pkey_obj_with_plain_rsa_key(self): pk = self.get_pk_obj('test_rsa.key') self.assertIsInstance(pk.get_pkey_obj(), paramiko.RSAKey) def test_get_pkey_obj_with_plain_ed25519_key(self): pk = self.get_pk_obj('test_ed25519.key') self.assertIsInstance(pk.get_pkey_obj(), paramiko.Ed25519Key) def test_get_pkey_obj_with_encrypted_rsa_key(self): fname = 'test_rsa_password.key' password = 'television' self._test_with_encrypted_key(fname, password, paramiko.RSAKey) def test_get_pkey_obj_with_encrypted_ed25519_key(self): fname = 'test_ed25519_password.key' password = 'abc123' self._test_with_encrypted_key(fname, password, paramiko.Ed25519Key) def test_get_pkey_obj_with_encrypted_new_rsa_key(self): fname = 'test_new_rsa_password.key' password = '123456' self._test_with_encrypted_key(fname, password, paramiko.RSAKey) def test_get_pkey_obj_with_plain_new_dsa_key(self): pk = self.get_pk_obj('test_new_dsa.key') self.assertIsInstance(pk.get_pkey_obj(), paramiko.DSSKey) def test_parse_name(self): key = u'-----BEGIN PRIVATE KEY-----' pk = PrivateKey(key) name, _ = pk.parse_name(pk.iostr, pk.tag_to_name) self.assertIsNone(name) key = u'-----BEGIN xxx PRIVATE KEY-----' pk = PrivateKey(key) name, _ = pk.parse_name(pk.iostr, pk.tag_to_name) self.assertIsNone(name) key = u'-----BEGIN RSA PRIVATE KEY-----' pk = PrivateKey(key) name, _ = pk.parse_name(pk.iostr, pk.tag_to_name) self.assertIsNone(name) key = u'-----BEGIN RSA PRIVATE KEY-----' pk = PrivateKey(key) name, _ = pk.parse_name(pk.iostr, pk.tag_to_name) self.assertIsNone(name) key = u'-----BEGIN RSA PRIVATE KEY-----' pk = PrivateKey(key) name, _ = pk.parse_name(pk.iostr, pk.tag_to_name) self.assertIsNone(name) for tag, to_name in PrivateKey.tag_to_name.items(): key = u'-----BEGIN {} PRIVATE KEY----- \r\n'.format(tag) pk = PrivateKey(key) name, length = pk.parse_name(pk.iostr, pk.tag_to_name) self.assertEqual(name, to_name) self.assertEqual(length, len(key)) class TestWsockHandler(unittest.TestCase): def test_check_origin(self): request = HTTPServerRequest(uri='/') obj = Mock(spec=WsockHandler, request=request) obj.origin_policy = 'same' request.headers['Host'] = 'www.example.com:4433' origin = 'https://www.example.com:4433' self.assertTrue(WsockHandler.check_origin(obj, origin)) origin = 'https://www.example.com' self.assertFalse(WsockHandler.check_origin(obj, origin)) obj.origin_policy = 'primary' self.assertTrue(WsockHandler.check_origin(obj, origin)) origin = 'https://blog.example.com' self.assertTrue(WsockHandler.check_origin(obj, origin)) origin = 'https://blog.example.org' self.assertFalse(WsockHandler.check_origin(obj, origin)) origin = 'https://blog.example.org' obj.origin_policy = {'https://blog.example.org'} self.assertTrue(WsockHandler.check_origin(obj, origin)) origin = 'http://blog.example.org' obj.origin_policy = {'http://blog.example.org'} self.assertTrue(WsockHandler.check_origin(obj, origin)) origin = 'http://blog.example.org' obj.origin_policy = {'https://blog.example.org'} self.assertFalse(WsockHandler.check_origin(obj, origin)) obj.origin_policy = '*' origin = 'https://blog.example.org' self.assertTrue(WsockHandler.check_origin(obj, origin)) def test_failed_weak_ref(self): request = HTTPServerRequest(uri='/') obj = Mock(spec=WsockHandler, request=request) obj.src_addr = ("127.0.0.1", 8888) class FakeWeakRef: def __init__(self): self.count = 0 def __call__(self): self.count += 1 return None ref = FakeWeakRef() obj.worker_ref = ref WsockHandler.on_message(obj, b'{"data": "somestuff"}') self.assertGreaterEqual(ref.count, 1) obj.close.assert_called_with(reason='No worker found') def test_worker_closed(self): request = HTTPServerRequest(uri='/') obj = Mock(spec=WsockHandler, request=request) obj.src_addr = ("127.0.0.1", 8888) class Worker: def __init__(self): self.closed = True class FakeWeakRef: def __call__(self): return Worker() ref = FakeWeakRef() obj.worker_ref = ref WsockHandler.on_message(obj, b'{"data": "somestuff"}') obj.close.assert_called_with(reason='Worker closed') class TestIndexHandler(unittest.TestCase): def test_get_args_keeps_password_for_saved_connection(self): obj = Mock(spec=IndexHandler) obj.get_hostname.return_value = '127.0.0.1' obj.get_port.return_value = 22 obj.get_value.return_value = 'root' obj.get_privatekey.return_value = ('', '') obj.get_directory.return_value = '' obj.policy = paramiko.WarningPolicy() obj.ssh_client = Mock() values = { 'password': 'root-secret', 'passphrase': '', 'totp': '', 'term': 'xterm-256color' } obj.get_argument.side_effect = lambda name, default=u'': values.get( name, default ) args = IndexHandler.get_args(obj) self.assertEqual( ('127.0.0.1', 22, 'root', 'root-secret'), args[:4] ) self.assertEqual('root-secret', obj.connection_info['password']) def test_get_args_carries_directory(self): obj = Mock(spec=IndexHandler) obj.get_hostname.return_value = '127.0.0.1' obj.get_port.return_value = 22 obj.get_value.return_value = 'root' obj.get_privatekey.return_value = ('', '') obj.get_directory.return_value = '/var/www' obj.policy = paramiko.WarningPolicy() obj.ssh_client = Mock() obj.get_argument.side_effect = lambda name, default=u'': default args = IndexHandler.get_args(obj) self.assertEqual('/var/www', args[6]) self.assertEqual('/var/www', obj.connection_info['directory']) def test_get_directory(self): obj = Mock(spec=IndexHandler) values = {} obj.get_argument.side_effect = lambda name, default=u'': values.get( name, default ) self.assertEqual('', IndexHandler.get_directory(obj)) values['directory'] = ' ' self.assertEqual('', IndexHandler.get_directory(obj)) values['directory'] = ' /var/www ' self.assertEqual('/var/www', IndexHandler.get_directory(obj)) values['directory'] = '/tmp\nrm -rf /' with self.assertRaises(InvalidValueError): IndexHandler.get_directory(obj) values['directory'] = '/tmp' + 'a' * 1024 with self.assertRaises(InvalidValueError): IndexHandler.get_directory(obj) def test_change_directory_quotes_the_path(self): obj = Mock(spec=IndexHandler) chan = Mock() IndexHandler.change_directory(obj, chan, '/tmp; reboot') chan.sendall.assert_called_with(b"cd '/tmp; reboot'\r") def test_change_directory_swallows_channel_errors(self): obj = Mock(spec=IndexHandler) chan = Mock() chan.sendall.side_effect = paramiko.SSHException('boom') # a failed cd must not prevent the session from starting IndexHandler.change_directory(obj, chan, '/var/www') def test_null_in_encoding(self): mock_handler = Mock(spec=IndexHandler) # This is a little nasty, but the index handler has a lot of # dependencies to mock. Mocking out everything but the bits # we want to test lets us test this case without needing to # refactor the relevant code out of IndexHandler def parse_encoding(data): return IndexHandler.parse_encoding(mock_handler, data) mock_handler.parse_encoding = parse_encoding def detect_encoding(ssh): return IndexHandler.detect_encoding(mock_handler, ssh) mock_handler.detect_encoding = detect_encoding ssh = Mock(spec=SSHClient) stdin = io.BytesIO() stdout = io.BytesIO(initial_bytes=b"UTF-8\0") stderr = io.BytesIO() ssh.exec_command.return_value = (stdin, stdout, stderr) self.assertIsNone(IndexHandler.detect_encoding(mock_handler, ssh)) encoding = IndexHandler.get_default_encoding(mock_handler, ssh) self.assertEqual("utf-8", encoding) class TestLoginRateLimiter(unittest.TestCase): def test_blocks_after_too_many_failures(self): limiter = LoginRateLimiter(max_attempts=3, window=300) self.assertFalse(limiter.is_blocked('1.1.1.1')) for _ in range(2): limiter.record_failure('1.1.1.1') self.assertFalse(limiter.is_blocked('1.1.1.1')) limiter.record_failure('1.1.1.1') self.assertTrue(limiter.is_blocked('1.1.1.1')) # other clients are unaffected self.assertFalse(limiter.is_blocked('2.2.2.2')) def test_success_resets_the_counter(self): limiter = LoginRateLimiter(max_attempts=2, window=300) limiter.record_failure('1.1.1.1') limiter.record_failure('1.1.1.1') self.assertTrue(limiter.is_blocked('1.1.1.1')) limiter.reset('1.1.1.1') self.assertFalse(limiter.is_blocked('1.1.1.1')) def test_entries_expire_after_the_window(self): limiter = LoginRateLimiter(max_attempts=1, window=300) limiter.record_failure('1.1.1.1') self.assertTrue(limiter.is_blocked('1.1.1.1')) # pretend the failure happened longer ago than the window count, last = limiter.attempts['1.1.1.1'] limiter.attempts['1.1.1.1'] = (count, last - 301) self.assertFalse(limiter.is_blocked('1.1.1.1')) self.assertEqual({}, limiter.attempts) def test_clear(self): limiter = LoginRateLimiter(max_attempts=1, window=300) limiter.record_failure('1.1.1.1') self.assertTrue(limiter.is_blocked('1.1.1.1')) limiter.clear() self.assertFalse(limiter.is_blocked('1.1.1.1')) class TestEncodingCache(unittest.TestCase): def test_get_and_set(self): cache = EncodingCache() self.assertIsNone(cache.get(('127.0.0.1', 22))) cache.set(('127.0.0.1', 22), 'GBK') self.assertEqual('GBK', cache.get(('127.0.0.1', 22))) # the port is part of the key self.assertIsNone(cache.get(('127.0.0.1', 2200))) def test_a_failed_probe_is_not_cached(self): cache = EncodingCache() cache.set(('127.0.0.1', 22), None) cache.set(('127.0.0.1', 22), '') self.assertIsNone(cache.get(('127.0.0.1', 22))) def test_cache_is_bounded(self): cache = EncodingCache(maxsize=2) cache.set(('a', 22), 'UTF-8') cache.set(('b', 22), 'UTF-8') cache.set(('c', 22), 'UTF-8') self.assertEqual(1, len(cache.cache)) self.assertEqual('UTF-8', cache.get(('c', 22))) class TestGetEncoding(unittest.TestCase): def setUp(self): handler.encoding_cache.clear() self.addCleanup(handler.encoding_cache.clear) def make_handler(self, detected): obj = Mock(spec=IndexHandler) obj.detect_encoding.return_value = detected return obj def test_probes_once_then_serves_from_cache(self): obj = self.make_handler('GBK') ssh = Mock() dst_addr = ('127.0.0.1', 22) self.assertEqual('GBK', IndexHandler.get_encoding(obj, ssh, dst_addr)) self.assertEqual('GBK', IndexHandler.get_encoding(obj, ssh, dst_addr)) self.assertEqual(1, obj.detect_encoding.call_count) def test_failed_detection_falls_back_without_caching(self): obj = self.make_handler(None) ssh = Mock() dst_addr = ('127.0.0.1', 22) self.assertEqual( 'utf-8', IndexHandler.get_encoding(obj, ssh, dst_addr) ) self.assertEqual( 'utf-8', IndexHandler.get_encoding(obj, ssh, dst_addr) ) self.assertEqual(2, obj.detect_encoding.call_count) class TestClientKey(unittest.TestCase): def make_handler(self, auth_manager, addr=('10.0.0.1', 1234)): obj = Mock(spec=MixinHandler) obj.auth_manager = auth_manager obj.get_client_addr.return_value = addr return obj def test_logged_in_users_get_their_own_bucket(self): alice = self.make_handler(Mock()) alice.current_user = 'alice' bob = self.make_handler(Mock()) bob.current_user = 'bob' # same source address, different buckets self.assertEqual('user:alice', MixinHandler.get_client_key(alice)) self.assertEqual('user:bob', MixinHandler.get_client_key(bob)) def test_without_auth_the_address_is_used(self): obj = self.make_handler(None) self.assertEqual('10.0.0.1', MixinHandler.get_client_key(obj)) def test_a_username_cannot_collide_with_an_address(self): spoofer = self.make_handler(Mock()) spoofer.current_user = '10.0.0.1' direct = self.make_handler(None) self.assertNotEqual( MixinHandler.get_client_key(spoofer), MixinHandler.get_client_key(direct) ) class TestClearWorker(unittest.TestCase): def make_worker(self, worker_id, client_key): obj = Mock() obj.id = worker_id obj.client_key = client_key return obj def test_removes_the_worker_and_prunes_empty_buckets(self): a = self.make_worker('a', 'user:admin') b = self.make_worker('b', 'user:admin') clients = {'user:admin': {'a': a, 'b': b}} worker.clear_worker(a, clients) self.assertEqual({'user:admin': {'b': b}}, clients) worker.clear_worker(b, clients) self.assertEqual({}, clients) def test_unregistered_worker_is_ignored(self): orphan = self.make_worker('a', None) clients = {} # must not raise even though the worker was never registered worker.clear_worker(orphan, clients) self.assertEqual({}, clients) clients = {'user:admin': {'b': 'other'}} orphan.client_key = 'user:admin' worker.clear_worker(orphan, clients) self.assertEqual({'user:admin': {'b': 'other'}}, clients)