add native websocket support

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

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

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:
+

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

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']

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')
+        ...     ws.recv(2)
+        ...     ws.recv(2)
+        b'te'
+        b'st'
+        b'te'
+        b'st'
+        b'te'
+        b'st'
+        >>> ws.sendline(b'test')
+        >>> ws.recv()
+        b'test\\n'
+        >>> ws.send(b'12345asdfg')
+        >>> ws.recvregex(b'[0-9]{5}')
+        b'12345'
+        >>> 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()

pyproject.toml

@@ -57,6 +57,7 @@ dependencies = [
     "pathlib2; python_version < '3.4'",
     "unix-ar; python_version >= '3'",
     "zstandard",
+    "websocket-client",
 ]
 
 [project.urls]