Colored messages

from colors import Color
from bitbuffers import BitBufferRead, BitBufferWrite
from engines.server import engine_server
from filters.recipients import RecipientFilter
from memory import Convention, DataType, get_object_pointer, make_object
from memory.hooks import PostHook, PreHook
from messages import get_message_index


HIGHLIGHT_COLOR = Color(255, 0, 150)
COLOR_MSG_COMMANDS = ("!c", "/c")

engine_server_ptr = get_object_pointer(engine_server)
saytext2_index = get_message_index('SayText2')


# CVEngineServer::UserMessageBegin(IRecipientFilter *, int)
user_message_begin = engine_server_ptr.make_virtual_function(
    43,
    Convention.THISCALL,
    [DataType.POINTER, DataType.POINTER, DataType.INT],
    DataType.POINTER
)

message_end = engine_server_ptr.make_virtual_function(
    44,
    Convention.THISCALL,
    [DataType.POINTER],
    DataType.VOID
)

user_message_data = None


@PostHook(user_message_begin)
def post_user_message_begin(args, ret_val):
    global user_message_data
    user_message_data = (
        make_object(RecipientFilter, args[1]),
        args[2],
        make_object(BitBufferWrite, ret_val)
    )


@PreHook(message_end)
def pre_message_end(args):
    # This happens when we initialize our hooks, while a user message is
    # currently being created
    if user_message_data is None:
        return

    recipients, message_index, buffer = user_message_data
    on_user_message_created(
        recipients, message_index, BitBufferRead(buffer), buffer)


def on_user_message_created(
        recipients, message_index, buffer_read, buffer_write):

    if message_index != saytext2_index:
        return

    index = buffer_read.read_byte()  # sender index (e.g. player's index)
    chat = buffer_read.read_byte()

    message_bits = buffer_read.num_bits_read # store current position
    message = buffer_read.read_string()

    param1 = buffer_read.read_string()
    param2 = buffer_read.read_string()
    param3 = buffer_read.read_string()
    param4 = buffer_read.read_string()

    for color_msg_cmd in COLOR_MSG_COMMANDS:
        if param2.startswith(color_msg_cmd):
            break
    else:
        return

    new_message = "\x01{color}{param2}".format(
        color=HIGHLIGHT_COLOR,
        param2=param2[len(color_msg_cmd):]
    )
    buffer_write.seek_to_bit(message_bits)
    buffer_write.write_string(new_message)

    if len(new_message) != len(message):
        buffer_write.write_string(param1)
        buffer_write.write_string(param2)
        buffer_write.write_string(param3)
        buffer_write.write_string(param4)