Introduce Token Plugins
Token plugins provide a generic interface for transforming a token into a `(host, port)` tuple. The plugin name is specified using the '--token-plugin' option, and may either be the name of a class from `websockify.token_plugins`, or a fully qualified python path to the token plugin class (see below). An optional plugin parameter can be specified using the '--token-source' option (a value of `None` will be used if no '--token-source' option is passed). Token plugins should inherit from `websockify.token_plugins.BasePlugin`, and should implement the `lookup(token)` method. The value of the '--token-source' option is available as `self.source`. Several plugins are included by default. The `ReadOnlyTokenFile` and `TokenFile` plugins implement functionality from '--target-config' (with the former only reading the file(s) once, and the latter reading them every time). The 'BaseTokenAPI' plugin fetches the value from an API, returning the result of `process_result(response_object)`. By default, `process_result` simply returns the text of the response, but may be overriden. The `JSONTokenAPI` does just this, returning the 'host' and 'port' values from the response JSON object. The old '--target-config' option is now deprecated, and maps to the `TokenFile` plugin under the hood. Also-Authored-By: James Portman (@james-portman) Closes #157
This commit is contained in:
Solly Ross
parent
23045cb212
commit
69a8b928aa
|
|
@ -0,0 +1,78 @@
|
|||
import os
|
||||
|
||||
class BasePlugin(object):
|
||||
def __init__(self, src):
|
||||
self.source = src
|
||||
|
||||
def lookup(self, token):
|
||||
return None
|
||||
|
||||
|
||||
class ReadOnlyTokenFile(BasePlugin):
|
||||
# source is a token file with lines like
|
||||
# token: host:port
|
||||
# or a directory of such files
|
||||
def _load_targets(self):
|
||||
if os.path.isdir(self.source):
|
||||
cfg_files = [os.path.join(self.source, f) for
|
||||
f in os.listdir(self.source)]
|
||||
else:
|
||||
cfg_files = [self.source]
|
||||
|
||||
self._targets = {}
|
||||
for f in cfg_files:
|
||||
for line in [l.strip() for l in open(f).readlines()]:
|
||||
if line and not line.startswith('#'):
|
||||
tok, target = line.split(': ')
|
||||
self._targets[tok] = target.strip().split(':')
|
||||
|
||||
def lookup(self, token):
|
||||
if self._targets is None:
|
||||
self._load_targets()
|
||||
|
||||
if token in self._targets:
|
||||
return self._targets[token]
|
||||
else:
|
||||
return None
|
||||
|
||||
|
||||
# the above one is probably more efficient, but this one is
|
||||
# more backwards compatible (although in most cases
|
||||
# ReadOnlyTokenFile should suffice)
|
||||
class TokenFile(ReadOnlyTokenFile):
|
||||
# source is a token file with lines like
|
||||
# token: host:port
|
||||
# or a directory of such files
|
||||
def lookup(self, token):
|
||||
self._load_targets()
|
||||
|
||||
return super(TokenFile, self).lookup(token)
|
||||
|
||||
|
||||
class BaseTokenAPI(BasePlugin):
|
||||
# source is a url with a '%s' in it where the token
|
||||
# should go
|
||||
|
||||
# we import things on demand so that other plugins
|
||||
# in this file can be used w/o unecessary dependencies
|
||||
|
||||
def process_result(self, resp):
|
||||
return resp.text.split(':')
|
||||
|
||||
def lookup(self, token):
|
||||
import requests
|
||||
|
||||
resp = requests.get(self.source % token)
|
||||
|
||||
if resp.ok:
|
||||
return self.process_result(resp)
|
||||
else:
|
||||
return None
|
||||
|
||||
|
||||
class JSONTokenApi(BaseTokenAPI):
|
||||
# source is a url with a '%s' in it where the token
|
||||
# should go
|
||||
|
||||
def process_result(self, resp):
|
||||
return (resp.json['host'], resp.json['port'])
|
||||
|
|
@ -44,8 +44,8 @@ Traffic Legend:
|
|||
"""
|
||||
# Checks if we receive a token, and look
|
||||
# for a valid target for it then
|
||||
if self.server.target_cfg:
|
||||
(self.server.target_host, self.server.target_port) = self.get_target(self.server.target_cfg, self.path)
|
||||
if self.server.token_plugin:
|
||||
(self.server.target_host, self.server.target_port) = self.get_target(self.server.token_plugin, self.path)
|
||||
|
||||
# Connect to the target
|
||||
if self.server.wrap_cmd:
|
||||
|
|
@ -73,15 +73,15 @@ Traffic Legend:
|
|||
if tsock:
|
||||
tsock.shutdown(socket.SHUT_RDWR)
|
||||
tsock.close()
|
||||
if self.verbose:
|
||||
if self.verbose:
|
||||
self.log_message("%s:%s: Closed target",
|
||||
self.server.target_host, self.server.target_port)
|
||||
raise
|
||||
|
||||
def get_target(self, target_cfg, path):
|
||||
def get_target(self, target_plugin, path):
|
||||
"""
|
||||
Parses the path, extracts a token, and looks for a valid
|
||||
target for that token in the configuration file(s). Sets
|
||||
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
|
||||
"""
|
||||
# The files in targets contain the lines
|
||||
|
|
@ -91,31 +91,16 @@ Traffic Legend:
|
|||
args = parse_qs(urlparse(path)[4]) # 4 is the query from url
|
||||
|
||||
if not 'token' in args or not len(args['token']):
|
||||
raise self.EClose("Token not present")
|
||||
raise self.server.EClose("Token not present")
|
||||
|
||||
token = args['token'][0].rstrip('\n')
|
||||
|
||||
# target_cfg can be a single config file or directory of
|
||||
# config files
|
||||
if os.path.isdir(target_cfg):
|
||||
cfg_files = [os.path.join(target_cfg, f)
|
||||
for f in os.listdir(target_cfg)]
|
||||
result_pair = target_plugin.lookup(token)
|
||||
|
||||
if result_pair is not None:
|
||||
return result_pair
|
||||
else:
|
||||
cfg_files = [target_cfg]
|
||||
|
||||
targets = {}
|
||||
for f in cfg_files:
|
||||
for line in [l.strip() for l in open(f).readlines()]:
|
||||
if line and not line.startswith('#'):
|
||||
ttoken, target = line.split(': ')
|
||||
targets[ttoken] = target.strip()
|
||||
|
||||
self.vmsg("Target config: %s" % repr(targets))
|
||||
|
||||
if token in targets:
|
||||
return targets[token].split(':')
|
||||
else:
|
||||
raise self.EClose("Token '%s' not found" % token)
|
||||
raise self.server.EClose("Token '%s' not found" % token)
|
||||
|
||||
def do_proxy(self, target):
|
||||
"""
|
||||
|
|
@ -147,7 +132,7 @@ Traffic Legend:
|
|||
|
||||
if closed:
|
||||
# TODO: What about blocking on client socket?
|
||||
if self.verbose:
|
||||
if self.verbose:
|
||||
self.log_message("%s:%s: Client closed connection",
|
||||
self.server.target_host, self.server.target_port)
|
||||
raise self.CClose(closed['code'], closed['reason'])
|
||||
|
|
@ -195,7 +180,23 @@ class WebSocketProxy(websocket.WebSocketServer):
|
|||
self.wrap_mode = kwargs.pop('wrap_mode', None)
|
||||
self.unix_target = kwargs.pop('unix_target', None)
|
||||
self.ssl_target = kwargs.pop('ssl_target', None)
|
||||
self.target_cfg = kwargs.pop('target_cfg', None)
|
||||
|
||||
token_plugin = kwargs.pop('token_plugin', None)
|
||||
token_source = kwargs.pop('token_source', None)
|
||||
|
||||
if token_plugin is not None:
|
||||
if '.' not in token_plugin:
|
||||
token_plugin = 'websockify.token_plugins.%s' % token_plugin
|
||||
|
||||
token_plugin_module, token_plugin_cls = token_plugin.rsplit('.', 1)
|
||||
|
||||
__import__(token_plugin_module)
|
||||
token_plugin_cls = getattr(sys.modules[token_plugin_module], token_plugin_cls)
|
||||
|
||||
self.token_plugin = token_plugin_cls(token_source)
|
||||
else:
|
||||
self.token_plugin = None
|
||||
|
||||
# Last 3 timestamps command was run
|
||||
self.wrap_times = [0, 0, 0]
|
||||
|
||||
|
|
@ -251,9 +252,9 @@ class WebSocketProxy(websocket.WebSocketServer):
|
|||
else:
|
||||
dst_string = "%s:%s" % (self.target_host, self.target_port)
|
||||
|
||||
if self.target_cfg:
|
||||
msg = " - proxying from %s:%s to targets in %s" % (
|
||||
self.listen_host, self.listen_port, self.target_cfg)
|
||||
if self.token_plugin:
|
||||
msg = " - proxying from %s:%s to targets generated by %s" % (
|
||||
self.listen_host, self.listen_port, type(self.token_plugin).__name__)
|
||||
else:
|
||||
msg = " - proxying from %s:%s to %s" % (
|
||||
self.listen_host, self.listen_port, dst_string)
|
||||
|
|
@ -352,20 +353,41 @@ def websockify_init():
|
|||
parser.add_option("--prefer-ipv6", "-6",
|
||||
action="store_true", dest="source_is_ipv6",
|
||||
help="prefer IPv6 when resolving source_addr")
|
||||
parser.add_option("--libserver", action="store_true",
|
||||
help="use Python library SocketServer engine")
|
||||
parser.add_option("--target-config", metavar="FILE",
|
||||
dest="target_cfg",
|
||||
help="Configuration file containing valid targets "
|
||||
"in the form 'token: host:port' or, alternatively, a "
|
||||
"directory containing configuration files of this form")
|
||||
parser.add_option("--libserver", action="store_true",
|
||||
help="use Python library SocketServer engine")
|
||||
"directory containing configuration files of this form "
|
||||
"(DEPRECATED: use `--token-plugin TokenFile --token-source "
|
||||
" path/to/token/file` instead)")
|
||||
parser.add_option("--token-plugin", default=None, metavar="PLUGIN",
|
||||
help="use the given Python class to process tokens "
|
||||
"into host:port pairs")
|
||||
parser.add_option("--token-source", default=None, metavar="ARG",
|
||||
help="an argument to be passed to the token plugin"
|
||||
"on instantiation")
|
||||
|
||||
(opts, args) = parser.parse_args()
|
||||
|
||||
if opts.verbose:
|
||||
logging.getLogger(WebSocketProxy.log_prefix).setLevel(logging.DEBUG)
|
||||
|
||||
if opts.token_source and not opts.token_plugin:
|
||||
parser.error("You must use --token-plugin to use --token-source")
|
||||
|
||||
# Transform to absolute path as daemon may chdir
|
||||
if opts.target_cfg:
|
||||
opts.target_cfg = os.path.abspath(opts.target_cfg)
|
||||
|
||||
if opts.target_cfg:
|
||||
opts.token_plugin = 'TokenFile'
|
||||
opts.token_source = opts.target_cfg
|
||||
del opts.target_cfg
|
||||
|
||||
# Sanity checks
|
||||
if len(args) < 2 and not (opts.target_cfg or opts.unix_target):
|
||||
if len(args) < 2 and not (opts.token_plugin or opts.unix_target):
|
||||
parser.error("Too few arguments")
|
||||
if sys.argv.count('--'):
|
||||
opts.wrap_cmd = args[1:]
|
||||
|
|
@ -390,7 +412,7 @@ def websockify_init():
|
|||
try: opts.listen_port = int(opts.listen_port)
|
||||
except: parser.error("Error parsing listen port")
|
||||
|
||||
if opts.wrap_cmd or opts.unix_target or opts.target_cfg:
|
||||
if opts.wrap_cmd or opts.unix_target or opts.token_plugin:
|
||||
opts.target_host = None
|
||||
opts.target_port = None
|
||||
else:
|
||||
|
|
@ -402,10 +424,6 @@ def websockify_init():
|
|||
try: opts.target_port = int(opts.target_port)
|
||||
except: parser.error("Error parsing target port")
|
||||
|
||||
# Transform to absolute path as daemon may chdir
|
||||
if opts.target_cfg:
|
||||
opts.target_cfg = os.path.abspath(opts.target_cfg)
|
||||
|
||||
# Create and start the WebSockets proxy
|
||||
libserver = opts.libserver
|
||||
del opts.libserver
|
||||
|
|
@ -456,8 +474,8 @@ class LibProxyServer(ForkingMixIn, HTTPServer):
|
|||
|
||||
if web:
|
||||
os.chdir(web)
|
||||
|
||||
HTTPServer.__init__(self, (listen_host, listen_port),
|
||||
|
||||
HTTPServer.__init__(self, (listen_host, listen_port),
|
||||
RequestHandlerClass)
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Reference in New Issue