jocay/webssh
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
c013a389fe2e2bd2cdad6ed63a4736377069d67b
webssh/tests/test_handler.py
T
Jocay c013a389fe feat: Implement directory validation and shell command quoting 
- Added is_valid_directory function to validate directory paths, ensuring they do not contain control characters and are within a specified length.
- Introduced quote_shell_arg to safely quote shell arguments, preventing command injection.
- Created build_cd_command to generate a command for changing directories in a shell.
- Enhanced the LoginHandler to utilize a login rate limiter, preventing brute-force attacks by tracking failed login attempts.
- Implemented an EncodingCache to optimize encoding detection for SSH connections.
- Updated the UI to include an input field for specifying an initial directory upon login, with appropriate validation and hints.
- Added a quickbar in the terminal interface for easy access to copy and paste functionality.
- Introduced a toast notification system to provide feedback on copy actions.
- Refactored connection storage to encrypt passwords at rest, improving security.
- Updated various templates and styles to accommodate new features and improve user experience.
2026-08-10 00:07:52 +08:00

606 lines
20 KiB
Python
Raw Blame History

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)
Reference in New Issue View Git Blame Copy Permalink