Source

mclib / connection.py

# -*- coding: utf-8 -*-

import gevent.socket as socket

from packet import *



class ProtocolError(Exception):
    pass

class Connection(object):
    '''
    Wrap basic operations on a connection to a server.
    '''

    PROTOCOL_VERSION = 23

    def __init__(self, username):
        self.username = username
        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

    def connect(self, address):
        try:
            self.socket.connect(address)
        except socket.error, e:
            self.on_disconnection(e)

        # TODO: handle properly Disconnect/Kick packets.
        self.send_packet(HandshakeClient(
            username=self.username
        ))
        response = self.recv_packet()
        self.expect_packet(response, HandshakeServer)
        self.connection_hash = response.connection_hash

        self.send_packet(LoginRequestClient(
            protocol_version=self.PROTOCOL_VERSION,
            username=self.username
        ))
        response = self.recv_packet()
        self.expect_packet(response, LoginRequestServer)

    def send_packet(self, packet):
        try:
            self.socket.sendall(packet.serialize())
        except socket.error, e:
            self.on_disconnection(e)
            raise

    def recv_packet(self):
        try:
            return Packet.parse_server(self.socket)
        except socket.error, e:
            self.on_disconnection(e)
            raise

    def expect_packet(self, packet, expected):
        if not isinstance(packet, expected):
            raise ProtocolError(
                "Expected a %s packet, but server sent a %s one" % (
                    expected.__name__,
                    packet.__class__.__name__
                )
            )

    def on_disconnection(self, exception):
        '''
        Simple hook, can be overridden by subclasses. 'exception' contains the
        exception that caused the disconnection, or None if the disconnection
        was expected.
        '''
        pass

    def close(self):
        self.socket.close()
        self.on_disconnection(None)