206 lines
6.4 KiB
Python
206 lines
6.4 KiB
Python
import asyncio
|
|
import logging
|
|
from asyncio import Task
|
|
from enum import Enum
|
|
|
|
from typing import Dict, List, Callable, Type
|
|
|
|
from .dispatcher import Dispatcher, ConnectionState
|
|
from .mc.mctypes import VarInt
|
|
from .mc.packet import Packet
|
|
from .mc.identity import Token
|
|
from .mc import proto, encryption
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
def _registry_from_state(state:ConnectionState) -> Dict[int, Dict[int, Type[Packet]]]:
|
|
if state == ConnectionState.HANDSHAKING:
|
|
return proto.handshaking.clientbound.REGISTRY
|
|
if state == ConnectionState.STATUS:
|
|
return proto.status.clientbound.REGISTRY
|
|
if state == ConnectionState.LOGIN:
|
|
return proto.login.clientbound.REGISTRY
|
|
if state == ConnectionState.PLAY:
|
|
return proto.play.clientbound.REGISTRY
|
|
return {}
|
|
|
|
_STATE_REGS = {
|
|
ConnectionState.HANDSHAKING : proto.handshaking,
|
|
ConnectionState.STATUS : proto.status,
|
|
ConnectionState.LOGIN : proto.login,
|
|
ConnectionState.PLAY : proto.play,
|
|
}
|
|
|
|
class Client:
|
|
host:str
|
|
port:int
|
|
token:Token
|
|
|
|
dispatcher : Dispatcher
|
|
_processing : bool
|
|
_worker : Task
|
|
|
|
_packet_callbacks : Dict[ConnectionState, Dict[Packet, List[Callable]]]
|
|
|
|
def __init__(
|
|
self,
|
|
token:Token,
|
|
host:str,
|
|
port:int = 25565,
|
|
):
|
|
self.host = host
|
|
self.port = port
|
|
self.token = token
|
|
|
|
self.dispatcher = Dispatcher(host, port)
|
|
self._processing = False
|
|
|
|
self._packet_callbacks = {}
|
|
|
|
def on(self, hook):
|
|
def wrapper(fun):
|
|
pass # TODO
|
|
return wrapper
|
|
|
|
def on_packet(self, packet:Type[Packet], state:ConnectionState) -> Callable:
|
|
def wrapper(fun):
|
|
if state not in self._packet_callbacks:
|
|
self._packet_callbacks[state] = {}
|
|
if packet not in self._packet_callbacks[state]:
|
|
self._packet_callbacks[state][packet] = []
|
|
self._packet_callbacks[state][packet].append(fun)
|
|
return fun
|
|
return wrapper
|
|
|
|
async def start(self):
|
|
await self.dispatcher.start()
|
|
self._processing = True
|
|
|
|
await self.dispatcher.write(
|
|
proto.handshaking.serverbound.PacketSetProtocol(
|
|
340, # TODO!!!!
|
|
protocolVersion=340,
|
|
serverHost=self.host,
|
|
serverPort=self.port,
|
|
nextState=2, # play
|
|
)
|
|
)
|
|
|
|
await self.dispatcher.write( # TODO PROTO!!!!
|
|
proto.login.serverbound.PacketLoginStart(340, username=self.token.profile.name)
|
|
)
|
|
|
|
self.dispatcher.state = ConnectionState.LOGIN
|
|
self._processing = True
|
|
self._worker = asyncio.get_event_loop().create_task(self._client_worker())
|
|
|
|
async def stop(self, block=True):
|
|
await self.dispatcher.stop()
|
|
self._processing = False
|
|
if block:
|
|
await self._worker
|
|
|
|
async def _client_worker(self):
|
|
while self._processing:
|
|
try:
|
|
# logger.info("Awaiting packet")
|
|
packet = await asyncio.wait_for(self.dispatcher.incoming.get(), timeout=5)
|
|
# logger.info("Client processing packet %s [state %s]", str(packet), str(self.dispatcher.state))
|
|
|
|
# Process packets? switch state, invoke callbacks? Maybe implement Reactors?
|
|
if self.dispatcher.state == ConnectionState.LOGIN:
|
|
await self.login_logic(packet)
|
|
elif self.dispatcher.state == ConnectionState.PLAY:
|
|
await self.play_logic(packet)
|
|
|
|
if self.dispatcher.state in self._packet_callbacks:
|
|
if Packet in self._packet_callbacks[self.dispatcher.state]: # callback for any packet
|
|
for cb in self._packet_callbacks[self.dispatcher.state][Packet]:
|
|
await cb(packet)
|
|
if packet.__class__ in self._packet_callbacks[self.dispatcher.state]: # callback for this packet
|
|
for cb in self._packet_callbacks[self.dispatcher.state][packet.__class__]:
|
|
await cb(packet)
|
|
self.dispatcher.incoming.task_done()
|
|
except asyncio.TimeoutError:
|
|
pass # need this to recheck self._processing periodically
|
|
except Exception:
|
|
logger.exception("Exception while processing packet %s", packet)
|
|
|
|
async def login_logic(self, packet:Packet):
|
|
if isinstance(packet, proto.login.clientbound.PacketEncryptionBegin):
|
|
logger.info("Encryption request")
|
|
secret = encryption.generate_shared_secret()
|
|
|
|
token, encrypted_secret = encryption.encrypt_token_and_secret(
|
|
packet.publicKey,
|
|
packet.verifyToken,
|
|
secret
|
|
)
|
|
|
|
if packet.serverId != '-' and self.token:
|
|
await self.token.join(
|
|
encryption.generate_verification_hash(
|
|
packet.serverId,
|
|
secret,
|
|
packet.publicKey
|
|
)
|
|
)
|
|
|
|
encryption_response = proto.login.serverbound.PacketEncryptionBegin(
|
|
340, # TODO!!!!
|
|
sharedSecret=encrypted_secret,
|
|
verifyToken=token
|
|
)
|
|
|
|
await self.dispatcher.write(encryption_response, wait=True)
|
|
|
|
await self.dispatcher.encrypt(secret)
|
|
|
|
elif isinstance(packet, proto.login.clientbound.PacketDisconnect):
|
|
logger.error("Disconnected while logging in")
|
|
await self.stop(False)
|
|
# raise Exception("Disconnected while logging in") # TODO make a more specific one, do some shit
|
|
|
|
elif isinstance(packet, proto.login.clientbound.PacketCompress):
|
|
logger.info("Set compression")
|
|
self.dispatcher.compression = packet.threshold
|
|
|
|
elif isinstance(packet, proto.login.clientbound.PacketSuccess):
|
|
logger.info("Login success")
|
|
self.dispatcher.state = ConnectionState.PLAY
|
|
|
|
elif isinstance(packet, proto.login.clientbound.PacketLoginPluginRequest):
|
|
pass # TODO ?
|
|
|
|
async def play_logic(self, packet:Packet):
|
|
if isinstance(packet, proto.play.clientbound.PacketSetCompression):
|
|
logger.info("Set compression")
|
|
self.dispatcher.compression = packet.threshold
|
|
|
|
elif isinstance(packet, proto.play.clientbound.PacketKeepAlive):
|
|
logger.info("Keep Alive")
|
|
keep_alive_packet = proto.play.serverbound.packet_keep_alive.PacketKeepAlive(340, keepAliveId=packet.keepAliveId)
|
|
await self.dispatcher.write(keep_alive_packet)
|
|
|
|
elif isinstance(packet, proto.play.clientbound.PacketPosition):
|
|
logger.info("PlayerPosLook")
|
|
# if self.connection.context.protocol_later_eq(107):
|
|
# teleport_confirm = serverbound.play.TeleportConfirmPacket()
|
|
# teleport_confirm.teleport_id = packet.teleport_id
|
|
# self.connection.write_packet(teleport_confirm)
|
|
# else:
|
|
# position_response = serverbound.play.PositionAndLookPacket()
|
|
# position_response.x = packet.x
|
|
# position_response.feet_y = packet.y
|
|
# position_response.z = packet.z
|
|
# position_response.yaw = packet.yaw
|
|
# position_response.pitch = packet.pitch
|
|
# position_response.on_ground = True
|
|
# self.connection.write_packet(position_response)
|
|
# self.connection.spawned = True
|
|
pass
|
|
|
|
elif isinstance(packet, proto.play.clientbound.PacketKickDisconnect):
|
|
logger.info("Play Disconnect")
|
|
raise Exception("Disconnected while playing") # TODO make a more specific one, do some shit
|
|
|