├── .gitignore ├── README.md ├── docs ├── demo-1.py ├── demo-2.py ├── demo-3.py ├── demo-4.py ├── hooks-mermaid.png └── screenshot.png ├── pyproject.toml ├── requirements.txt └── wsrepl ├── MessageHandler.py ├── Ping0x1Thread.py ├── PingThread.py ├── Plugin.py ├── WSMessage.py ├── WebsocketConnection.py ├── __init__.py ├── __main__.py ├── app.css ├── app.py ├── cli.py ├── constants.py ├── log.py ├── utils.py └── widgets ├── CopyButton.py ├── DirectionSign.py ├── FormattedMessage.py ├── History.py ├── HistoryIndex.py ├── HistoryRow.py ├── Parent.py └── __init__.py /.gitignore: -------------------------------------------------------------------------------- 1 | poetry.lock 2 | __pycache__ 3 | venv 4 | dist 5 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # `wsrepl` - Websocket REPL for pentesters 2 | 3 | `wsrepl` is an interactive websocket REPL designed specifically for penetration testing. It provides an interface for observing incoming websocket messages and sending new ones, with an easy-to-use framework for automating this communication. 4 | 5 | ![Screenshot](docs/screenshot.png) 6 | 7 | ## Features 8 | 9 | - Interactively send and receive websocket messages 10 | - Customize headers, ping/pong messages, and other parameters 11 | - Handle SSL verification and reconnections 12 | - Plug-in support for automating complex interaction scenarios 13 | - Full logging and message history 14 | - Supports curl command line arguments for easy onboarding from Developer Tools or Burp Suite (use 'Copy as Curl' menu and replace `curl` with `wsrepl`) 15 | 16 | ## Installation 17 | 18 | You can download and install wsrepl using pip: 19 | 20 | ``` 21 | pip install wsrepl 22 | ``` 23 | 24 | Alternatively, you can clone this repository and install it from source: 25 | 26 | ``` 27 | git clone https://github.com/doyensec/wsrepl 28 | cd wsrepl 29 | pip install . 30 | ``` 31 | 32 | ## Usage 33 | 34 | The basic command for starting wsrepl is as follows: 35 | 36 | ``` 37 | wsrepl -u URL 38 | ``` 39 | 40 | Replace URL with your target websocket URL, e.g. `wss://echo.websocket.org`. For more options and settings, you can use the -h or --help option: 41 | 42 | ``` 43 | usage: wsrepl [-h] [-u URL] [-i] [-s] [-k] [-X REQUEST] [-H HEADER] [-b COOKIE] [--compressed] [-S] [-A USER_AGENT] 44 | [-O ORIGIN] [-F HEADERS_FILE] [--no-native-ping] [--ping-interval PING_INTERVAL] [--hide-ping-pong] 45 | [--ping-0x1-interval PING_0X1_INTERVAL] [--ping-0x1-payload PING_0X1_PAYLOAD] 46 | [--pong-0x1-payload PONG_0X1_PAYLOAD] [--hide-0x1-ping-pong] [-t TTL] [-p HTTP_PROXY] 47 | [-r RECONNECT_INTERVAL] [-I INITIAL_MESSAGES] [-P PLUGIN] [--plugin-provided-url] [-v VERBOSE] 48 | [url_positional] 49 | 50 | Websocket Client 51 | 52 | positional arguments: 53 | url_positional Websocket URL (e.g. wss://echo.websocket.org) 54 | 55 | options: 56 | -h, --help show this help message and exit 57 | -u URL, --url URL Websocket URL (e.g. wss://echo.websocket.org) 58 | -i, --include No effect, just for curl compatibility 59 | -s, --silent No effect, just for curl compatibility 60 | -k, --insecure Disable SSL verification 61 | -X REQUEST, --request REQUEST 62 | No effect, just for curl compatibility 63 | -H HEADER, --header HEADER 64 | Additional header (e.g. "X-Header: value"), can be used multiple times 65 | -b COOKIE, --cookie COOKIE 66 | Cookie header (e.g. "name=value"), can be used multiple times 67 | --compressed No effect, just for curl compatibility 68 | -S, --small Smaller UI 69 | -A USER_AGENT, --user-agent USER_AGENT 70 | User-Agent header 71 | -O ORIGIN, --origin ORIGIN 72 | Origin header 73 | -F HEADERS_FILE, --headers-file HEADERS_FILE 74 | Additional headers file (e.g. "headers.txt") 75 | --no-native-ping Disable native ping/pong messages 76 | --ping-interval PING_INTERVAL 77 | Ping interval (seconds) 78 | --hide-ping-pong Hide ping/pong messages 79 | --ping-0x1-interval PING_0X1_INTERVAL 80 | Fake ping (0x1 opcode) interval (seconds) 81 | --ping-0x1-payload PING_0X1_PAYLOAD 82 | Fake ping (0x1 opcode) payload 83 | --pong-0x1-payload PONG_0X1_PAYLOAD 84 | Fake pong (0x1 opcode) payload 85 | --hide-0x1-ping-pong Hide fake ping/pong messages 86 | -t TTL, --ttl TTL Heartbeet interval (seconds) 87 | -p HTTP_PROXY, --http-proxy HTTP_PROXY 88 | HTTP Proxy Address (e.g. 127.0.0.1:8080) 89 | -r RECONNECT_INTERVAL, --reconnect-interval RECONNECT_INTERVAL 90 | Reconnect interval (seconds, default: 2) 91 | -I INITIAL_MESSAGES, --initial-messages INITIAL_MESSAGES 92 | Send the messages from this file on connect 93 | -P PLUGIN, --plugin PLUGIN 94 | Plugin file to load 95 | --plugin-provided-url 96 | Indicates if plugin provided dynamic url for websockets 97 | -v VERBOSE, --verbose VERBOSE 98 | Verbosity level, 1-4 default: 3 (errors, warnings, info), 4 adds debug 99 | ``` 100 | 101 | ## Automating with Plugins 102 | 103 | To automate your websocket communication, you can create a Python plugin by extending the [Plugin](wsrepl/Plugin.py) class in wsrepl. This class allows you to define various hooks that are triggered during different stages of the websocket communication. 104 | 105 | Here is an outline of how to define a plugin: 106 | 107 | ```python 108 | class MyPlugin(Plugin): 109 | # Messages that will be sent to the server on (re-)connect. 110 | messages = ["message1", "message2", "message3"] 111 | 112 | def init(self): 113 | # This method is called when the plugin is loaded. 114 | # Use it to set initial settings or generate self.messages dynamically. 115 | pass 116 | 117 | async def on_connect(self): 118 | # This method is called when the websocket connection is established. 119 | pass 120 | 121 | async def on_message_received(self, message: WSMessage): 122 | # This method is called when a message is received from the server. 123 | pass 124 | 125 | # ... Other hooks can be defined here. 126 | ``` 127 | 128 | Refer to the source of [Plugin class](wsrepl/Plugin.py) for the full list of hooks you can use and what they do: 129 | 130 | ![Plugin Hooks](docs/hooks-mermaid.png) 131 | 132 | ### Example Plugin 133 | 134 | Here is an example of a plugin that sends a predefined authentication message to a server: 135 | 136 | ```python 137 | from wsrepl import Plugin 138 | import json 139 | 140 | class AuthPlugin(Plugin): 141 | def init(self): 142 | auth_message = { 143 | "messageType": "auth", 144 | "auth": { 145 | "user": "user-1234"", 146 | "password": "password-1234" 147 | } 148 | } 149 | self.messages = [json.dumps(auth_message)] 150 | ``` 151 | 152 | This plugin can be used by specifying it when running wsrepl: 153 | 154 | ``` 155 | wsrepl -u URL -P auth_plugin.py 156 | ``` 157 | 158 | Replace URL with your target websocket URL and auth_plugin.py with the path to the Python file containing your plugin. 159 | 160 | [docs/](./docs/) directory contains a few more example plugins. 161 | 162 | ## Contributing 163 | 164 | Contributions to wsrepl are welcome! Please, [create an issue](https://github.com/doyensec/wsrepl/issues) or submit a pull request if you have any ideas or suggestions. In particular, adding more plugin examples would be very helpful. 165 | 166 | ## Credits 167 | 168 | This project has been sponsored by [Doyensec LLC](https://www.doyensec.com/). 169 | 170 | ![Doyensec Research](https://raw.githubusercontent.com/doyensec/inql/master/docs/doyensec_logo.svg "Doyensec Logo") 171 | 172 | -------------------------------------------------------------------------------- /docs/demo-1.py: -------------------------------------------------------------------------------- 1 | from wsrepl import Plugin 2 | 3 | MESSAGES = [ 4 | "hello", 5 | "world" 6 | ] 7 | 8 | class Demo(Plugin): 9 | 10 | def init(self): 11 | """Initialization method that is called when the plugin is loaded. 12 | 13 | In this demo, we're simply populating the self.messages list with predefined messages. 14 | These messages will be sent to the server once a websocket connection is established. 15 | """ 16 | self.messages = MESSAGES 17 | -------------------------------------------------------------------------------- /docs/demo-2.py: -------------------------------------------------------------------------------- 1 | from wsrepl import Plugin 2 | from wsrepl.WSMessage import WSMessage 3 | 4 | import json 5 | import requests 6 | 7 | class Demo(Plugin): 8 | 9 | def init(self): 10 | """Initialization method that is called when the plugin is loaded. 11 | 12 | In this example, we are dynamically populating self.messages list by getting a session token from a HTTP endpoint. 13 | This is a typical scenario when interacting with APIs that require user authentication before the WebSocket connection is established. 14 | """ 15 | 16 | # Here we simulate an API request to get a session token by supplying a username and password. 17 | # For the demo, we're using a dummy endpoint "https://hb.cran.dev/uuid" that returns a UUID. 18 | # In a real-life scenario, replace this with your own authentication endpoint and provide necessary credentials. 19 | token = requests.get("https://hb.cran.dev/uuid").json()["uuid"] 20 | 21 | # The acquired session token is then used to populate self.messages with an authentication message. 22 | # The exact format of this message will depend on your WebSocket server requirements. 23 | self.messages = [ 24 | json.dumps({ 25 | "auth": "session", 26 | "sessionId": token 27 | }) 28 | ] 29 | 30 | async def on_connect(self) -> None: 31 | """Called when the websocket connection is established. 32 | 33 | In this demo, we're simply logging the successful connection event. 34 | """ 35 | self.log.info("Successfully connected to the server with the session token!") 36 | 37 | async def on_message_received(self, message: WSMessage) -> None: 38 | """Called when a (text) message is received from the server. 39 | 40 | In this demo, we're logging the original message received from the server. 41 | """ 42 | self.log.info(f"Received message from server: {message.msg}") 43 | # Further processing of the message can go here 44 | 45 | async def on_error(self, exception: Exception) -> None: 46 | """Called when an error message is received from the server 47 | 48 | In this demo, we're logging the error. 49 | """ 50 | self.log.error(f"An error occurred: {exception}") 51 | # Additional error handling can go here 52 | -------------------------------------------------------------------------------- /docs/demo-3.py: -------------------------------------------------------------------------------- 1 | from wsrepl import Plugin 2 | 3 | import json 4 | from wsrepl.WSMessage import WSMessage 5 | 6 | class Demo(Plugin): 7 | 8 | async def on_message_sent(self, message: WSMessage) -> None: 9 | """This method is called on every message that the user enters into the REPL. 10 | It modifies the message before it's sent to the server by wrapping it into a specific JSON format. 11 | 12 | Here we're demonstrating how to wrap user messages into a more complex structure that some websocket servers might require. 13 | """ 14 | 15 | # Grab the original message entered by the user 16 | original = message.msg 17 | 18 | # Prepare a more complex message structure that our server requires. 19 | # The exact structure here will depend on your websocket server's requirements. 20 | message.msg = json.dumps({ 21 | "type": "message", 22 | "data": { 23 | "text": original 24 | } 25 | }) 26 | 27 | # Short and long versions of the message are used for display purposes in REPL UI. 28 | # By default they are the same as 'message.msg', but here we modify them for better UX. 29 | message.short = original 30 | message.long = message.msg 31 | 32 | 33 | async def on_message_received(self, message: WSMessage) -> None: 34 | """This method is called every time a message is received from the server. 35 | It extracts the core message out of the JSON object received from the server. 36 | 37 | Here we're demonstrating how to unwrap the received messages and display more user-friendly information. 38 | """ 39 | 40 | # Get the original message received from the server 41 | original = message.msg 42 | 43 | try: 44 | # Try to parse the received message and extract meaningful data. 45 | # The exact structure here will depend on your websocket server's responses. 46 | message.short = json.loads(original)["data"]["text"] 47 | except: 48 | # In case of a parsing error, let's inform the user about it in the history view. 49 | message.short = "Error: could not parse message" 50 | 51 | # Show the original message when the user focuses on it in the UI. 52 | message.long = original 53 | -------------------------------------------------------------------------------- /docs/demo-4.py: -------------------------------------------------------------------------------- 1 | from wsrepl import Plugin 2 | import requests 3 | 4 | # Set to None if no proxy required 5 | plugin_proxy = { 6 | "http": "http://127.0.0.1:8080", 7 | "https": "http://127.0.0.1:8080" 8 | } 9 | 10 | class MultiStepAuthDemo(Plugin): 11 | url = None # Required to pass dynamic wss url to MessageHandler.py 12 | 13 | def init(self): 14 | # Step one: Get a JWT 15 | response = self.send_jwt_request() 16 | jwt_token = self.extract_jwt_from_response(response) 17 | if jwt_token: 18 | # Step 2: Get a access/bearer token 19 | response_jwt = self.request_access_token(jwt_token) 20 | access_token = self.extract_access_token_from_response(response_jwt) 21 | if access_token: 22 | # Step 3: Get the dynamic wss link 23 | wss_start = self.get_wss_endpoint(access_token) 24 | self.url = self.extract_url_from_response(wss_start) 25 | 26 | def send_jwt_request(self): 27 | url = "https://example.com/users/auth" 28 | headers = { 29 | "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:109.0) Gecko/20100101 Firefox/112.0", 30 | "Accept": "application/json", 31 | "Accept-Language": "en-US,en;q=0.5", 32 | "Accept-Encoding": "gzip, deflate", 33 | "Content-Type": "application/x-www-form-urlencoded", 34 | "Origin": "https://example-origin.com", 35 | "Referer": "https://example-referer.com/", 36 | "Sec-Fetch-Dest": "empty", 37 | "Sec-Fetch-Mode": "cors", 38 | "Sec-Fetch-Site": "cross-site" 39 | } 40 | 41 | data = { 42 | "clientId": "your-client-id", 43 | "secret": "your-client-secret", 44 | "identity": "user-identity", 45 | "aud": "public", 46 | "isAnonymous": "true", 47 | "sid": "session-id", 48 | "page": "contactus", 49 | "lang": "en_US", 50 | "role": "VISITOR" 51 | } 52 | 53 | response = requests.post(url, headers=headers, data=data, proxies=plugin_proxy, verify=False) 54 | return response 55 | 56 | def extract_jwt_from_response(self, response): 57 | try: 58 | json_data = response.json() 59 | jwt_token = json_data.get("token") 60 | return jwt_token 61 | except ValueError: 62 | return None 63 | 64 | def request_access_token(self, jwt_token): 65 | url = "https://example.com/api/token/jwtgrant" 66 | headers = { 67 | "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:109.0) Gecko/20100101 Firefox/112.0", 68 | "Accept": "*/*", 69 | "Accept-Language": "en-US,en;q=0.5", 70 | "Accept-Encoding": "gzip, deflate", 71 | "Content-Type": "application/json", 72 | "Origin": "https://example-origin.com", 73 | "Referer": "https://example-referer.com/", 74 | "Sec-Fetch-Dest": "empty", 75 | "Sec-Fetch-Mode": "cors", 76 | "Sec-Fetch-Site": "cross-site" 77 | } 78 | 79 | data = { 80 | "assertion": jwt_token, 81 | "botInfo": { 82 | "chatBot": "example-bot", 83 | "botId": "task-bot-id" 84 | } 85 | } 86 | 87 | response = requests.post(url, headers=headers, json=data, proxies=plugin_proxy, verify=False) 88 | return response 89 | 90 | def extract_access_token_from_response(self, response): 91 | try: 92 | json_data = response.json() 93 | access_token = json_data["authorization"]["accessToken"] 94 | return access_token 95 | except (ValueError, KeyError): 96 | return None 97 | 98 | def get_wss_endpoint(self, access_token): 99 | url = "https://example.com/api/chat/start" 100 | headers = { 101 | "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:109.0) Gecko/20100101 Firefox/112.0", 102 | "Accept": "application/json", 103 | "Accept-Language": "en-US,en;q=0.5", 104 | "Accept-Encoding": "gzip, deflate", 105 | "Content-Type": "application/json", 106 | "Authorization": f"Bearer {access_token}", 107 | "Origin": "https://example-origin.com", 108 | "Referer": "https://example-referer.com/", 109 | "Sec-Fetch-Dest": "empty", 110 | "Sec-Fetch-Mode": "cors", 111 | "Sec-Fetch-Site": "cross-site" 112 | } 113 | 114 | data = { 115 | "botInfo": { 116 | "chatBot": "example-bot", 117 | "botId": "task-bot-id" 118 | } 119 | } 120 | 121 | response = requests.post(url, headers=headers, json=data, proxies=plugin_proxy, verify=False) 122 | return response 123 | 124 | def extract_url_from_response(self, response): 125 | try: 126 | json_data = response.json() 127 | url = json_data["endpoint"] 128 | return url 129 | except (ValueError, KeyError): 130 | return None 131 | -------------------------------------------------------------------------------- /docs/hooks-mermaid.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/doyensec/wsrepl/f8c56635119b4edef2e3ac9f21961e7574170ac8/docs/hooks-mermaid.png -------------------------------------------------------------------------------- /docs/screenshot.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/doyensec/wsrepl/f8c56635119b4edef2e3ac9f21961e7574170ac8/docs/screenshot.png -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [tool.poetry] 2 | name = "wsrepl" 3 | version = "0.2.0" 4 | description = "Websocket REPL for pentesters" 5 | authors = ["Andrew Konstantinov <105389353+execveat@users.noreply.github.com>"] 6 | license = "MIT" 7 | readme = "README.md" 8 | homepage = "https://github.com/doyensec/wsrepl" 9 | keywords = ["websocket", "pentesting"] 10 | 11 | [tool.poetry.dependencies] 12 | python = ">3.10,<4.0" 13 | websocket-client = "^1.5.2" 14 | pyperclip = "^1.8.2" 15 | rich = "^13.4.2" 16 | textual = "^0.30.0" 17 | 18 | [build-system] 19 | requires = ["poetry-core"] 20 | build-backend = "poetry.core.masonry.api" 21 | 22 | [tool.poetry.scripts] 23 | wsrepl = "wsrepl:cli" 24 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | importlib-metadata==6.8.0 ; python_full_version > "3.10.0" and python_version < "4.0" 2 | linkify-it-py==2.0.2 ; python_full_version > "3.10.0" and python_version < "4.0" 3 | markdown-it-py==3.0.0 ; python_full_version > "3.10.0" and python_version < "4.0" 4 | markdown-it-py[linkify,plugins]==3.0.0 ; python_full_version > "3.10.0" and python_version < "4.0" 5 | mdit-py-plugins==0.4.0 ; python_full_version > "3.10.0" and python_version < "4.0" 6 | mdurl==0.1.2 ; python_full_version > "3.10.0" and python_version < "4.0" 7 | pygments==2.15.1 ; python_full_version > "3.10.0" and python_version < "4.0" 8 | pyperclip==1.8.2 ; python_full_version > "3.10.0" and python_version < "4.0" 9 | rich==13.4.2 ; python_full_version > "3.10.0" and python_version < "4.0" 10 | textual==0.30.0 ; python_full_version > "3.10.0" and python_version < "4.0" 11 | typing-extensions==4.7.1 ; python_full_version > "3.10.0" and python_version < "4.0" 12 | uc-micro-py==1.0.2 ; python_full_version > "3.10.0" and python_version < "4.0" 13 | websocket-client==1.6.1 ; python_full_version > "3.10.0" and python_version < "4.0" 14 | zipp==3.16.2 ; python_full_version > "3.10.0" and python_version < "4.0" 15 | -------------------------------------------------------------------------------- /wsrepl/MessageHandler.py: -------------------------------------------------------------------------------- 1 | import asyncio 2 | import textual 3 | import threading 4 | import ssl 5 | from collections import OrderedDict 6 | from urllib.parse import urlparse 7 | 8 | from wsrepl.log import log 9 | from wsrepl.utils import load_plugin 10 | from wsrepl.WSMessage import WSMessage, Direction 11 | from wsrepl.WebsocketConnection import WebsocketConnection 12 | from wsrepl.PingThread import PingThread 13 | from wsrepl.Ping0x1Thread import Ping0x1Thread 14 | 15 | 16 | class MessageHandler: 17 | def __init__(self, 18 | app: textual.app.App, 19 | 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, 36 | plugin_provided_url: bool | None = None) -> None: 37 | 38 | self.app = app 39 | self.plugin = load_plugin(plugin_path)(message_handler=self) 40 | if(plugin_provided_url): 41 | try: 42 | url = self.plugin.url 43 | print("URL from plugin = " + url) 44 | except: 45 | print("Failed to get URL path from plugin. Exiting...") 46 | exit() 47 | 48 | self.initial_messages: list[WSMessage] = self._load_initial_messages(initial_msgs_file) 49 | processed_headers: OrderedDict = self._process_headers(headers, headers_file, user_agent, origin, cookies) 50 | 51 | self._ws = WebsocketConnection( 52 | # Stuff WebsocketConnection needs to call back to us 53 | async_handler=self, 54 | # WebSocketApp args 55 | url=url, 56 | header=processed_headers 57 | ) 58 | 59 | # Args passed to websocket.WebSocketApp.run_forever() 60 | if isinstance(proxy, str): 61 | proxy = 'http://' + proxy if '://' not in proxy else proxy 62 | 63 | self._ws.connect_args = { 64 | 'suppress_origin': 'Origin' in processed_headers, 65 | 'sslopt': {'cert_reqs': ssl.CERT_NONE} if not verify_tls else {}, 66 | 'ping_interval': 0, # Disable websocket-client's autoping because it doesn't provide feedback 67 | 'http_proxy_host': urlparse(proxy).hostname if proxy else None, 68 | 'http_proxy_port': urlparse(proxy).port if proxy else None, 69 | 'proxy_type': 'http' if proxy else None, 70 | 'reconnect': reconnect_interval 71 | } 72 | 73 | self.is_stopped = threading.Event() 74 | 75 | # Regular ping thread, conforming to RFC 6455 (ping uses opcode 0x9, pong uses 0xA, data is arbitrary but must be the same) 76 | if ping_interval: 77 | self.ping_thread = PingThread(self, ping_interval, self.is_stopped) 78 | else: 79 | self.ping_thread = None 80 | 81 | # Fake ping thread, using opcode 0x (TEXT) and arbitrary ping / pong messages 82 | self.ping_0x1_interval = ping_0x1_payload and pong_0x1_payload and ping_0x1_interval 83 | self.ping_0x1_data = ping_0x1_payload or self.plugin.ping_0x1_payload 84 | self.pong_0x1_data = pong_0x1_payload or self.plugin.pong_0x1_payload 85 | if self.ping_0x1_interval: 86 | self.ping_0x1_thread = Ping0x1Thread(self, ping_0x1_interval, ping_0x1_payload, self.is_stopped) 87 | else: 88 | self.ping_0x1_thread = None 89 | 90 | # Whether to show ping / pong messages in the history 91 | self.hide_ping_pong = hide_ping_pong 92 | self.hide_0x1_ping_pong = hide_0x1_ping_pong 93 | 94 | def _process_headers(self, headers: list[str] | None, headers_file: str | None, 95 | user_agent: str | None, origin: str | None, cookies: list[str]) -> OrderedDict: 96 | """Process headers and return an OrderedDict of headers.""" 97 | result = OrderedDict() 98 | cookie_headers = [] 99 | 100 | # Blacklisted headers that should be removed to avoid duplication 101 | blacklisted_headers = [ 102 | "Host", 103 | "Upgrade", 104 | "Connection" 105 | ] 106 | 107 | # Headers from command line take precedence 108 | if headers: 109 | for header in headers: 110 | name, value = map(str.strip, header.split(":", 1)) 111 | if name in blacklisted_headers: 112 | continue 113 | 114 | if name.lower().strip() == "cookie": 115 | cookie_headers.append(value.strip()) 116 | else: 117 | result[name] = value 118 | 119 | # Headers from file are next 120 | if headers_file: 121 | with open(headers_file, "r") as f: 122 | for header in f.read().splitlines(): 123 | name, value = map(str.strip, header.split(":", 1)) 124 | if name in blacklisted_headers: 125 | continue 126 | 127 | if name.lower().strip() == "cookie": 128 | cookie_headers.append(value.strip()) 129 | elif name not in result: 130 | result[name] = value 131 | 132 | # Add User-Agent if not already present 133 | if user_agent and "User-Agent" not in result: 134 | result["User-Agent"] = user_agent 135 | 136 | # Add Origin if not already present 137 | if origin and "Origin" not in result: 138 | result["Origin"] = origin 139 | 140 | # Merge and add Cookies 141 | if cookies: 142 | cookie_value = "; ".join(cookies) 143 | if cookie_headers: 144 | result['Cookie'] = cookie_value + "; " + "; ".join(cookie_headers) 145 | else: 146 | result['Cookie'] = cookie_value 147 | 148 | return result 149 | 150 | def _load_initial_messages(self, initial_msgs_file: str | None) -> list[WSMessage]: 151 | messages = [] 152 | if initial_msgs_file: 153 | with open(initial_msgs_file, "r") as f: 154 | for msg in f.readlines(): 155 | messages.append(WSMessage.outgoing(msg.strip())) 156 | return messages 157 | 158 | async def init(self, event_loop) -> None: 159 | """Start the websocket thread.""" 160 | # Run self._ws.run_forever in a separate thread 161 | self._ws.event_loop = event_loop 162 | threading.Thread(target=self._ws.run_forever, daemon=True).start() 163 | 164 | # Start the ping thread 165 | if self.ping_thread: 166 | self.ping_thread.event_loop = asyncio.get_running_loop() 167 | self.ping_thread.start() 168 | 169 | if self.ping_0x1_thread: 170 | self.ping_0x1_thread.event_loop = asyncio.get_running_loop() 171 | self.ping_0x1_thread.start() 172 | 173 | async def on_connect(self) -> None: 174 | """Called when the websocket connects.""" 175 | log.info("Websocket connected") 176 | self.is_stopped.clear() 177 | 178 | await self._send_initial_messages() 179 | await self.plugin._on_connect() 180 | 181 | # Enable the input widget and focus it 182 | self.app.enable_input() 183 | 184 | async def _send_initial_messages(self): 185 | """Send the initial messages to the server""" 186 | for message in self.initial_messages: 187 | await self.send(message) 188 | 189 | async def on_disconnect(self, status_code, reason) -> None: 190 | """Called when the websocket disconnects.""" 191 | log.error(f"Websocket disconnected with status code {status_code} and reason {reason}") 192 | 193 | self.is_stopped.set() 194 | await self.plugin.on_disconnect() 195 | 196 | async def on_error(self, exception: Exception) -> None: 197 | """Called when the websocket encounters an error.""" 198 | log.error(f"Websocket error: {exception}") 199 | await self.plugin.on_error(exception) 200 | 201 | async def on_message_received(self, msg: str) -> bool: 202 | """Called when the websocket receives a text message.""" 203 | message = WSMessage.incoming(message) 204 | 205 | await self.plugin.on_message_received(message) 206 | await self.plugin.after_message_received(message) 207 | 208 | # NOTE: This is what sends received messages to the history for 1) text; 2) binary and 3) continuation frames 209 | async def on_data_received(self, data: str | bytes, opcode: int, fin: bool) -> None: 210 | log.debug(f"Received data: {data}, opcode: {opcode}, fin: {fin}") 211 | 212 | message = WSMessage.incoming(data, opcode=opcode, fin=fin) 213 | if opcode == 0x1 and data in (self.ping_0x1_data, self.pong_0x1_data): 214 | message.is_fake_ping_pong = True 215 | if self.hide_0x1_ping_pong: 216 | message.is_hidden = True 217 | 218 | await self.plugin.on_message_received(message) 219 | if not message.is_hidden: 220 | self.app.history.add_message(message) 221 | await self.plugin.after_message_received(message) 222 | 223 | # Auto respond with a fake pong if the received message is a fake ping 224 | if self.ping_0x1_data and self.pong_0x1_data and opcode == 0x1 and data == self.ping_0x1_data: 225 | log.debug("Received fake ping - responding with fake pong") 226 | pong_message = WSMessage.outgoing(self.pong_0x1_data, opcode=0x1) 227 | pong_message.is_fake_ping_pong = True 228 | if self.hide_0x1_ping_pong: 229 | pong_message.is_hidden = True 230 | await self.send(pong_message) 231 | 232 | async def on_ping_received(self, data: str | bytes) -> None: 233 | """Called when the websocket receives a ping.""" 234 | log.debug("Received ping") 235 | message = WSMessage.ping_in(data) 236 | await self.plugin.on_ping_received(data) 237 | if not message.is_hidden and not self.hide_ping_pong: 238 | self.app.history.add_message(message) 239 | await self.plugin.after_ping_received(data) 240 | 241 | # Auto respond with a pong 242 | log.debug("Responding with pong") 243 | pong_message = WSMessage.pong_out(data) 244 | await self.send(pong_message) 245 | 246 | async def on_pong_received(self, data: str | bytes) -> None: 247 | """Called when the websocket receives a pong.""" 248 | message = WSMessage.pong_in(data) 249 | await self.plugin.on_pong_received(data) 250 | if not message.is_hidden and not self.hide_ping_pong: 251 | self.app.history.add_message(message) 252 | await self.plugin.after_pong_received(data) 253 | 254 | async def on_continuation_received(self, data: str | bytes, fin: bool) -> None: 255 | """Called when the websocket receives a continuation frame.""" 256 | message = WSMessage.incoming(data, is_continuation=True, fin=fin) 257 | await self.plugin.on_continuation_received(data) 258 | if not message.is_hidden: 259 | self.app.history.add_message(message) 260 | await self.plugin.after_continuation_received(data) 261 | 262 | async def send(self, message: WSMessage) -> None: 263 | """Send a message to the server and push it to the history.""" 264 | log.debug(f"Sending message: {message.msg}") 265 | skip_msg_if_false = await self.plugin.on_message_sent(message) 266 | if skip_msg_if_false == False: 267 | message.direction = Direction.INFO 268 | message.short = f"Message skipped: {message.short}" 269 | else: 270 | # NOTE: This should be the only place where we send messages! 271 | self._ws.send(message.msg, opcode=message.opcode.value) 272 | 273 | # Decide whether to hide the message or not 274 | if (message.is_ping or message.is_pong) and self.hide_ping_pong: 275 | message.is_hidden = True 276 | 277 | if message.is_fake_ping_pong and self.hide_0x1_ping_pong: 278 | message.is_hidden = True 279 | 280 | if not message.is_hidden: 281 | self.app.history.add_message(message) 282 | 283 | # TODO: Spawn this without await maybe? 284 | await self.plugin.after_message_sent(message) 285 | 286 | async def send_str(self, msg: str) -> None: 287 | """Send a string message to the server.""" 288 | await self.send(WSMessage.outgoing(msg)) 289 | 290 | async def log(self, msg: str) -> None: 291 | """Log a message to the history.""" 292 | log.debug(msg) 293 | -------------------------------------------------------------------------------- /wsrepl/Ping0x1Thread.py: -------------------------------------------------------------------------------- 1 | from wsrepl.WSMessage import WSMessage, Opcode 2 | from wsrepl.PingThread import PingThread 3 | import threading 4 | 5 | class Ping0x1Thread(PingThread): 6 | def __init__(self, 7 | message_handler: 'MessageHandler', 8 | ping_interval: int, 9 | ping_data: str, 10 | is_stopped: threading.Event) -> None: 11 | super().__init__(message_handler, ping_interval, is_stopped) 12 | self.data = ping_data 13 | 14 | def send_ping(self) -> None: 15 | """Send a ping packet to the server""" 16 | ping_message = WSMessage.outgoing(self.data, opcode=Opcode.TEXT) 17 | self.send(ping_message) 18 | self.update_last_ping() 19 | -------------------------------------------------------------------------------- /wsrepl/PingThread.py: -------------------------------------------------------------------------------- 1 | import asyncio 2 | import time 3 | import threading 4 | from datetime import datetime 5 | from wsrepl.WSMessage import WSMessage 6 | 7 | 8 | class PingThread(threading.Thread): 9 | def __init__(self, 10 | message_handler: 'MessageHandler', 11 | ping_interval: int | float, 12 | is_stopped: threading.Event) -> None: 13 | super().__init__(daemon=True) 14 | 15 | self.ping_interval = ping_interval 16 | self.message_handler = message_handler 17 | self.is_stopped = is_stopped 18 | self.update_last_ping() 19 | 20 | def update_last_ping(self) -> None: 21 | """Updates last ping time""" 22 | self.last_ping = datetime.now() 23 | 24 | def seconds_since_last_ping(self) -> int: 25 | """Returns number of seconds since last ping""" 26 | return (datetime.now() - self.last_ping).seconds 27 | 28 | def run(self) -> None: 29 | self.log(f"Starting ping thread with interval of {self.ping_interval} seconds") 30 | while True: 31 | seconds_passed = self.seconds_since_last_ping() 32 | self.log(f"Seconds passed since last ping: {seconds_passed}") 33 | if seconds_passed < self.ping_interval: 34 | # Wait for ping interval using time.sleep 35 | seconds_to_sleep = max(1, self.ping_interval - seconds_passed) 36 | time.sleep(seconds_to_sleep) 37 | continue 38 | 39 | # Check if thread is stopped 40 | if self.is_stopped.is_set(): 41 | self.log(f"Ping thread is stopped, sleeping for {self.ping_interval} seconds") 42 | self.update_last_ping() 43 | continue 44 | 45 | # Send ping packet 46 | self.send_ping() 47 | 48 | def log(self, msg: str) -> None: 49 | """Log a message""" 50 | asyncio.run_coroutine_threadsafe( 51 | self.message_handler.log(msg), 52 | self.event_loop) 53 | 54 | def send(self, *args, **kwargs) -> None: 55 | """Send a message to the server""" 56 | return asyncio.run_coroutine_threadsafe( 57 | self.message_handler.send(*args, **kwargs), 58 | self.event_loop) 59 | 60 | def send_ping(self) -> None: 61 | """Send a ping packet to the server""" 62 | ping_message = WSMessage.ping_out() 63 | self.send(ping_message) 64 | self.update_last_ping() 65 | -------------------------------------------------------------------------------- /wsrepl/Plugin.py: -------------------------------------------------------------------------------- 1 | import asyncio 2 | 3 | from wsrepl.log import log 4 | from wsrepl.WSMessage import WSMessage 5 | 6 | class Plugin: 7 | def __init__(self, message_handler) -> None: 8 | self.log = log 9 | self.handler = message_handler 10 | self.messages = [] 11 | self.ping_0x1_payload = "" 12 | self.pong_0x1_payload = "" 13 | self.init() 14 | 15 | def init(self): 16 | """Called when the plugin is loaded""" 17 | pass 18 | 19 | async def send(self, message: WSMessage) -> None: 20 | """Send a message to the server""" 21 | await self.handler.send(message) 22 | 23 | async def send_str(self, msg: str) -> None: 24 | """Send a string to the server""" 25 | message = WSMessage.outgoing(msg) 26 | await self.send(message) 27 | 28 | async def on_connect(self) -> None: 29 | """Called when the websocket connection is established""" 30 | pass 31 | 32 | async def _on_connect(self) -> None: 33 | """Internal method to send messages from the self.messages list and call the public on_connect hook""" 34 | await self.on_connect() 35 | if self.messages: 36 | self.log.debug(f"Sending initial messages ({len(self.messages)})") 37 | for msg in self.messages: 38 | if msg.strip() == "": 39 | # Empty message signals 1 second pause 40 | self.log.debug("Empty message, sleeping for 1 second") 41 | await asyncio.sleep(1) 42 | else: 43 | self.log.debug(f"Sending initial message: {msg}") 44 | await self.send_str(msg) 45 | await self.after_connect() 46 | 47 | async def after_connect(self) -> None: 48 | pass 49 | 50 | async def on_disconnect(self) -> None: 51 | """Called when the websocket connection is lost""" 52 | pass 53 | 54 | # Message frame 55 | async def on_message_received(self, message: WSMessage) -> None: 56 | """Called when a (text) message is received from the server. Can modify the message before it is pushed to the history. 57 | 58 | - message.msg contains the original message received from the server (should not be modified) 59 | - message.short contains a shortened version of the message that will be displayed in the history (can be overwritten) 60 | - message.long contains a long version of the message that will be displayed upon selection (can be overwritten) 61 | 62 | - message.is_hidden - if True, the message will not be displayed in the history 63 | 64 | ! Do not use this hook to send responses to the server. Use the after_message_received hook instead ! 65 | """ 66 | pass 67 | 68 | async def after_message_received(self, message: WSMessage) -> None: 69 | """Called after a message is received from the server. Use this hook to send responses to the server.""" 70 | pass 71 | 72 | # Data frame 73 | async def on_data_received(self, message: WSMessage) -> None: 74 | """Called before on_message_received is fired, and also on binary messages and continuation frames.""" 75 | pass 76 | 77 | async def after_data_received(self, message: WSMessage) -> None: 78 | """Called after on_message_received is fired, and also on binary messages and continuation frames.""" 79 | pass 80 | 81 | # Continuation frame 82 | async def on_continuation_received(self, message: WSMessage) -> None: 83 | """Called when a continuation frame is received from the server.""" 84 | pass 85 | 86 | async def after_continuation_received(self, message: WSMessage) -> None: 87 | """Called after a continuation frame is received from the server.""" 88 | pass 89 | 90 | # Ping frame (received) 91 | async def on_ping_received(self, message: WSMessage) -> None: 92 | """Called when a ping frame is received from the server.""" 93 | pass 94 | 95 | async def after_ping_received(self, message: WSMessage) -> None: 96 | """Called after a ping frame is received from the server.""" 97 | pass 98 | 99 | # Pong frame (received) 100 | async def on_pong_received(self, message: WSMessage) -> None: 101 | """Called when a pong frame is received from the server.""" 102 | pass 103 | 104 | async def after_pong_received(self, message: WSMessage) -> None: 105 | """Called after a pong frame is received from the server.""" 106 | pass 107 | 108 | # Sending messages to the server 109 | async def on_message_sent(self, message: WSMessage) -> bool | None: 110 | """Called when a message is sent to the server. Can modify the message before it is sent and pushed to history. 111 | 112 | See on_message_received for more information on the message object. 113 | 114 | Similary to on_message_received, setting message.is_hidden to True will prevent the message from being displayed in the history. 115 | 116 | Returning False from this hook will prevent the message from being sent to the server. 117 | 118 | ! Do not use this hook to send more messages to the server. Use the after_message_sent hook instead ! 119 | """ 120 | pass 121 | 122 | async def after_message_sent(self, message: WSMessage) -> None: 123 | """Called after a message is sent to the server. Use this hook to send responses to the server.""" 124 | pass 125 | 126 | async def on_error(self, exception: Exception) -> None: 127 | """Called when an error message is received from the server""" 128 | pass 129 | -------------------------------------------------------------------------------- /wsrepl/WSMessage.py: -------------------------------------------------------------------------------- 1 | import json 2 | from enum import Enum 3 | from base64 import b64decode, b64encode 4 | 5 | from rich.syntax import Syntax 6 | from rich.markup import escape 7 | 8 | from wsrepl.constants import ACCENT_COLOR 9 | 10 | 11 | class Direction(Enum): 12 | """Direction of a websocket message.""" 13 | INCOMING = '<' # Received message 14 | OUTGOING = '>' # Sent message 15 | DEBUG = '.' # Debug message 16 | INFO = '?' # Information message 17 | WARNING = '#' # Warning message 18 | ERROR = '!' # Error message 19 | 20 | class Opcode(Enum): 21 | """Opcode of a websocket message.""" 22 | CONTINUATION = 0x0 23 | TEXT = 0x1 24 | BINARY = 0x2 25 | CLOSE = 0x8 26 | PING = 0x9 27 | PONG = 0xA 28 | 29 | class WSMessage: 30 | """A websocket message. 31 | 32 | The original message is stored in `msg`, and optionally two more versions can be provided: 33 | 34 | - `short` is a shortened version of the message, that is displayed in the history widget 35 | - `long` is a formatted version of the message, that is displayed in the message details widget 36 | 37 | If these are not provided, the original message is used for `short` and JSON formatted version is used for `long`. 38 | """ 39 | _short_value = None 40 | _long_value = None 41 | is_fake_ping_pong = False 42 | 43 | def __init__(self, msg: str, direction: Direction, hidden = False, is_binary: bool = False, is_ping: bool = False, 44 | is_pong: bool = False, is_continuation: bool = False, opcode: int | Opcode = Opcode.TEXT, fin: bool = True) -> None: 45 | self.msg: str = msg 46 | self.direction: Direction = direction 47 | 48 | # Hide the message from the history 49 | self.is_hidden: bool = hidden 50 | 51 | # Set the opcode and fin flags 52 | self.opcode: Opcode = opcode if isinstance(opcode, Opcode) else Opcode(opcode) 53 | if is_binary: 54 | self.is_binary = True 55 | if is_ping: 56 | self.is_ping = True 57 | if is_pong: 58 | self.is_pong = True 59 | if is_continuation: 60 | self.is_continuation = True 61 | 62 | # FIN flag means this is the last frame in a message 63 | self.fin : bool = fin 64 | 65 | @property 66 | def is_binary(self) -> bool: 67 | return self.opcode == Opcode.BINARY 68 | 69 | @is_binary.setter 70 | def is_binary(self, value: bool) -> None: 71 | self.opcode = Opcode.BINARY if value else Opcode.TEXT 72 | 73 | @property 74 | def is_ping(self) -> bool: 75 | return self.opcode == Opcode.PING 76 | 77 | @is_ping.setter 78 | def is_ping(self, value: bool) -> None: 79 | self.opcode = Opcode.PING if value else Opcode.TEXT 80 | 81 | @property 82 | def is_pong(self) -> bool: 83 | return self.opcode == Opcode.PONG 84 | 85 | @is_pong.setter 86 | def is_pong(self, value: bool) -> None: 87 | self.opcode = Opcode.PONG if value else Opcode.TEXT 88 | 89 | @property 90 | def is_continuation(self) -> bool: 91 | return self.opcode == Opcode.CONTINUATION 92 | 93 | @is_continuation.setter 94 | def is_continuation(self, value: bool) -> None: 95 | self.opcode = Opcode.CONTINUATION if value else Opcode.TEXT 96 | 97 | @property 98 | def short(self) -> str: 99 | return escape(self._short) 100 | 101 | @short.setter 102 | def short(self, value: str) -> None: 103 | self._short_value = value 104 | 105 | @property 106 | def _short(self) -> str: 107 | """Not escaped version of the short message, for internal consumption.""" 108 | if self._short_value: 109 | return self._short_value 110 | 111 | if isinstance(self.msg, bytes): 112 | try: 113 | decoded = self.msg.decode('utf-8') 114 | return decoded 115 | except: 116 | return 'b64:' + b64encode(self.msg).decode('utf-8') 117 | elif isinstance(self.msg, str): 118 | return self.msg 119 | else: 120 | raise ValueError('Message is not a string or bytes') 121 | 122 | @property 123 | def long(self) -> str: 124 | msg = self._long 125 | 126 | # Try to format the message as JSON, if it fails, just return the original message 127 | try: 128 | parsed = json.loads(msg) 129 | pretty = json.dumps(parsed, indent=4) 130 | return Syntax(pretty, lexer="json", theme='native', background_color=ACCENT_COLOR) 131 | except: 132 | return escape(msg) 133 | 134 | @long.setter 135 | def long(self, value: str) -> None: 136 | self._long_value = value 137 | 138 | @property 139 | def _long(self) -> str: 140 | """Not escaped version of the long message, for internal consumption.""" 141 | return self._long_value if self._long_value else self._short 142 | 143 | @property 144 | def binary(self) -> bytes: 145 | """Return the binary message as bytes (check self.is_binary flag first).""" 146 | if self.is_binary: 147 | # Validate that the message is a base64 string by checking 'b64:' prefix 148 | if not self.msg.startswith('b64:'): 149 | raise ValueError('Binary message does not start with "b64:" prefix') 150 | return b64decode(self.msg[4:]) 151 | return self.msg.encode() 152 | 153 | # Convenient methods for creating messages 154 | @classmethod 155 | def incoming(cls, msg: str, *args, **kwargs) -> 'WSMessage': 156 | return cls(msg, Direction.INCOMING, *args, **kwargs) 157 | 158 | @classmethod 159 | def outgoing(cls, msg: str, *args, **kwargs) -> 'WSMessage': 160 | return cls(msg, Direction.OUTGOING, *args, **kwargs) 161 | 162 | @classmethod 163 | def ping_out(cls, msg: str = "") -> 'WSMessage': 164 | return cls(msg, Direction.OUTGOING, is_ping=True) 165 | 166 | @classmethod 167 | def ping_in(cls, msg: str = "") -> 'WSMessage': 168 | return cls(msg, Direction.INCOMING, is_ping=True) 169 | 170 | @classmethod 171 | def pong_out(cls, msg: str = "") -> 'WSMessage': 172 | return cls(msg, Direction.OUTGOING, is_pong=True) 173 | 174 | @classmethod 175 | def pong_in(cls, msg: str = "") -> 'WSMessage': 176 | return cls(msg, Direction.INCOMING, is_pong=True) 177 | 178 | @classmethod 179 | def debug(cls, msg: str) -> 'WSMessage': 180 | return cls(msg, Direction.DEBUG) 181 | 182 | @classmethod 183 | def info(cls, msg: str) -> 'WSMessage': 184 | return cls(msg, Direction.INFO) 185 | 186 | @classmethod 187 | def warning(cls, msg: str) -> 'WSMessage': 188 | return cls(msg, Direction.WARNING) 189 | 190 | @classmethod 191 | def error(cls, msg: str) -> 'WSMessage': 192 | return cls(msg, Direction.ERROR) 193 | 194 | @property 195 | def is_traffic(self) -> bool: 196 | """Return True if this message is a websocket traffic message.""" 197 | return self.direction in (Direction.INCOMING, Direction.OUTGOING) 198 | 199 | @property 200 | def opcode_hex(self) -> str: 201 | """Return the opcode as a hex string.""" 202 | if not self.is_traffic: 203 | return "-" 204 | 205 | opcode = self.opcode.value 206 | return f"0x{opcode:X}" 207 | -------------------------------------------------------------------------------- /wsrepl/WebsocketConnection.py: -------------------------------------------------------------------------------- 1 | import asyncio 2 | import websocket 3 | import time 4 | 5 | ## websocket-client does not support async, so we need to use threads. 6 | ## At the same type Textual isn't thread-safe, so we need to take care 7 | ## to only communicate with it through post_message: 8 | ## https://textual.textualize.io/guide/workers/#posting-messages 9 | 10 | 11 | class WebsocketConnection: 12 | def __init__(self, async_handler, **init_args) -> None: 13 | # Async MessageHandler that will actually handle all messages 14 | self.async_handler = async_handler 15 | # Main asyncio event loop where the async handlers runs (set during init()) 16 | self.event_loop = None 17 | 18 | self._ws = websocket.WebSocketApp( 19 | on_open=self.on_open, 20 | on_message=self.on_message, 21 | on_data=self.on_data, 22 | on_error=self.on_error, 23 | on_close=self.on_close, 24 | on_ping=self.on_ping, 25 | on_pong=self.on_pong, 26 | on_cont_message=self.on_cont_message, 27 | **init_args) 28 | self.init_args = init_args 29 | 30 | def run_forever(self) -> None: 31 | """Connect to the websocket and reconnect if disconnected.""" 32 | self.log(f"Websocket init arguments: {self.init_args}") 33 | self.log(f"run_forever arguments: {self.connect_args}") 34 | 35 | timeout = self.connect_args['reconnect'] 36 | while True: 37 | try: 38 | # Blocking call, will return when the websocket is closed 39 | # It returns False if the self._ws.close() was called or KeyboardInterrupt was raised 40 | there_were_errors = self._ws.run_forever(**self.connect_args) == False 41 | except: 42 | there_were_errors = True 43 | # TODO: Should we exit on there_were_errors or any other condition? 44 | 45 | self.log(f"Lost connection, reconnecting in {timeout} seconds") 46 | time.sleep(timeout) 47 | 48 | def _async_proxy(self, func_name, *args, **kwargs): 49 | """Run the given function in the asyncio event loop.""" 50 | func = getattr(self.async_handler, func_name) 51 | return asyncio.run_coroutine_threadsafe(func(*args, **kwargs), self.event_loop) 52 | 53 | # Blocks the websocket thread until the async handler is done, use sparingly 54 | def _async_proxy_block(self, func_name, *args, **kwargs): 55 | """Run the given function in the asyncio event loop and block until it's done.""" 56 | future = self._async_proxy(func_name, *args, **kwargs) 57 | future.result() 58 | 59 | def on_open(self, _) -> None: 60 | """Executed when the websocket is opened (also after reconnect).""" 61 | self.log("Connected to websocket") 62 | self._async_proxy('on_connect') 63 | 64 | def on_message(self, _, message: str) -> None: 65 | """Executed when a message is received. Gets fired *on text messages only*.""" 66 | self.log(f"Message: {message}") 67 | self._async_proxy('on_message_received', message) 68 | 69 | def on_data(self, _, data: str | bytes, opcode: int, fin: bool) -> None: 70 | """Executed when a message is received. Gets fired *on text, binary and continuation messages*. 71 | 72 | Execution order: 73 | - text and continuation messages: this callback will be called *before* on_message and on_cont_message 74 | - binary messages: this callback is the only one that will be called 75 | 76 | Type of data: 77 | - text messages: decoded automatically (from utf-8), so you get a string 78 | - binary messages: not decoded, so you get a bytes object 79 | 80 | The other parameters are the same as the ones of the websocket-client callback: 81 | - opcode: the websocket opcode of the message: 0x1 for text, 0x2 for binary, 0x0 for continuation 82 | - fin: True if the message is the last one of the frame, False otherwise. 83 | """ 84 | self.log(f"Data: {data}") 85 | self._async_proxy('on_data_received', data, opcode, fin) 86 | 87 | def on_error(self, _, exception: Exception) -> None: 88 | self.log(f"Error: {exception}") 89 | self._async_proxy('on_error', exception) 90 | 91 | def on_close(self, _, status_code, reason) -> None: 92 | """Executed *AFTER* the websocket is closed.""" 93 | self.log(f"Closed: {status_code} {reason}") 94 | self._async_proxy('on_disconnect', status_code, reason) 95 | 96 | def on_ping(self, _, data) -> None: 97 | self.log(f"Ping: {data}") 98 | self._async_proxy('on_ping_received', data) 99 | 100 | def on_pong(self, _, data) -> None: 101 | self.log(f"Pong: {data}") 102 | self._async_proxy('on_pong_received', data) 103 | 104 | def on_cont_message(self, _, data, fin) -> None: 105 | self.log(f"Continuation message: {data}") 106 | self._async_proxy('on_cont_message', data, fin) 107 | 108 | # NOTE: This is the only function that will be called from the main thread (that runs async UI code), everything 109 | # else works in the dedicated websocket thread, but all handlers proxy requests to the async handlers in the main thread. 110 | def send(self, data: str | bytes, opcode: int = 0x1) -> None: 111 | """Send a message to the websocket. Default opcode is 0x1 (text) and expects utf-8 string.""" 112 | self.log(f"Sending message: {data}") 113 | self._ws.send(data, opcode) 114 | 115 | def log(self, msg: str) -> None: 116 | """Log the given message.""" 117 | self._async_proxy('log', msg) 118 | -------------------------------------------------------------------------------- /wsrepl/__init__.py: -------------------------------------------------------------------------------- 1 | from wsrepl.app import WSRepl 2 | from wsrepl.cli import cli 3 | from wsrepl.WSMessage import WSMessage 4 | from wsrepl.Plugin import Plugin 5 | -------------------------------------------------------------------------------- /wsrepl/__main__.py: -------------------------------------------------------------------------------- 1 | from wsrepl.cli import cli 2 | 3 | cli() 4 | -------------------------------------------------------------------------------- /wsrepl/app.css: -------------------------------------------------------------------------------- 1 | /* The default accent color is too light which breaks syntax highlighting, this is roughly accent-darken-3 */ 2 | $accent: #005c45; 3 | $background: #005c45; 4 | $primary: #005c45; 5 | $primary-background: #005c45; 6 | $accent-darken-1: #005c45; 7 | $accent-darken-2: #005c45; 8 | $accent-darken-3: #005c45; 9 | 10 | History { 11 | scrollbar-color: $accent-darken-3; 12 | } 13 | 14 | /* By default, set line height to 3, this will be overwritten with '.small' class if --small was provided */ 15 | ListItem { 16 | min-height: 3; 17 | } 18 | 19 | /* Line counter on the left */ 20 | .history-index { 21 | /* I don't think the app will be able to handle 99999 messages anyway */ 22 | width: 5; 23 | /* Less accented */ 24 | color: $text 50%; 25 | /* Make numbers nicely line up */ 26 | text-align: right; 27 | /* 1 char padding above and below */ 28 | padding: 1 0; 29 | } 30 | 31 | /* Timestamp */ 32 | .history-time { 33 | width: 10; 34 | color: $text 50%; 35 | padding: 1; 36 | } 37 | 38 | /* Opcode */ 39 | .history-opcode { 40 | width: 3; 41 | padding: 1 0; 42 | text-align: center; 43 | } 44 | 45 | /* Indicates the type of the history entry */ 46 | .history-sign { 47 | text-align: right; 48 | width: 2; 49 | text-style: bold; 50 | padding: 1 0 0 0; 51 | } 52 | /* Need to rethink colors, so that they work with all color schemes */ 53 | .OUTGOING .history-opcode, .OUTGOING .history-sign { 54 | color: orange; 55 | } 56 | .INCOMING .history-opcode, .INCOMING .history-sign { 57 | color: green; 58 | } 59 | .DEBUG .history-opcode, .DEBUG .history-sign { 60 | color: gray; 61 | } 62 | .INFO .history-opcode, .INFO .history-sign { 63 | color: $text; 64 | } 65 | .WARNING .history-opcode, .WARNING .history-sign { 66 | color: yellow; 67 | } 68 | .ERROR .history-opcode, .ERROR .history-sign { 69 | color: red; 70 | } 71 | 72 | /* The actual history entry, in the middle, could be multiline */ 73 | .history-text { 74 | text-style: bold; 75 | padding: 1 2; 76 | } 77 | .history-text.selected { 78 | outline: round #888; 79 | } 80 | 81 | /* Copy button on the right */ 82 | .history-btn { 83 | dock: right; 84 | margin-right: 2; 85 | width: 17; 86 | outline: round $success; 87 | border: none; 88 | color: $success; 89 | background: $boost; 90 | } 91 | 92 | /* --small switch enables smaller UI */ 93 | .small ListItem { 94 | min-height: 1; 95 | } 96 | 97 | .small .history-time { 98 | padding: 0 1; 99 | } 100 | 101 | .small .history-index, .small .history-sign, .small .history-btn, .small .history-opcode { 102 | padding: 0; 103 | } 104 | 105 | .small .history-text { 106 | padding: 0 2; 107 | } 108 | .small .history-text.selected { 109 | outline: none !important; 110 | } 111 | 112 | .small .history-btn { 113 | height: 1; 114 | outline: none; 115 | border: none; 116 | } 117 | 118 | .small Input { 119 | min-height: 1; 120 | outline: none; 121 | border: none; 122 | } 123 | -------------------------------------------------------------------------------- /wsrepl/app.py: -------------------------------------------------------------------------------- 1 | import asyncio 2 | import inspect 3 | import logging 4 | 5 | from textual import events, on 6 | from textual.app import App, ComposeResult 7 | from textual.logging import TextualHandler 8 | from textual.widgets import Input 9 | 10 | from wsrepl.log import log, register_log_handler 11 | from wsrepl.widgets import History, Parent 12 | from wsrepl.MessageHandler import MessageHandler 13 | 14 | 15 | logging.basicConfig( 16 | level="DEBUG", 17 | # https://textual.textualize.io/guide/devtools/ 18 | handlers=[TextualHandler()], 19 | ) 20 | 21 | 22 | class WSRepl(App): 23 | CSS_PATH = "app.css" 24 | 25 | def __init__(self, 26 | # URL is the only required argument 27 | url: str, 28 | # UI settings 29 | small: bool = False, 30 | # 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, 45 | # Other 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: 50 | super().__init__() 51 | 52 | # Small UI 53 | self.small = small 54 | 55 | # Verbosity for logging level 56 | self.verbosity = verbosity 57 | 58 | # Message handler, spawns a thread to handle the websocket connection 59 | self.message_handler = MessageHandler( 60 | app=self, 61 | url=url, user_agent=user_agent, origin=origin, cookies=cookies, headers=headers, headers_file=headers_file, 62 | ping_interval=ping_interval, hide_ping_pong=hide_ping_pong, 63 | ping_0x1_interval=ping_0x1_interval, ping_0x1_payload=ping_0x1_payload, pong_0x1_payload=pong_0x1_payload, 64 | hide_0x1_ping_pong=hide_0x1_ping_pong, 65 | reconnect_interval=reconnect_interval, proxy=proxy, verify_tls=verify_tls, 66 | initial_msgs_file=initial_msgs_file, plugin_path=plugin_path, plugin_provided_url=plugin_provided_url 67 | ) 68 | 69 | # These are set in compose() 70 | self.history = None 71 | self.input_widget = None 72 | 73 | @on(History.Ready) 74 | async def _history_mount(self, event: events.Mount) -> None: 75 | """Called when the history widget is mounted and we're ready to connect to the websocket.""" 76 | # Set up logging, allows adding messages to UI by logging them 77 | register_log_handler(self, self.verbosity) 78 | # Pass asyncio event loop to the message handler so that it can schedule tasks on main thread 79 | await self.message_handler.init(asyncio.get_running_loop()) 80 | 81 | def compose(self) -> ComposeResult: 82 | """Compose the Textual app layout.""" 83 | self.history = History(self.small) 84 | self.input_widget = Input(placeholder="Enter websocket message", disabled=True) 85 | 86 | classes = ["app"] 87 | if self.small: 88 | classes.append("small") 89 | 90 | yield Parent(classes=classes) 91 | 92 | async def on_input_submitted(self, event) -> None: 93 | await self.message_handler.send_str(event.value) 94 | self.input_widget.value = '' 95 | 96 | def disable_input(self) -> None: 97 | """Disable the input widget.""" 98 | self.input_widget.disabled = True 99 | 100 | def enable_input(self) -> None: 101 | """Enable the input widget.""" 102 | self.input_widget.disabled = False 103 | self.input_widget.focus() 104 | 105 | -------------------------------------------------------------------------------- /wsrepl/cli.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | from wsrepl import WSRepl 4 | 5 | parser = argparse.ArgumentParser(description='Websocket Client') 6 | # Pass URL either as -u or as a positional argument 7 | parser.add_argument('-u', '--url', type=str, help='Websocket URL (e.g. wss://echo.websocket.org)') 8 | parser.add_argument('url_positional', nargs='?', type=str, help='Websocket URL (e.g. wss://echo.websocket.org)') 9 | 10 | # curl compatible args, this is the stuff that Burp uses in a 'Copy as curl command' menu action 11 | parser.add_argument('-i', '--include', action='store_true', default=False, help='No effect, just for curl compatibility') 12 | parser.add_argument('-s', '--silent', action='store_true', default=False, help='No effect, just for curl compatibility') 13 | parser.add_argument('-k', '--insecure', action='store_true', default=False, help='Disable SSL verification') 14 | parser.add_argument('-X', '--request', type=str, help='No effect, just for curl compatibility') 15 | parser.add_argument('-H', '--header', action='append', help='Additional header (e.g. "X-Header: value"), can be used multiple times') 16 | parser.add_argument('-b', '--cookie', action='append', help='Cookie header (e.g. "name=value"), can be used multiple times') 17 | # curl compatible args, used by Chrome's 'Copy as cURL' menu action 18 | parser.add_argument('--compressed', action='store_true', default=False, help='No effect, just for curl compatibility') 19 | 20 | parser.add_argument('-S', '--small', action='store_true', default=False, help='Smaller UI') 21 | parser.add_argument('-A', '--user-agent', type=str, help='User-Agent header') 22 | parser.add_argument('-O', '--origin', type=str, help='Origin header') 23 | parser.add_argument('-F', '--headers-file', type=str, help='Additional headers file (e.g. "headers.txt")') 24 | parser.add_argument( '--no-native-ping', action='store_true', default=False, help='Disable native ping/pong messages') 25 | parser.add_argument( '--ping-interval', type=int, default=24, help='Ping interval (seconds)') 26 | parser.add_argument( '--hide-ping-pong', action='store_true', default=False, help='Hide ping/pong messages') 27 | parser.add_argument( '--ping-0x1-interval', type=int, default=24, help='Fake ping (0x1 opcode) interval (seconds)') 28 | parser.add_argument( '--ping-0x1-payload', type=str, help='Fake ping (0x1 opcode) payload') 29 | parser.add_argument( '--pong-0x1-payload', type=str, help='Fake pong (0x1 opcode) payload') 30 | parser.add_argument( '--hide-0x1-ping-pong', action='store_true', default=False, help='Hide fake ping/pong messages') 31 | parser.add_argument('-t', '--ttl', type=int, help='Heartbeet interval (seconds)') 32 | parser.add_argument('-p', '--http-proxy', type=str, help='HTTP Proxy Address (e.g. 127.0.0.1:8080)') 33 | parser.add_argument('-r', '--reconnect-interval', type=int, default=2, help='Reconnect interval (seconds, default: 2)') 34 | parser.add_argument('-I', '--initial-messages', type=str, help='Send the messages from this file on connect') 35 | 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='Indicates if plugin provided dynamic url for websockets') 37 | parser.add_argument('-v', '--verbose', type=int, default=3, help='Verbosity level, 1-4 default: 3 (errors, warnings, info), 4 adds debug') 38 | 39 | def cli(): 40 | args = parser.parse_args() 41 | url = args.url or args.url_positional 42 | if url and args.plugin_provided_url is False: 43 | # Check and modify the URL protocol if necessary 44 | if url.startswith('http://'): 45 | url = url.replace('http://', 'ws://', 1) 46 | elif url.startswith('https://'): 47 | url = 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 is False and args.plugin is not None: 51 | parser.error('Please provide a WebSocket URL using -u or use --plugin-provided-url if the WebSocket URL provided in a plugin') 52 | elif args.plugin_provided_url is False and args.plugin is None: 53 | parser.error('Please provide either a WebSocket URL using -u or use --plugin-provided-url with --plugin if the WebSocket URL provided in a plugin') 54 | 55 | app = WSRepl( 56 | url=url, 57 | small=args.small, 58 | user_agent=args.user_agent, 59 | origin=args.origin, 60 | cookies=args.cookie if isinstance(args.cookie, list) else [args.cookie] if args.cookie else [], 61 | headers=args.header if isinstance(args.header, list) else [args.header] if args.header else [], 62 | headers_file=args.headers_file, 63 | ping_interval=args.ping_interval, 64 | hide_ping_pong=args.hide_ping_pong, 65 | ping_0x1_interval=args.ping_0x1_interval, 66 | ping_0x1_payload=args.ping_0x1_payload, 67 | pong_0x1_payload=args.pong_0x1_payload, 68 | hide_0x1_ping_pong=args.hide_0x1_ping_pong, 69 | reconnect_interval=args.reconnect_interval, 70 | proxy=args.http_proxy, 71 | verify_tls=not args.insecure, 72 | initial_msgs_file=args.initial_messages, 73 | plugin_path=args.plugin, 74 | plugin_provided_url=args.plugin_provided_url, 75 | verbosity=args.verbose, 76 | ) 77 | app.run() 78 | 79 | if __name__ == '__main__': 80 | cli() 81 | -------------------------------------------------------------------------------- /wsrepl/constants.py: -------------------------------------------------------------------------------- 1 | ACCENT_COLOR = "#005c45" 2 | -------------------------------------------------------------------------------- /wsrepl/log.py: -------------------------------------------------------------------------------- 1 | import logging 2 | 3 | from wsrepl.WSMessage import WSMessage 4 | 5 | class WSReplHandler(logging.Handler): 6 | def __init__(self, app, level): 7 | super().__init__() 8 | self.history = app.history 9 | self.setLevel(level) 10 | 11 | def emit(self, record): 12 | msg = record.getMessage() 13 | 14 | if record.levelname == 'DEBUG': 15 | message = WSMessage.debug(msg) 16 | elif record.levelname == 'INFO': 17 | message = WSMessage.info(msg) 18 | elif record.levelname == 'WARNING': 19 | message = WSMessage.warning(msg) 20 | elif record.levelname == 'ERROR': 21 | message = WSMessage.error(msg) 22 | else: 23 | raise ValueError(f'Unknown log level: {record.levelname}') 24 | 25 | self.history.add_message(message) 26 | 27 | log = logging.getLogger('wsrepl') 28 | 29 | def register_log_handler(app, verbosity: int): 30 | log_level = (5 - verbosity) * 10 31 | handler = WSReplHandler(app, log_level) 32 | log.addHandler(handler) 33 | -------------------------------------------------------------------------------- /wsrepl/utils.py: -------------------------------------------------------------------------------- 1 | import importlib.util 2 | import os 3 | import os.path 4 | import inspect 5 | from urllib.parse import urlparse 6 | 7 | from wsrepl.Plugin import Plugin 8 | 9 | 10 | # Sanitize url and construct likely Origin, return a tuple of (url, origin) 11 | def sanitize_url(url: str, origin: str = None) -> tuple[str, str]: 12 | parsed_url = urlparse(url) 13 | scheme = parsed_url.scheme 14 | 15 | if scheme == "http": 16 | parsed_url.scheme = "ws" 17 | elif scheme == "https": 18 | parsed_url.scheme = "wss" 19 | elif not scheme: 20 | # No scheme specified, assume wss 21 | parsed_url.scheme = "wss" 22 | url = parsed_url.geturl() 23 | 24 | # Set Origin header to match wss url, unless user specified a custom one or opted out 25 | if origin is None: 26 | if scheme == "wss": 27 | origin = "https://" + parsed_url.hostname 28 | else: 29 | origin = "http://" + parsed_url.hostname 30 | elif not origin: 31 | # False or empty string means no origin header 32 | origin = None 33 | 34 | return (url, origin) 35 | 36 | 37 | def build_ws_args(url: str, origin: str | None = None, user_agent: str | None = None, 38 | extra_headers: list[str] | None = None, extra_headers_file: str | None = None, 39 | native_ping_enabled: bool = True, heartbeet_interval: int = 24, proxy: str | None = None, 40 | skip_verify: bool = False): 41 | 42 | # Convert http / https into ws / wss if needed 43 | url, origin = sanitize_url(url, origin) 44 | 45 | args= { 46 | "url": url, 47 | "origin": origin 48 | } 49 | 50 | # --headers-file takes precedence over -H, --header 51 | if extra_headers_file: 52 | with open(extra_headers_file, "r") as f: 53 | extra_headers = [line.strip() for line in f.read().splitlines()] 54 | 55 | # Headers are provided as a list of strings, each string is a header in the format "Header-Name: Header-Value" 56 | if extra_headers: 57 | headers = {} 58 | for header in extra_headers: 59 | name, value = header.split(":", 1) 60 | headers[name.strip()] = value.strip() 61 | args["headers"] = headers 62 | 63 | # Set User-Agent header if specified 64 | if user_agent: 65 | headers = headers if headers else {} 66 | headers["User-Agent"] = user_agent 67 | 68 | # Disable native ping/pong 69 | if not native_ping_enabled: 70 | args["autoping"] = False 71 | 72 | # How often to initiate ping/pong 73 | if heartbeet_interval: 74 | args["heartbeat"] = heartbeet_interval 75 | 76 | # HTTP proxy 77 | if proxy: 78 | args["proxy"] = proxy 79 | 80 | # Disable SSL verification 81 | if skip_verify: 82 | args["ssl"] = False 83 | 84 | return args 85 | 86 | 87 | def _get_plugin_name(plugin_path): 88 | # Extracts the module name from the plugin file path 89 | plugin_file = os.path.basename(plugin_path) 90 | plugin_name = os.path.splitext(plugin_file)[0] 91 | return plugin_name 92 | 93 | 94 | def load_plugin(plugin_path) -> type[Plugin]: 95 | """Loads a plugin from a file path or returns an empty plugin if no path is specified""" 96 | if not plugin_path: 97 | return Plugin 98 | 99 | if not os.path.isfile(plugin_path): 100 | raise Exception("Plugin not found: {}".format(plugin_path)) 101 | 102 | plugin_name = _get_plugin_name(plugin_path) 103 | spec = importlib.util.spec_from_file_location(plugin_name, plugin_path) 104 | module = importlib.util.module_from_spec(spec) 105 | spec.loader.exec_module(module) 106 | 107 | # Find a subclass of Plugin in the module 108 | for name, obj in inspect.getmembers(module): 109 | if inspect.isclass(obj) and issubclass(obj, Plugin) and obj is not Plugin: 110 | # Instantiate it and return 111 | return obj 112 | 113 | raise ValueError(f"No subclass of Plugin found in {plugin_path}") 114 | -------------------------------------------------------------------------------- /wsrepl/widgets/CopyButton.py: -------------------------------------------------------------------------------- 1 | import pyperclip 2 | 3 | from textual.widgets import Button 4 | 5 | from wsrepl.WSMessage import WSMessage 6 | 7 | 8 | class CopyButton(Button): 9 | """A button that copies its data to the clipboard when pressed""" 10 | def __init__(self, message: WSMessage, small) -> None: 11 | if small: 12 | name = "[Click to copy]" 13 | else: 14 | name = "Click to copy" 15 | 16 | super().__init__(name, classes="history-btn") 17 | self.message = message 18 | 19 | def on_button_pressed(self, event: Button.Pressed) -> None: 20 | pyperclip.copy(self.message.msg) 21 | self.blur() 22 | -------------------------------------------------------------------------------- /wsrepl/widgets/DirectionSign.py: -------------------------------------------------------------------------------- 1 | from textual.widgets import Label 2 | 3 | from wsrepl.WSMessage import Direction 4 | 5 | 6 | class DirectionSign(Label): 7 | """A label that displays a direction sign""" 8 | def __init__(self, sign: Direction) -> None: 9 | super().__init__(sign.value, classes=f"history-sign {sign.name}") 10 | -------------------------------------------------------------------------------- /wsrepl/widgets/FormattedMessage.py: -------------------------------------------------------------------------------- 1 | from textual.widgets import Label 2 | 3 | from wsrepl.WSMessage import WSMessage 4 | 5 | 6 | class FormattedMessage(Label): 7 | """A label that formats its data and displays it""" 8 | 9 | def __init__(self, message: WSMessage, message_width: int = 20) -> None: 10 | self.message = message 11 | super().__init__(message.short, classes="history-text") 12 | 13 | # History widget has a handler that gets triggered on resize, iterates over all messages and fixes their width 14 | self.styles.width = message_width 15 | 16 | def show_short(self) -> str: 17 | self.remove_class('selected') 18 | return self.update(self.message.short) 19 | 20 | def show_long(self) -> str: 21 | self.add_class('selected') 22 | return self.update(self.message.long) 23 | 24 | -------------------------------------------------------------------------------- /wsrepl/widgets/History.py: -------------------------------------------------------------------------------- 1 | import shutil 2 | 3 | import textual 4 | from textual import events 5 | from textual import on 6 | from textual.widgets import ListView 7 | from textual.message import Message 8 | 9 | from wsrepl.WSMessage import WSMessage, Direction 10 | from wsrepl.widgets import HistoryRow 11 | 12 | 13 | class History(ListView): 14 | message_width = 20 15 | last_highlighted = None 16 | 17 | def __init__(self, small: bool = False) -> None: 18 | self.small = small 19 | super().__init__() 20 | 21 | self.id = "history" 22 | 23 | def on_list_view_highlighted(self, event) -> None: 24 | 25 | if self.last_highlighted is not None: 26 | # Update the last highlighted message with its original text before the update 27 | self.last_highlighted.show_short() 28 | 29 | # Update the new highlighted message with a "changed" text 30 | if event.item is None: 31 | return 32 | 33 | self.last_highlighted = event.item.text 34 | self.last_highlighted.show_long() 35 | 36 | class Ready(Message): 37 | """Sent when the history is ready to receive messages""" 38 | pass 39 | 40 | def on_mount(self) -> None: 41 | self.post_message(self.Ready()) 42 | 43 | def add_message(self, message: WSMessage) -> None: 44 | assert(isinstance(message, WSMessage)) 45 | 46 | on_last_element = (self.index is None) or (self.index == len(self.children) - 1) 47 | self.append(HistoryRow(message, self.message_width, self.small)) 48 | 49 | if on_last_element: 50 | # Scroll to the bottom if the user is already there 51 | self.index = len(self.children) - 1 52 | self.scroll_end(animate=False) 53 | 54 | def on_key(self, event: events.Key) -> None: 55 | """Vim style keybindings for scrolling through the history""" 56 | if event.key == 'j': 57 | self.action_cursor_down() 58 | if event.key == 'k': 59 | self.action_cursor_up() 60 | if event.key == 'g': 61 | self.index = 0 62 | if event.key == 'G': 63 | self.index = len(self.children) - 1 64 | 65 | @on(events.Resize) 66 | def handle_resize(self, event: events.Resize) -> None: 67 | """Handle terminal resize events by adjusting the width of the history message column""" 68 | self.message_width = self.calculate_message_width() 69 | 70 | # FIXME: This is ugly af, but I don't know how to update CSS styles dynamically 71 | for child in self.children: 72 | child.text.styles.width = self.message_width 73 | 74 | def calculate_message_width(self): 75 | # Get the new terminal width 76 | terminal_width = shutil.get_terminal_size().columns 77 | 78 | # Calculate the new width of the history message column 79 | return (terminal_width 80 | - 5 # The index column 81 | - 10 # The time column 82 | - 3 # The opcode column 83 | - 3 # The direction sign column 84 | - 17 # The copy button column 85 | - 2 # Padding on the right of the copy button 86 | - 1 # An additional padding for some breathing space 87 | ) 88 | 89 | -------------------------------------------------------------------------------- /wsrepl/widgets/HistoryIndex.py: -------------------------------------------------------------------------------- 1 | from textual.widgets import Label 2 | 3 | 4 | class HistoryIndex(Label): 5 | """A label that counts how many instances of it have been created""" 6 | __internal_counter = 0 7 | 8 | def __init__(self): 9 | HistoryIndex.__internal_counter += 1 10 | super().__init__(str(self.__internal_counter), classes="history-index") 11 | 12 | -------------------------------------------------------------------------------- /wsrepl/widgets/HistoryRow.py: -------------------------------------------------------------------------------- 1 | from datetime import datetime 2 | 3 | from textual.containers import Horizontal 4 | from textual.widgets import ListItem, Label 5 | from textual.app import ComposeResult 6 | 7 | from wsrepl.WSMessage import WSMessage 8 | from wsrepl.widgets import CopyButton, DirectionSign, FormattedMessage, HistoryIndex 9 | 10 | 11 | class HistoryRow(ListItem): 12 | """A row in the history widget""" 13 | 14 | def __init__(self, message: WSMessage, message_width: int = 20, small: bool = False) -> None: 15 | # Just an auto-incremental counter 16 | self.index = HistoryIndex() 17 | 18 | # Label with current time 19 | current_time = datetime.now().strftime("%H:%M:%S") 20 | self.time = Label(current_time, classes="history-time") 21 | 22 | # Label with the opcode 23 | self.opcode = Label(message.opcode_hex, classes="history-opcode") 24 | 25 | # The direction sign (> and < for sent and received messages, respectively, ! for errors) 26 | self.sign = DirectionSign(message.direction) 27 | 28 | # The message itself, pretty printed and syntax highlighted 29 | self.text = FormattedMessage(message, message_width) 30 | 31 | # A button that copies the message to the clipboard when pressed 32 | self.button = CopyButton(message, small) 33 | 34 | self.row = Horizontal(self.index, self.time, self.opcode, self.sign, self.text, self.button, 35 | classes=message.direction.name) 36 | 37 | super().__init__(self.row) 38 | 39 | def compose(self) -> ComposeResult: 40 | yield self.row 41 | -------------------------------------------------------------------------------- /wsrepl/widgets/Parent.py: -------------------------------------------------------------------------------- 1 | from textual.widget import Widget 2 | from textual.app import ComposeResult 3 | 4 | class Parent(Widget): 5 | """A parent class for all app widgets, to make them easier to style via CSS""" 6 | def __init__(self, classes: list[str]): 7 | class_string = " ".join(classes) 8 | super().__init__(classes=class_string) 9 | 10 | def compose(self) -> ComposeResult: 11 | yield self.app.history 12 | yield self.app.input_widget 13 | -------------------------------------------------------------------------------- /wsrepl/widgets/__init__.py: -------------------------------------------------------------------------------- 1 | from .CopyButton import CopyButton 2 | from .DirectionSign import DirectionSign 3 | from .FormattedMessage import FormattedMessage 4 | from .HistoryIndex import HistoryIndex 5 | from .HistoryRow import HistoryRow 6 | from .History import History 7 | from .Parent import Parent 8 | --------------------------------------------------------------------------------