From 20b8c798c191fa9aefb0690da2fa42a0b2bebd19 Mon Sep 17 00:00:00 2001 From: Gerhard Schaden Date: Sun, 26 Feb 2023 14:17:12 +0100 Subject: [PATCH] Add traffic plugins --- README.md | 4 ++++ websockify/token_plugins.py | 2 +- websockify/traffic_plugins.py | 8 ++++++++ websockify/websocketproxy.py | 27 ++++++++++++++++++++++++++- websockify/websockifyserver.py | 1 + 5 files changed, 40 insertions(+), 2 deletions(-) create mode 100644 websockify/traffic_plugins.py diff --git a/README.md b/README.md index c5d14cc..182144d 100644 --- a/README.md +++ b/README.md @@ -110,6 +110,10 @@ These are not necessary for the basic operation. options, where CLASS is usually one from token_plugins.py and ARG is the plugin's configuration. +* Traffic plugins: an instance per connection to intercept and modify traffic + This functionality is activated with the `--traffic-plugin CLASS`, where + CLASS is an implementation of the traffic_plugin BasePlugin. + ### Other implementations of websockify The primary implementation of websockify is in python. There are diff --git a/websockify/token_plugins.py b/websockify/token_plugins.py index d42414e..b61a848 100644 --- a/websockify/token_plugins.py +++ b/websockify/token_plugins.py @@ -95,7 +95,7 @@ class JSONTokenApi(BaseTokenAPI): def process_result(self, resp): resp_json = resp.json() - return (resp_json['host'], resp_json['port']) + return (resp_json['host'], resp_json['port'], resp_json) class JWTTokenApi(BasePlugin): diff --git a/websockify/traffic_plugins.py b/websockify/traffic_plugins.py new file mode 100644 index 0000000..de6abbf --- /dev/null +++ b/websockify/traffic_plugins.py @@ -0,0 +1,8 @@ +class BasePlugin(): + def __init__(self, handler, tsock): + self.handler = handler + self.tsock = tsock + def from_client(self, s): + return s + def from_target(self, s): + return s diff --git a/websockify/websocketproxy.py b/websockify/websocketproxy.py index 9b4595b..8d79d7b 100644 --- a/websockify/websocketproxy.py +++ b/websockify/websocketproxy.py @@ -48,13 +48,15 @@ Traffic Legend: if not self.server.token_plugin: return - host, port = self.get_target(self.server.token_plugin) + host, port, *rest = self.get_target(self.server.token_plugin) if host == 'unix_socket': self.server.unix_target = port else: self.server.target_host = host self.server.target_port = port + if len(rest) > 0: # rest is an array of all remaining parts + self.target_attribtues = rest[0] def auth_connection(self): if not self.server.auth_plugin: @@ -122,6 +124,8 @@ Traffic Legend: # Start proxying try: + if self.traffic_plugin: # replace class with instance + self.traffic_plugin = self.traffic_plugin(self, tsock) self.do_proxy(tsock) finally: if tsock: @@ -246,6 +250,9 @@ Traffic Legend: if target in outs: # Send queued client data to the target dat = tqueue.pop(0) + if self.traffic_plugin: + dat = self.traffic_plugin.from_client(dat) + if not dat: continue sent = target.send(dat) if sent == len(dat): self.print_traffic(">") @@ -258,6 +265,9 @@ Traffic Legend: if target in ins: # Receive target data, encode it and queue for client buf = target.recv(self.buffer_size) + if self.traffic_plugin: + buf = self.traffic_plugin.from_target(buf) + if not buf: continue if len(buf) == 0: # Target socket closed, flushing queues and closing client-side websocket @@ -298,6 +308,7 @@ class WebSocketProxy(websockifyserver.WebSockifyServer): self.token_plugin = kwargs.pop('token_plugin', None) self.host_token = kwargs.pop('host_token', None) self.auth_plugin = kwargs.pop('auth_plugin', None) + self.traffic_plugin = kwargs.pop('traffic_plugin', None) # Last 3 timestamps command was run self.wrap_times = [0, 0, 0] @@ -537,6 +548,9 @@ def websockify_init(): 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") + parser.add_option("--traffic-plugin", default=None, metavar="CLASS", + help="use a Python class, usually one from websockify.traffic_plugins, " + "to modify the traffic") parser.add_option("--auth-source", default=None, metavar="ARG", help="an argument to be passed to the auth plugin " "on instantiation") @@ -732,6 +746,17 @@ def websockify_init(): del opts.auth_source + if opts.traffic_plugin is not None: + if '.' not in opts.traffic_plugin: + opts.traffic_plugin = 'websockify.traffic_plugins.%s' % opts.traffic_plugin + + traffic_plugin_module, traffic_plugin_cls = opts.traffic_plugin.rsplit('.', 1) + + __import__(traffic_plugin_module) + traffic_plugin_cls = getattr(sys.modules[traffic_plugin_module], traffic_plugin_cls) + + opts.traffic_plugin = traffic_plugin_cls + # Create and start the WebSockets proxy libserver = opts.libserver del opts.libserver diff --git a/websockify/websockifyserver.py b/websockify/websockifyserver.py index 74f9f53..2c8715a 100644 --- a/websockify/websockifyserver.py +++ b/websockify/websockifyserver.py @@ -84,6 +84,7 @@ class WebSockifyRequestHandler(WebSocketRequestHandlerMixIn, SimpleHTTPRequestHa if self.logger is None: self.logger = WebSockifyServer.get_logger() + self.traffic_plugin = getattr(server, "traffic_plugin", None) super().__init__(req, addr, server) def log_message(self, format, *args):