Skip to content

Commit 4aee6c5

Browse files
ahanel13ahanel13
committed
Added an additional variable plugin_provided_url to handle dynamic urls.
This allows a user to create an auth plugin that retrieves a new url for each websocket connection.
1 parent 1b8b720 commit 4aee6c5

File tree

3 files changed

+61
-44
lines changed

3 files changed

+61
-44
lines changed

‎wsrepl/MessageHandler.py

Lines changed: 29 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -17,33 +17,43 @@ class MessageHandler:
1717
def __init__(self,
1818
app: textual.app.App,
1919
url: str,
20-
user_agent: str | None = None,
21-
origin: str | None = None,
22-
cookies: list[str] | None = None,
23-
headers: list[str] | None = None,
24-
headers_file: str | None = None,
25-
ping_interval: int | float = 24,
26-
hide_ping_pong: bool = False,
27-
ping_0x1_interval: int | float = 24,
28-
ping_0x1_payload: str | None = None,
29-
pong_0x1_payload: str | None = None,
30-
hide_0x1_ping_pong: bool = False,
31-
reconnect_interval: int = 0,
32-
proxy: str | None = None,
33-
verify_tls: bool = True,
34-
initial_msgs_file: str | None = None,
35-
plugin_path: str | None = None) -> None:
20+
user_agent: str | None = None,
21+
origin: str | None = None,
22+
cookies: list[str] | None = None,
23+
headers: list[str] | None = None,
24+
headers_file: str | None = None,
25+
ping_interval: int | float = 24,
26+
hide_ping_pong: bool = False,
27+
ping_0x1_interval: int | float = 24,
28+
ping_0x1_payload: str | None = None,
29+
pong_0x1_payload: str | None = None,
30+
hide_0x1_ping_pong: bool = False,
31+
reconnect_interval: int = 0,
32+
proxy: str | None = None,
33+
verify_tls: bool = True,
34+
initial_msgs_file: str | None = None,
35+
plugin_path: str | None = None,
36+
plugin_provided_url: bool | None = None) -> None:
3637

3738
self.app = app
38-
self.plugin = load_plugin(plugin_path)(message_handler=self)
39+
if(plugin_provided_url):
40+
self.plugin = load_plugin(plugin_path)(message_handler=self)
41+
try:
42+
self.url = self.plugin.url
43+
print("URL from plugin = " + self.url)
44+
except:
45+
print("Failed to get URL path from plugin. Exiting...")
46+
exit()
47+
else:
48+
self.plugin = load_plugin(plugin_path)(message_handler=self)
3949
self.initial_messages: list[WSMessage] = self._load_initial_messages(initial_msgs_file)
4050
processed_headers: OrderedDict = self._process_headers(headers, headers_file, user_agent, origin, cookies)
4151

4252
self._ws = WebsocketConnection(
4353
# Stuff WebsocketConnection needs to call back to us
4454
async_handler=self,
4555
# WebSocketApp args
46-
url=url,
56+
url=self.url,
4757
header=processed_headers
4858
)
4959

@@ -61,6 +71,7 @@ def __init__(self,
6171
'reconnect': reconnect_interval
6272
}
6373

74+
print("test")
6475
self.is_stopped = threading.Event()
6576

6677
# Regular ping thread, conforming to RFC 6455 (ping uses opcode 0x9, pong uses 0xA, data is arbitrary but must be the same)

‎wsrepl/app.py

