Source

txMysql / txmysql / packet.py

Full commit
"""
Includes format strings that describe packets
and classes to read/write packets.

"""

import struct

from zope import interface as zinterface

from txmysql import imysql


# packers
packetHeaderLength = 4

def pack_uint8(x):
    bytes = struct.pack("B", x)
    return bytes

def pack_uint24(x):
    bytes = struct.pack("BBB", x&0xFF, (x>>8)&0xFF, (x>>16)&0xFF)
    return bytes

def packHeader(length, order):
    return pack_uint24(length) + pack_uint8(order)

# unpackers
def unpack_uint8(byte):
    return struct.unpack("B", byte)[0]

def unpack_uint24(bytes):
    return ((struct.unpack("B", bytes[0])[0]) +
            (struct.unpack("B", bytes[1])[0]<<8) +
            (struct.unpack("B", bytes[2])[0]<<16))

def unpackHeader(bytes):
    return (unpack_uint24(bytes[0:3]),
            unpack_uint8(bytes[3]))


# Packet Class
class Packet(object):
    zinterface.implements(imysql.IPacket)

    def __init__(self, length, order, remainingBytes=""):
        self.length = length
        self.order = order
        self.bytes = remainingBytes

    @classmethod
    def fromBytes(cls, bytes):
        length, order = unpackHeader(bytes[0:4])
        remainingBytes = bytes[4:]
        return Packet(length, order, remainingBytes)

    def toBytes(self):
        return packHeader(self.length, self.order)


handshakeInitializationFormat = "=B%ssxL8sxHBHxxxxxxxxxxxxx13s"


def packHandshakeInitialization(protocolVersion,
                                serverVersion,
                                threadId,
                                scrambleBuffer,
                                serverCapabilities,
                                serverLanguage,
                                serverStatus,
                                restOfScramble):
    bytes = ""
    bytes += struct.pack(handshakeInitializationFormat % (
            len(serverVersion),
            ),
                         protocolVersion,
                         serverVersion,
                         threadId,
                         scrambleBuffer,
                         serverCapabilities,
                         serverLanguage,
                         serverStatus,
                         restOfScramble)
    return bytes


def unpackHandshakeInitialization(bytes):
    index = bytes.find("\x00")
    values = struct.unpack(handshakeInitializationFormat % (index-1,), bytes)
    return values


# class HandshakeInitializationPacket(Packet):
#     def __init__(self, bytes):
#         super(HandshakeInitializationPacket, self).__init__(bytes)
#         handshakeValues = unpackHandshakeInitialization(self.bytes)
#         self.serverVersion = handshakeValues[1]

class HandshakeInitializationAdapter(Packet):
    zinterface.implements(imysql.IHandshakeInitializationPacket)

    __used_for__ = imysql.IPacket

    def __init__(self, packet):
        Packet.__init__(self, packet.length, packet.order, packet.bytes)
        self.context = packet
        handshakeValues = unpackHandshakeInitialization(self.bytes)
        self.protocolVersion = handshakeValues[0]
        self.serverVersion = handshakeValues[1]
        self.threadId = handshakeValues[2]
        self.scrambleBuffer = handshakeValues[3]
        self.serverCapabilities = handshakeValues[4]
        self.serverLanguage = handshakeValues[5]
        self.serverStatus = handshakeValues[6]
        self.restOfScrambleBuffer = handshakeValues[7]

imysql.registry.registerAdapter(HandshakeInitializationAdapter,
                                (imysql.IPacket,),
                                imysql.IHandshakeInitializationPacket)


class AuthenticationAdapter(Packet):
    zinterface.implements(imysql.IAuthenticationPacket)

    __used_for__ = imysql.IHandshakeInitializationPacket

    def __init__(self, packet):
        Packet.__init__(self, packet.length, packet.order, packet.bytes)
        self.context = packet

    def toBytes(self):
        pass

imysql.registry.registerAdapter(AuthenticationAdapter,
                                (imysql.IHandshakeInitializationPacket,),
                                imysql.IAuthenticationPacket)