mirror of https://github.com/huashengdun/webssh
250 lines
8.4 KiB
Python
250 lines
8.4 KiB
Python
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.handler import (
|
|
MixinHandler, IndexHandler, WsockHandler, InvalidValueError
|
|
)
|
|
|
|
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', 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', 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 TestIndexHandler(unittest.TestCase):
|
|
|
|
def test_get_specific_pkey_with_plain_key(self):
|
|
fname = 'test_rsa.key'
|
|
cls = paramiko.RSAKey
|
|
key = read_file(make_tests_data_path(fname))
|
|
|
|
pkey = IndexHandler.get_specific_pkey(cls, key, None)
|
|
self.assertIsInstance(pkey, cls)
|
|
|
|
pkey = IndexHandler.get_specific_pkey(cls, key, 'iginored')
|
|
self.assertIsInstance(pkey, cls)
|
|
|
|
pkey = IndexHandler.get_specific_pkey(cls, 'x'+key, None)
|
|
self.assertIsNone(pkey)
|
|
|
|
def test_get_specific_pkey_with_encrypted_key(self):
|
|
fname = 'test_rsa_password.key'
|
|
cls = paramiko.RSAKey
|
|
password = 'television'
|
|
|
|
key = read_file(make_tests_data_path(fname))
|
|
pkey = IndexHandler.get_specific_pkey(cls, key, password)
|
|
self.assertIsInstance(pkey, cls)
|
|
|
|
pkey = IndexHandler.get_specific_pkey(cls, 'x'+key, None)
|
|
self.assertIsNone(pkey)
|
|
|
|
with self.assertRaises(paramiko.PasswordRequiredException):
|
|
pkey = IndexHandler.get_specific_pkey(cls, key, None)
|
|
|
|
def test_get_pkey_obj_with_plain_key(self):
|
|
fname = 'test_ed25519.key'
|
|
cls = paramiko.Ed25519Key
|
|
key = read_file(make_tests_data_path(fname))
|
|
|
|
pkey = IndexHandler.get_pkey_obj(key, None, fname)
|
|
self.assertIsInstance(pkey, cls)
|
|
|
|
pkey = IndexHandler.get_pkey_obj(key, 'iginored', fname)
|
|
self.assertIsInstance(pkey, cls)
|
|
|
|
with self.assertRaises(InvalidValueError) as ctx:
|
|
pkey = IndexHandler.get_pkey_obj('x'+key, None, fname)
|
|
self.assertIn('Invalid private key', str(ctx.exception))
|
|
|
|
def test_get_pkey_obj_with_encrypted_key(self):
|
|
fname = 'test_ed25519_password.key'
|
|
password = 'abc123'
|
|
cls = paramiko.Ed25519Key
|
|
key = read_file(make_tests_data_path(fname))
|
|
|
|
pkey = IndexHandler.get_pkey_obj(key, password, fname)
|
|
self.assertIsInstance(pkey, cls)
|
|
|
|
with self.assertRaises(InvalidValueError) as ctx:
|
|
pkey = IndexHandler.get_pkey_obj(key, 'wrongpass', fname)
|
|
self.assertIn('Wrong password', str(ctx.exception))
|
|
|
|
with self.assertRaises(InvalidValueError) as ctx:
|
|
pkey = IndexHandler.get_pkey_obj('x'+key, '', fname)
|
|
self.assertIn('Invalid private key', str(ctx.exception))
|
|
|
|
with self.assertRaises(paramiko.PasswordRequiredException):
|
|
pkey = IndexHandler.get_pkey_obj(key, '', fname)
|
|
|
|
|
|
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))
|