Lines changed: 20 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -26,26 +26,27 @@ def __init__(self,
2626
# URL is the only required argument
2727
url: str,
2828
# UI settings
29-
small: bool = False,
29+
small: bool = False,
3030
# Websocket settings
31-
user_agent: str | None = None,
32-
origin: str | None = None,
33-
cookies: list[str] | None = None,
34-
headers: list[str] | None = None,
35-
headers_file: str | None = None,
36-
ping_interval: int | float = 24, # 0 -> disable auto ping
37-
hide_ping_pong: bool = False,
38-
ping_0x1_interval: int | float = 24, # 0 -> disable fake ping
39-
ping_0x1_payload: str | None = None,
40-
pong_0x1_payload: str | None = None,
41-
hide_0x1_ping_pong: bool = False,
42-
reconnect_interval: int = 0,
43-
proxy: str | None = None,
44-
verify_tls: bool = True,
31+
user_agent: str | None = None,
32+
origin: str | None = None,
33+
cookies: list[str] | None = None,
34+
headers: list[str] | None = None,
35+
headers_file: str | None = None,
36+
ping_interval: int | float = 24, # 0 -> disable auto ping
37+
hide_ping_pong: bool = False,
38+
ping_0x1_interval: int | float = 24, # 0 -> disable fake ping
39+
ping_0x1_payload: str | None = None,
40+
pong_0x1_payload: str | None = None,
41+
hide_0x1_ping_pong: bool = False,
42+
reconnect_interval: int = 0,
43+
proxy: str | None = None,
44+
verify_tls: bool = True,
4545
# Other
46-
initial_msgs_file: str | None = None,
47-
plugin_path: str | None = None,
48-
verbosity: int = 3) -> None:
46+
initial_msgs_file: str | None = None,
47+
plugin_path: str | None = None,
48+
plugin_provided_url: bool | None = None,
49+
verbosity: int = 3) -> None:
4950
super().__init__()
5051

5152
# Small UI
@@ -62,7 +63,7 @@ def __init__(self,
6263
ping_0x1_interval=ping_0x1_interval, ping_0x1_payload=ping_0x1_payload, pong_0x1_payload=pong_0x1_payload,
6364
hide_0x1_ping_pong=hide_0x1_ping_pong,
6465
reconnect_interval=reconnect_interval, proxy=proxy, verify_tls=verify_tls,
65-
initial_msgs_file=initial_msgs_file, plugin_path=plugin_path
66+
initial_msgs_file=initial_msgs_file, plugin_path=plugin_path, plugin_provided_url=plugin_provided_url
6667
)
6768

6869
# These are set in compose()

‎wsrepl/cli.py

Lines changed: 12 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -33,18 +33,22 @@
3333
parser.add_argument('-r', '--reconnect-interval', type=int, default=2, help='Reconnect interval (seconds, default: 2)')
3434
parser.add_argument('-I', '--initial-messages', type=str, help='Send the messages from this file on connect')
3535
parser.add_argument('-P', '--plugin', type=str, help='Plugin file to load')
36+
parser.add_argument( '--plugin-provided-url', action='store_true', default=False, help='Plugin file to load')
3637
parser.add_argument('-v', '--verbose', type=int, default=3, help='Verbosity level, 1-4 default: 3 (errors, warnings, info), 4 adds debug')
3738

3839
def cli():
3940
args = parser.parse_args()
4041
url = args.url or args.url_positional
41-
# Check and modify the URL protocol if necessary
42-
if url.startswith('http://'):
43-
url = url.replace('http://', 'ws://', 1)
44-
elif url.startswith('https://'):
45-
url = url.replace('https://', 'wss://', 1)
46-
elif not url.startswith(('ws://', 'wss://')):
47-
parser.error('Invalid protocol. Supported protocols are http://, https://, ws://, and wss://.')
42+
if url and not args.plugin_provided_url:
43+
# Check and modify the URL protocol if necessary
44+
if url.startswith('http://'):
45+
return url.replace('http://', 'ws://', 1)
46+
elif url.startswith('https://'):
47+
return url.replace('https://', 'wss://', 1)
48+
elif not url.startswith(('ws://', 'wss://')):
49+
parser.error('Invalid protocol. Supported protocols are http://, https://, ws://, and wss://.')
50+
elif args.plugin_provided_url and not args.plugin:
51+
parser.error('Please provide a WebSocket URL using either -u as a positional argument or use --plugin-provided-url if the WebSocket URL provided in a plugin')
4852

4953
app = WSRepl(
5054
url=url,
@@ -65,6 +69,7 @@ def cli():
6569
verify_tls=not args.insecure,
6670
initial_msgs_file=args.initial_messages,
6771
plugin_path=args.plugin,
72+
plugin_provided_url=args.plugin_provided_url,
6873
verbosity=args.verbose,
6974
)
7075
app.run()

0 commit comments

Comments
 (0)