|
|
|
|
|
|
|
|
""" |
|
|
wsdump.py |
|
|
websocket - WebSocket client library for Python |
|
|
|
|
|
Copyright 2024 engn33r |
|
|
|
|
|
Licensed under the Apache License, Version 2.0 (the "License"); |
|
|
you may not use this file except in compliance with the License. |
|
|
You may obtain a copy of the License at |
|
|
|
|
|
http://www.apache.org/licenses/LICENSE-2.0 |
|
|
|
|
|
Unless required by applicable law or agreed to in writing, software |
|
|
distributed under the License is distributed on an "AS IS" BASIS, |
|
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
|
|
See the License for the specific language governing permissions and |
|
|
limitations under the License. |
|
|
""" |
|
|
|
|
|
import argparse |
|
|
import code |
|
|
import gzip |
|
|
import ssl |
|
|
import sys |
|
|
import threading |
|
|
import time |
|
|
import zlib |
|
|
from urllib.parse import urlparse |
|
|
|
|
|
import websocket |
|
|
|
|
|
try: |
|
|
import readline |
|
|
except ImportError: |
|
|
pass |
|
|
|
|
|
|
|
|
def get_encoding() -> str: |
|
|
encoding = getattr(sys.stdin, "encoding", "") |
|
|
if not encoding: |
|
|
return "utf-8" |
|
|
else: |
|
|
return encoding.lower() |
|
|
|
|
|
|
|
|
OPCODE_DATA = (websocket.ABNF.OPCODE_TEXT, websocket.ABNF.OPCODE_BINARY) |
|
|
ENCODING = get_encoding() |
|
|
|
|
|
|
|
|
class VAction(argparse.Action): |
|
|
def __call__( |
|
|
self, |
|
|
parser: argparse.Namespace, |
|
|
args: tuple, |
|
|
values: str, |
|
|
option_string: str = None, |
|
|
) -> None: |
|
|
if values is None: |
|
|
values = "1" |
|
|
try: |
|
|
values = int(values) |
|
|
except ValueError: |
|
|
values = values.count("v") + 1 |
|
|
setattr(args, self.dest, values) |
|
|
|
|
|
|
|
|
def parse_args() -> argparse.Namespace: |
|
|
parser = argparse.ArgumentParser(description="WebSocket Simple Dump Tool") |
|
|
parser.add_argument( |
|
|
"url", metavar="ws_url", help="websocket url. ex. ws://echo.websocket.events/" |
|
|
) |
|
|
parser.add_argument("-p", "--proxy", help="proxy url. ex. http://127.0.0.1:8080") |
|
|
parser.add_argument( |
|
|
"-v", |
|
|
"--verbose", |
|
|
default=0, |
|
|
nargs="?", |
|
|
action=VAction, |
|
|
dest="verbose", |
|
|
help="set verbose mode. If set to 1, show opcode. " |
|
|
"If set to 2, enable to trace websocket module", |
|
|
) |
|
|
parser.add_argument( |
|
|
"-n", "--nocert", action="store_true", help="Ignore invalid SSL cert" |
|
|
) |
|
|
parser.add_argument("-r", "--raw", action="store_true", help="raw output") |
|
|
parser.add_argument("-s", "--subprotocols", nargs="*", help="Set subprotocols") |
|
|
parser.add_argument("-o", "--origin", help="Set origin") |
|
|
parser.add_argument( |
|
|
"--eof-wait", |
|
|
default=0, |
|
|
type=int, |
|
|
help="wait time(second) after 'EOF' received.", |
|
|
) |
|
|
parser.add_argument("-t", "--text", help="Send initial text") |
|
|
parser.add_argument( |
|
|
"--timings", action="store_true", help="Print timings in seconds" |
|
|
) |
|
|
parser.add_argument("--headers", help="Set custom headers. Use ',' as separator") |
|
|
|
|
|
return parser.parse_args() |
|
|
|
|
|
|
|
|
class RawInput: |
|
|
def raw_input(self, prompt: str = "") -> str: |
|
|
line = input(prompt) |
|
|
|
|
|
if ENCODING and ENCODING != "utf-8" and not isinstance(line, str): |
|
|
line = line.decode(ENCODING).encode("utf-8") |
|
|
elif isinstance(line, str): |
|
|
line = line.encode("utf-8") |
|
|
|
|
|
return line |
|
|
|
|
|
|
|
|
class InteractiveConsole(RawInput, code.InteractiveConsole): |
|
|
def write(self, data: str) -> None: |
|
|
sys.stdout.write("\033[2K\033[E") |
|
|
|
|
|
sys.stdout.write("\033[34m< " + data + "\033[39m") |
|
|
sys.stdout.write("\n> ") |
|
|
sys.stdout.flush() |
|
|
|
|
|
def read(self) -> str: |
|
|
return self.raw_input("> ") |
|
|
|
|
|
|
|
|
class NonInteractive(RawInput): |
|
|
def write(self, data: str) -> None: |
|
|
sys.stdout.write(data) |
|
|
sys.stdout.write("\n") |
|
|
sys.stdout.flush() |
|
|
|
|
|
def read(self) -> str: |
|
|
return self.raw_input("") |
|
|
|
|
|
|
|
|
def main() -> None: |
|
|
start_time = time.time() |
|
|
args = parse_args() |
|
|
if args.verbose > 1: |
|
|
websocket.enableTrace(True) |
|
|
options = {} |
|
|
if args.proxy: |
|
|
p = urlparse(args.proxy) |
|
|
options["http_proxy_host"] = p.hostname |
|
|
options["http_proxy_port"] = p.port |
|
|
if args.origin: |
|
|
options["origin"] = args.origin |
|
|
if args.subprotocols: |
|
|
options["subprotocols"] = args.subprotocols |
|
|
opts = {} |
|
|
if args.nocert: |
|
|
opts = {"cert_reqs": ssl.CERT_NONE, "check_hostname": False} |
|
|
if args.headers: |
|
|
options["header"] = list(map(str.strip, args.headers.split(","))) |
|
|
ws = websocket.create_connection(args.url, sslopt=opts, **options) |
|
|
if args.raw: |
|
|
console = NonInteractive() |
|
|
else: |
|
|
console = InteractiveConsole() |
|
|
print("Press Ctrl+C to quit") |
|
|
|
|
|
def recv() -> tuple: |
|
|
try: |
|
|
frame = ws.recv_frame() |
|
|
except websocket.WebSocketException: |
|
|
return websocket.ABNF.OPCODE_CLOSE, "" |
|
|
if not frame: |
|
|
raise websocket.WebSocketException(f"Not a valid frame {frame}") |
|
|
elif frame.opcode in OPCODE_DATA: |
|
|
return frame.opcode, frame.data |
|
|
elif frame.opcode == websocket.ABNF.OPCODE_CLOSE: |
|
|
ws.send_close() |
|
|
return frame.opcode, "" |
|
|
elif frame.opcode == websocket.ABNF.OPCODE_PING: |
|
|
ws.pong(frame.data) |
|
|
return frame.opcode, frame.data |
|
|
|
|
|
return frame.opcode, frame.data |
|
|
|
|
|
def recv_ws() -> None: |
|
|
while True: |
|
|
opcode, data = recv() |
|
|
msg = None |
|
|
if opcode == websocket.ABNF.OPCODE_TEXT and isinstance(data, bytes): |
|
|
data = str(data, "utf-8") |
|
|
if ( |
|
|
isinstance(data, bytes) and len(data) > 2 and data[:2] == b"\037\213" |
|
|
): |
|
|
try: |
|
|
data = "[gzip] " + str(gzip.decompress(data), "utf-8") |
|
|
except: |
|
|
pass |
|
|
elif isinstance(data, bytes): |
|
|
try: |
|
|
data = "[zlib] " + str( |
|
|
zlib.decompress(data, -zlib.MAX_WBITS), "utf-8" |
|
|
) |
|
|
except: |
|
|
pass |
|
|
|
|
|
if isinstance(data, bytes): |
|
|
data = repr(data) |
|
|
|
|
|
if args.verbose: |
|
|
msg = f"{websocket.ABNF.OPCODE_MAP.get(opcode)}: {data}" |
|
|
else: |
|
|
msg = data |
|
|
|
|
|
if msg is not None: |
|
|
if args.timings: |
|
|
console.write(f"{time.time() - start_time}: {msg}") |
|
|
else: |
|
|
console.write(msg) |
|
|
|
|
|
if opcode == websocket.ABNF.OPCODE_CLOSE: |
|
|
break |
|
|
|
|
|
thread = threading.Thread(target=recv_ws) |
|
|
thread.daemon = True |
|
|
thread.start() |
|
|
|
|
|
if args.text: |
|
|
ws.send(args.text) |
|
|
|
|
|
while True: |
|
|
try: |
|
|
message = console.read() |
|
|
ws.send(message) |
|
|
except KeyboardInterrupt: |
|
|
return |
|
|
except EOFError: |
|
|
time.sleep(args.eof_wait) |
|
|
return |
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
try: |
|
|
main() |
|
|
except Exception as e: |
|
|
print(e) |
|
|
|
