codekingpro/portable-devtools

1#!/usr/bin/env python32 3"""4_wsdump.py5websocket - WebSocket client library for Python6 7Copyright 2025 engn33r8 9Licensed under the Apache License, Version 2.0 (the "License");10you may not use this file except in compliance with the License.11You may obtain a copy of the License at12 13    http://www.apache.org/licenses/LICENSE-2.014 15Unless required by applicable law or agreed to in writing, software16distributed under the License is distributed on an "AS IS" BASIS,17WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.18See the License for the specific language governing permissions and19limitations under the License.20"""21 22import argparse23import code24import gzip25import ssl26import sys27import threading28import time29import zlib30from urllib.parse import urlparse31 32import websocket33 34try:35    import readline  # noqa: F40136except ImportError:37    pass38 39 40def get_encoding() -> str:41    encoding = getattr(sys.stdin, "encoding", "")42    if not encoding:43        return "utf-8"44    else:45        return encoding.lower()46 47 48OPCODE_DATA = (websocket.ABNF.OPCODE_TEXT, websocket.ABNF.OPCODE_BINARY)49ENCODING = get_encoding()50 51 52class VAction(argparse.Action):53    def __call__(54        self,55        parser: argparse.Namespace,56        args: tuple,57        values: str,58        option_string: str = None,59    ) -> None:60        if values is None:61            values = "1"62        try:63            values = int(values)64        except ValueError:65            values = values.count("v") + 166        setattr(args, self.dest, values)67 68 69def parse_args() -> argparse.Namespace:70    parser = argparse.ArgumentParser(description="WebSocket Simple Dump Tool")71    parser.add_argument(72        "url", metavar="ws_url", help="websocket url. ex. ws://echo.websocket.events/"73    )74    parser.add_argument("-p", "--proxy", help="proxy url. ex. http://127.0.0.1:8080")75    parser.add_argument(76        "-v",77        "--verbose",78        default=0,79        nargs="?",80        action=VAction,81        dest="verbose",82        help="set verbose mode. If set to 1, show opcode. "83        "If set to 2, enable to trace  websocket module",84    )85    parser.add_argument(86        "-n", "--nocert", action="store_true", help="Ignore invalid SSL cert"87    )88    parser.add_argument("-r", "--raw", action="store_true", help="raw output")89    parser.add_argument("-s", "--subprotocols", nargs="*", help="Set subprotocols")90    parser.add_argument("-o", "--origin", help="Set origin")91    parser.add_argument(92        "--eof-wait",93        default=0,94        type=int,95        help="wait time(second) after 'EOF' received.",96    )97    parser.add_argument("-t", "--text", help="Send initial text")98    parser.add_argument(99        "--timings", action="store_true", help="Print timings in seconds"100    )101    parser.add_argument("--headers", help="Set custom headers. Use ',' as separator")102 103    return parser.parse_args()104 105 106class RawInput:107    def raw_input(self, prompt: str = "") -> str:108        line = input(prompt)109 110        if ENCODING and ENCODING != "utf-8" and not isinstance(line, str):111            line = line.decode(ENCODING).encode("utf-8")112        elif isinstance(line, str):113            line = line.encode("utf-8")114 115        return line116 117 118class InteractiveConsole(RawInput, code.InteractiveConsole):119    def write(self, data: str) -> None:120        sys.stdout.write("\033[2K\033[E")121        # sys.stdout.write("\n")122        sys.stdout.write("\033[34m< " + data + "\033[39m")123        sys.stdout.write("\n> ")124        sys.stdout.flush()125 126    def read(self) -> str:127        return self.raw_input("> ")128 129 130class NonInteractive(RawInput):131    def write(self, data: str) -> None:132        sys.stdout.write(data)133        sys.stdout.write("\n")134        sys.stdout.flush()135 136    def read(self) -> str:137        return self.raw_input("")138 139 140def main() -> None:141    start_time = time.time()142    args = parse_args()143    if args.verbose > 1:144        websocket.enableTrace(True)145    options = {}146    if args.proxy:147        p = urlparse(args.proxy)148        options["http_proxy_host"] = p.hostname149        options["http_proxy_port"] = p.port150    if args.origin:151        options["origin"] = args.origin152    if args.subprotocols:153        options["subprotocols"] = args.subprotocols154    opts = {}155    if args.nocert:156        opts = {"cert_reqs": ssl.CERT_NONE, "check_hostname": False}157    if args.headers:158        options["header"] = list(map(str.strip, args.headers.split(",")))159    ws = websocket.create_connection(args.url, sslopt=opts, **options)160    if args.raw:161        console = NonInteractive()162    else:163        console = InteractiveConsole()164        print("Press Ctrl+C to quit")165 166    def recv() -> tuple:167        try:168            frame = ws.recv_frame()169        except websocket.WebSocketException:170            return websocket.ABNF.OPCODE_CLOSE, ""171        if not frame:172            raise websocket.WebSocketException(f"Not a valid frame {frame}")173        elif frame.opcode in OPCODE_DATA:174            return frame.opcode, frame.data175        elif frame.opcode == websocket.ABNF.OPCODE_CLOSE:176            ws.send_close()177            return frame.opcode, ""178        elif frame.opcode == websocket.ABNF.OPCODE_PING:179            ws.pong(frame.data)180            return frame.opcode, frame.data181 182        return frame.opcode, frame.data183 184    def recv_ws() -> None:185        while True:186            opcode, data = recv()187            msg = None188            if opcode == websocket.ABNF.OPCODE_TEXT and isinstance(data, bytes):189                data = str(data, "utf-8")190            if (191                isinstance(data, bytes) and len(data) > 2 and data[:2] == b"\037\213"192            ):  # gzip magick193                try:194                    data = "[gzip] " + str(gzip.decompress(data), "utf-8")195                except:196                    pass197            elif isinstance(data, bytes):198                try:199                    data = "[zlib] " + str(200                        zlib.decompress(data, -zlib.MAX_WBITS), "utf-8"201                    )202                except:203                    pass204 205            if isinstance(data, bytes):206                data = repr(data)207 208            if args.verbose:209                msg = f"{websocket.ABNF.OPCODE_MAP.get(opcode)}: {data}"210            else:211                msg = data212 213            if msg is not None:214                if args.timings:215                    console.write(f"{time.time() - start_time}: {msg}")216                else:217                    console.write(msg)218 219            if opcode == websocket.ABNF.OPCODE_CLOSE:220                break221 222    thread = threading.Thread(target=recv_ws)223    thread.daemon = True224    thread.start()225 226    if args.text:227        ws.send(args.text)228 229    while True:230        try:231            message = console.read()232            ws.send(message)233        except KeyboardInterrupt:234            return235        except EOFError:236            time.sleep(args.eof_wait)237            return238 239 240if __name__ == "__main__":241    try:242        main()243    except Exception as e:244        print(e)245