minecraft-bot/mosfet/protocol/managers.py

import os
from math import floor
import json
import time

from minecraft.networking.packets import clientbound, serverbound
from mosfet.protocol import packets

from mosfet import utils
from mosfet.info import blocks

class DataManager:
    def __init__(self, directory):
        self.blocks = {}
        self.blocks_states = {}
        self.blocks_properties = {}
        self.registries = {}
        self.biomes = {}
        self.entity_type = {}
        
        if not os.path.isdir(directory):
            raise FileNotFoundError("%s is not a valid directory")

        if not os.path.isfile("%s/registries.json"%(directory)):
            raise FileNotFoundError("%s is not a valid minecraft data directory")
        
        with open("%s/blocks.json"%(directory)) as f:
            blocks = json.loads(f.read())
        for x in blocks:
            for s in blocks[x]['states']:
                self.blocks_states[s['id']] = x
                self.blocks_properties[s['id']] = s.get('properties', {})

        with open("%s/registries.json"%(directory)) as f:
            registries = json.loads(f.read())
        #for x in registries["minecraft:biome"]["entries"]:
        #    self.biomes[registries["minecraft:biome"]["entries"][x]["protocol_id"]] = x
        for x in registries["minecraft:entity_type"]["entries"]:
            self.entity_type[registries["minecraft:entity_type"]["entries"][x]["protocol_id"]] = x


class ChunksManager:
    def __init__(self, data_manager):
        self.data = data_manager
        self.chunks = {}
        #self.biomes = {}
        self.index = {}
        self.loading = False

    def handle_block(self, block_packet):
        self.set_block_at(block_packet.location.x, block_packet.location.y, block_packet.location.z, block_packet.block_state_id)

    def handle_multiblock(self, multiblock_packet):
        dx = 16 * multiblock_packet.chunk_section_pos.x
        dy = 16 * multiblock_packet.chunk_section_pos.y
        dz = 16 * multiblock_packet.chunk_section_pos.z
        for b in multiblock_packet.records:
            self.set_block_at(dx+b.x, dy+b.y, dz+b.z, b.block_state_id)

    def handle_chunk(self, chunk_packet):
        for i in chunk_packet.chunks:
            chunk = chunk_packet.chunks[i]
            self.chunks[(chunk_packet.x, i, chunk_packet.z)] = chunk

            if chunk.sub_index:
                dx = chunk.x * 16
                dy = chunk.y * 16
                dz = chunk.z * 16

                for item_id, locations in chunk.sub_index.items():
                    if item_id not in self.index:
                        self.index[item_id] = []
                    for l in locations:
                        coords = (dx + l%16, dy + l//256, dz + l%256//16)
                        self.index[item_id].append(coords)

        #self.biomes[(chunk_packet.x, None, chunk_packet.z)] = chunk_packet.biomes # FIXME
        if self.loading:
            print('.', end='', flush=True)

    def register(self, connection):
        connection.register_packet_listener(self.handle_block, clientbound.play.BlockChangePacket)
        connection.register_packet_listener(self.handle_multiblock, clientbound.play.MultiBlockChangePacket)
        connection.register_packet_listener(self.handle_chunk, packets.ChunkDataPacket)

    def get_chunk(self, x, y, z):
        index = (x, y, z)
        if not index in self.chunks:
            return None
            #raise ChunkNotLoadedException(index)
        return self.chunks[index]

    def get_loaded_area(self, ignore_empty=False):
        first = next(iter(self.chunks.keys()))
        x0 = x1 = first[0]
        y0 = y1 = first[1]
        z0 = z1 = first[2]
        for k in self.chunks.keys():
            if ignore_empty and self.chunks[k].empty:
                continue
            x0 = min(x0, k[0])
            x1 = max(x1, k[0])
            y0 = min(y0, k[1])
            y1 = max(y1, k[1])
            z0 = min(z0, k[2])
            z1 = max(z1, k[2])
        return ((x0,y0,z0),(x1,y1,z1))

    def get_block_at(self, x, y, z):
        c = self.get_chunk(floor(x/16), floor(y/16), floor(z/16))
        if not c: return None
        return c.get_block_at(x%16, y%16, z%16)

    def set_block_at(self, x, y, z, block):
        c = self.get_chunk(floor(x/16), floor(y/16), floor(z/16))
        if not c: return None
        c.set_block_at(x%16, y%16, z%16, block)

        if block in blocks.INDEXED_IDS:
            if block not in self.index:
                self.index[block] = []
            self.index[block].append((x, y, z))

    def check_loaded(self, chunk_distance):
        num = (chunk_distance * 2 + 1) ** 2
        num_subchunks = num * 16
        return len(self.chunks) >= num_subchunks

    def unload_chunk(self, x, z):
        for y in range(16):
            try:
                del self.chunks[(x, y, z)]
            except KeyError:
                pass

    def unload_all_chunks(self):
        self.chunks = {}
        self.index = {}

class ChunkNotLoadedException(Exception):
    def __str__(self):
        pos = self.args[0]
        return "Chunk at %d %d %d not loaded (yet?)"%(pos[0], pos[1], pos[2])


class ChatManager:
    def __init__(self, global_state):
        self.g = global_state
        self.handler = None

        self.g.connection.register_packet_listener(self.print_chat, clientbound.play.ChatMessagePacket)

    def translate_chat(self, data):
        if isinstance(data, str):
            return data

        result = data.get('text', '')
        result += data.get('translate', '')

        if 'with' in data:
            result += ' ' + ' '.join([self.translate_chat(x) for x in data['with']])
        elif 'extra' in data:
            result += ''.join([self.translate_chat(x) for x in data['extra']])

        return result

    def print_chat(self, chat_packet):
        #print(chat_packet)
        #print(chat_packet.json_data)
        #import json
        #print(json.dumps(json.loads(chat_packet.json_data), indent=4))
        try:
            source = chat_packet.field_string('position')
            sender = chat_packet.sender
            text = self.translate_chat(json.loads(chat_packet.json_data))
            print('[%s] %s'%(source, text))
        except Exception as ex:
            print('Exception %r on message (%s): %s' % (ex, chat_packet.field_string('position'), chat_packet.json_data))
            return

        if self.handler:
            self.handler((source, sender, text))

    def set_handler(self, func):
        self.handler = func

    def send(self, text):
        if not text:
            # Prevents connection bug when sending empty chat message
            return
        packet = serverbound.play.ChatPacket()
        packet.message = text
        self.g.connection.write_packet(packet)