Added --host-token to allow choosing target by hostname.
This commit is contained in:
parent
44bb213a72
commit
af85184e28
|
@ -88,7 +88,7 @@ class ProxyRequestHandlerTestCase(unittest.TestCase):
|
|||
return ("some host", "some port")
|
||||
|
||||
host, port = self.handler.get_target(
|
||||
TestPlugin(None), self.handler.path)
|
||||
TestPlugin(None))
|
||||
|
||||
self.assertEqual(host, "some host")
|
||||
self.assertEqual(port, "some port")
|
||||
|
@ -99,7 +99,7 @@ class ProxyRequestHandlerTestCase(unittest.TestCase):
|
|||
return ("unix_socket", "/tmp/socket")
|
||||
|
||||
_, socket = self.handler.get_target(
|
||||
TestPlugin(None), self.handler.path)
|
||||
TestPlugin(None))
|
||||
|
||||
self.assertEqual(socket, "/tmp/socket")
|
||||
|
||||
|
@ -109,7 +109,7 @@ class ProxyRequestHandlerTestCase(unittest.TestCase):
|
|||
return None
|
||||
|
||||
self.assertRaises(FakeServer.EClose, self.handler.get_target,
|
||||
TestPlugin(None), "https://localhost:6080/websockify?token=blah")
|
||||
TestPlugin(None))
|
||||
|
||||
def test_token_plugin(self):
|
||||
class TestPlugin(token_plugins.BasePlugin):
|
||||
|
|
|
@ -135,23 +135,31 @@ Traffic Legend:
|
|||
self.log_message("%s:%s: Closed target",
|
||||
self.server.target_host, self.server.target_port)
|
||||
|
||||
def get_target(self, target_plugin, path):
|
||||
def get_target(self, target_plugin):
|
||||
"""
|
||||
Parses the path, extracts a token, and looks up a target
|
||||
for that token using the token plugin. Sets
|
||||
target_host and target_port if successful
|
||||
Gets a token from either the path or the host,
|
||||
depending on --host-token, and looks up a target
|
||||
for that token using the token plugin. Used by
|
||||
validate_connection() to set target_host and target_port.
|
||||
"""
|
||||
# The files in targets contain the lines
|
||||
# in the form of token: host:port
|
||||
|
||||
# Extract the token parameter from url
|
||||
args = parse_qs(urlparse(path)[4]) # 4 is the query from url
|
||||
if self.host_token:
|
||||
token = self.headers.get('Host')
|
||||
|
||||
if not 'token' in args or not len(args['token']):
|
||||
else:
|
||||
# Extract the token parameter from url
|
||||
args = parse_qs(urlparse(self.path)[4]) # 4 is the query from url
|
||||
|
||||
if 'token' in args and len(args['token']):
|
||||
token = args['token'][0].rstrip('\n')
|
||||
else:
|
||||
token = None
|
||||
|
||||
if token is None:
|
||||
raise self.server.EClose("Token not present")
|
||||
|
||||
token = args['token'][0].rstrip('\n')
|
||||
|
||||
result_pair = target_plugin.lookup(token)
|
||||
|
||||
if result_pair is not None:
|
||||
|
@ -263,6 +271,7 @@ class WebSocketProxy(websockifyserver.WebSockifyServer):
|
|||
self.heartbeat = kwargs.pop('heartbeat', None)
|
||||
|
||||
self.token_plugin = kwargs.pop('token_plugin', None)
|
||||
self.host_token = kwargs.pop('host_token', None)
|
||||
self.auth_plugin = kwargs.pop('auth_plugin', None)
|
||||
|
||||
# Last 3 timestamps command was run
|
||||
|
@ -451,6 +460,9 @@ def websockify_init():
|
|||
parser.add_option("--token-source", default=None, metavar="ARG",
|
||||
help="an argument to be passed to the token plugin "
|
||||
"on instantiation")
|
||||
parser.add_option("--host-token", action="store_true",
|
||||
help="use the host HTTP header as token instead of the "
|
||||
"token URL query parameter")
|
||||
parser.add_option("--auth-plugin", default=None, metavar="CLASS",
|
||||
help="use a Python class, usually one from websockify.auth_plugins, "
|
||||
"such as BasicHTTPAuth, to determine if a connection is allowed")
|
||||
|
@ -517,6 +529,9 @@ def websockify_init():
|
|||
if opts.token_source and not opts.token_plugin:
|
||||
parser.error("You must use --token-plugin to use --token-source")
|
||||
|
||||
if opts.host_token and not opts.token_plugin:
|
||||
parser.error("You must use --token-plugin to use --host-token")
|
||||
|
||||
if opts.auth_source and not opts.auth_plugin:
|
||||
parser.error("You must use --auth-plugin to use --auth-source")
|
||||
|
||||
|
|
|
@ -93,6 +93,7 @@ class WebSockifyRequestHandler(WebSocketRequestHandler, SimpleHTTPRequestHandler
|
|||
self.file_only = getattr(server, "file_only", False)
|
||||
self.traffic = getattr(server, "traffic", False)
|
||||
self.web_auth = getattr(server, "web_auth", False)
|
||||
self.host_token = getattr(server, "host_token", False)
|
||||
|
||||
self.logger = getattr(server, "logger", None)
|
||||
if self.logger is None:
|
||||
|
|
Loading…
Reference in New Issue