#!/usr/bin/python3
# -*- coding: utf-8 -*-

import frida
import sys
import json
import argparse
import subprocess
import base64

def validateMode(mode):
	if mode not in 'rw':
		raise argparse.ArgumentTypeError('Unknown mode')
	return mode

parser = argparse.ArgumentParser(description='bedrock_server packet tracer')
parser.add_argument('mode', help='"r" - read, "w" - write', type=validateMode)
parser.add_argument('packets', help='write packets that interest you', type=str, nargs='+')
args = parser.parse_args()

try:
	session = frida.attach('bedrock_server')
except frida.ProcessNotFoundError:
	sys.exit('Could not find bedrock_server')
except frida.PermissionDeniedError as e:
	sys.exit(e)

process = subprocess.Popen(['php', 'dump.php'], stdin=subprocess.PIPE, stdout=sys.stdout)

def onMessage(message, data):
	print(message['payload'])
	process.stdin.write(base64.b64encode(data) + b'\n')
	process.stdin.flush()

try:
	script = session.create_script("""
var stringLength = new NativeFunction(Module.findExportByName(null, '_ZNKSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEE6lengthEv'), 'int', ['pointer']);

recv('input', function(message) {
	var mode = message.mode;
	var doRead = mode.includes('r');
	var doWrite = mode.includes('w');
	JSON.parse(message.payload).forEach(function(packet) {
		var read = Module.findExportByName(null, '_ZN' + packet.length + packet + '4readER20ReadOnlyBinaryStream');
		var write = Module.findExportByName(null, '_ZNK' + packet.length + packet + '5writeER12BinaryStream');
		if (read === null || write === null) {
			console.log('Could not find symbol for ' + packet);
		} else {
			if (doRead) {
				Interceptor.attach(read, {
					onEnter: function(args) {
						this.pointer = args[1];
					},
					onLeave: function(retval) {
						var realAddr = Memory.readPointer(this.pointer.add(56));
						var rlen = stringLength(realAddr);
						send('from Client:', Memory.readByteArray(Memory.readPointer(realAddr), rlen));
					}
				});
			}
			if (doWrite) {
				Interceptor.attach(write, {
					onEnter: function(args) {
						this.pointer = args[1];
					},
					onLeave: function(retval) {
						var realAddr = Memory.readPointer(this.pointer.add(56));
						var rlen = stringLength(realAddr);
						send('from Server:', Memory.readByteArray(Memory.readPointer(realAddr), rlen));
					}
				});
			}
		}
	});
});
	""")

	script.on('message', onMessage)
	script.load()
	script.post({
		'type': 'input',
		'payload': json.dumps(args.packets),
		'mode': args.mode
	})

	sys.stdin.read()
except KeyboardInterrupt:
	sys.exit(0)