From e0fd0e6c2d29fbfd38731abe56ac1ddf198969ca Mon Sep 17 00:00:00 2001 From: ksshen0000 Date: Sun, 18 Feb 2024 00:34:45 -0600 Subject: [PATCH 1/5] add native websocket support --- docs/source/tubes/wstube.rst | 13 ++++ pwnlib/tubes/__init__.py | 3 +- pwnlib/tubes/wstube.py | 129 +++++++++++++++++++++++++++++++++++ pyproject.toml | 1 + 4 files changed, 145 insertions(+), 1 deletion(-) create mode 100644 docs/source/tubes/wstube.rst create mode 100644 pwnlib/tubes/wstube.py diff --git a/docs/source/tubes/wstube.rst b/docs/source/tubes/wstube.rst new file mode 100644 index 000000000..5848dcc07 --- /dev/null +++ b/docs/source/tubes/wstube.rst @@ -0,0 +1,13 @@ +.. testsetup:: * + + from pwn import * + from websocket import WebSocket, ABNF, WebSocketException, WebSocketTimeoutException + +:mod:`pwnlib.tubes.wstube` --- Wstube +=========================================================== + +.. automodule:: pwnlib.tubes.wstube + + .. autoclass:: pwnlib.tubes.wstube.wstube + :members: + diff --git a/pwnlib/tubes/__init__.py b/pwnlib/tubes/__init__.py index 9f55c3321..6f8db211d 100644 --- a/pwnlib/tubes/__init__.py +++ b/pwnlib/tubes/__init__.py @@ -19,5 +19,6 @@ from pwnlib.tubes import sock from pwnlib.tubes import ssh from pwnlib.tubes import tube +from pwnlib.tubes import wstube -__all__ = ['tube', 'sock', 'remote', 'listen', 'process', 'serialtube', 'server', 'ssh'] +__all__ = ['tube', 'sock', 'remote', 'listen', 'process', 'serialtube', 'server', 'ssh', 'wstube'] diff --git a/pwnlib/tubes/wstube.py b/pwnlib/tubes/wstube.py new file mode 100644 index 000000000..dc4c63553 --- /dev/null +++ b/pwnlib/tubes/wstube.py @@ -0,0 +1,129 @@ + +# * Portions of this file are derived from pwntools-tube-websocket by frankli0324 +# * under the MIT License. +# * +# * Copyright (c) frankli0324. +# * https://gist.github.com/frankli0324/795162a14be988a01e0efa0531f7ac5a +# * +# * Permission is hereby granted, free of charge, to any person obtaining a copy +# * of this software and associated documentation files (the "Software"), to deal +# * in the Software without restriction, including without limitation the rights +# * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# * copies of the Software, and to permit persons to whom the Software is +# * furnished to do so, subject to the following conditions: +# * +# * The above copyright notice and this permission notice shall be included in all +# * copies or substantial portions of the Software. +from __future__ import absolute_import +from __future__ import division +from websocket import WebSocket, ABNF, WebSocketException, WebSocketTimeoutException + +from pwnlib.tubes.tube import tube + + +class wstube(tube): + """ + A basic websocket interface that wrapped as a tube. + + Arguments: + url (str): The websocket server's URL to connect to. + headers (dict): The same headers as the websocket protocol. + + Examples: + + >>> ws = wstube('wss://echo.websocket.events') + >>> ws.recv() + b'echo.websocket.events sponsored by Lob.com' + >>> for i in range(3): + ... ws.send(b'test') + ... print(ws.recv(2)) + ... print(ws.recv(2)) + b'te' + b'st' + b'te' + b'st' + b'te' + b'st' + >>> ws.sendline(b'test') + >>> print(ws.recv()) + b'test\\n' + >>> ws.send(b'12345asdfg') + >>> print(ws.recvregex(b'[0-9]{5}')) + b'12345' + >>> print(ws.recv()) + b'asdfg' + >>> ws.close() + """ + def __init__(self, url, headers=None, *args, **kwargs): + if headers is None: + headers = {} + super(wstube, self).__init__(*args, **kwargs) + self.closed = False + self.sock = WebSocket() + self.url = url + self.sock.connect(url, header=headers) + + + def recv_raw(self, numb): + if self.closed: + raise EOFError + + while True: + try: + data = self.sock.recv() + if isinstance(data, str): + data = data.encode() + break + except WebSocketTimeoutException: + return None + except WebSocketException: + self.shutdown("recv") + raise EOFError + + if not data: + self.shutdown() + raise EOFError('Recv Error') + + return data + + def send_raw(self, data): + if self.closed: + raise EOFError + + try: + self.sock.send_binary(data) + except WebSocketException as e: + self.shutdown() + raise EOFError('Send Error') + + def settimeout_raw(self, timeout): + if getattr(self, 'sock', None): + self.sock.settimeout(timeout) + + def connected_raw(self, direction): + try: + self.sock.ping() + opcode, data = self.sock.recv_data(True) + return opcode == ABNF.OPCODE_PONG + except: + return False + + def close(self): + if not getattr(self, 'sock', None): + return + + self.closed = True + + self.sock.close() + self.sock = None + self._close_msg() + + def _close_msg(self): + self.info('Closed connection to %s', self.url) + + def shutdown_raw(self, direction): + if self.closed: + return + + self.closed = True + self.sock.shutdown() \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index ce59233f9..0642f8ffd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -57,6 +57,7 @@ dependencies = [ "pathlib2; python_version < '3.4'", "unix-ar; python_version >= '3'", "zstandard", + "websocket-client", ] [project.urls] From 1118bcd62391df5c1580e8c275ff9e82be240077 Mon Sep 17 00:00:00 2001 From: ksshen0000 Date: Sun, 18 Feb 2024 01:27:42 -0600 Subject: [PATCH 2/5] add native websocket support --- pwn/toplevel.py | 1 + 1 file changed, 1 insertion(+) diff --git a/pwn/toplevel.py b/pwn/toplevel.py index 92b01ad1d..4c61b8281 100644 --- a/pwn/toplevel.py +++ b/pwn/toplevel.py @@ -52,6 +52,7 @@ from pwnlib.tubes.server import server from pwnlib.tubes.ssh import ssh from pwnlib.tubes.tube import tube +from pwnlib.tubes.wstube import wstube from pwnlib.ui import * from pwnlib.util import crc from pwnlib.util import iters From 24c5d5d701d26f70a5f491e98e5e07a53f436671 Mon Sep 17 00:00:00 2001 From: ksshen0000 Date: Sun, 18 Feb 2024 01:30:24 -0600 Subject: [PATCH 3/5] add pr to change log --- CHANGELOG.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 62647edbf..674de50bd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -69,7 +69,7 @@ The table below shows which release corresponds to each branch, and what date th | [2.2.0](#220) | | Jan 5, 2015 ## 4.13.0 (`dev`) - +- [#2350][2350] add native websocket support - [#2242][2242] Term module revamp: activating special handling of terminal only when necessary - [#2277][2277] elf: Resolve more relocations into GOT entries - [#2281][2281] FIX: Getting right amount of data for search fix @@ -99,6 +99,7 @@ The table below shows which release corresponds to each branch, and what date th - [#2347][2347] Fix/workaround Unicorn Engine 1GB limit that calls exit() - [#2233][2233] Fix gdb.debug: exe parameter now respected, allow empty argv +[2350]: https://github.com/Gallopsled/pwntools/pull/2350 [2242]: https://github.com/Gallopsled/pwntools/pull/2242 [2277]: https://github.com/Gallopsled/pwntools/pull/2277 [2281]: https://github.com/Gallopsled/pwntools/pull/2281 From 5536d9382ad3960d89375685b54714fd7782f345 Mon Sep 17 00:00:00 2001 From: ksshen0000 Date: Sun, 18 Feb 2024 01:45:14 -0600 Subject: [PATCH 4/5] fix bytewarning in doctest by removing print() --- pwnlib/tubes/wstube.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/pwnlib/tubes/wstube.py b/pwnlib/tubes/wstube.py index dc4c63553..b4c24104f 100644 --- a/pwnlib/tubes/wstube.py +++ b/pwnlib/tubes/wstube.py @@ -36,8 +36,8 @@ class wstube(tube): b'echo.websocket.events sponsored by Lob.com' >>> for i in range(3): ... ws.send(b'test') - ... print(ws.recv(2)) - ... print(ws.recv(2)) + ... ws.recv(2) + ... ws.recv(2) b'te' b'st' b'te' @@ -45,12 +45,12 @@ class wstube(tube): b'te' b'st' >>> ws.sendline(b'test') - >>> print(ws.recv()) + >>> ws.recv() b'test\\n' >>> ws.send(b'12345asdfg') - >>> print(ws.recvregex(b'[0-9]{5}')) + >>> ws.recvregex(b'[0-9]{5}') b'12345' - >>> print(ws.recv()) + >>> ws.recv() b'asdfg' >>> ws.close() """ From 2bde202eb1f16803e635f8bd049dbe376e611bb9 Mon Sep 17 00:00:00 2001 From: ksshen0000 Date: Sun, 18 Feb 2024 11:04:32 -0600 Subject: [PATCH 5/5] add websocket-client to docs/requirement.txt --- docs/requirements.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/requirements.txt b/docs/requirements.txt index cca064106..0f23c8256 100755 --- a/docs/requirements.txt +++ b/docs/requirements.txt @@ -20,3 +20,4 @@ sphinx==1.8.6; python_version<'3' sphinx>=4.5.0; python_version>='3' sphinx_rtd_theme sphinxcontrib-autoprogram<=0.1.5 +websocket-client