Re: mccp module

"""
Client decompression module for the mud client compression protocol.
See http://homepages.ihug.co.nz/~icecube/compress/ for more details.

Original C version by Oliver Jowett <icecube$ihug.co.nz>.
Python version by Conan Brink <conanb$psnw.com>.
Demangle addresses as needed.

This python code is placed in the public domain, as was the original
C source code from which it is derived.

The code was then converted into a Lyntin plugin module for Lyntin 4.*
by Will Guaraldi <willg$bluesock.org>.

Notes:
 - tested with lensmoor.org 3500 (uses v2)
"""
__author__ = "Will Guaraldi (et al)"
__version__ = "0.5 (23 March, 2004)"
__description__ = "MCCP module for Lyntin.  Supports v1 and v2."

from lyntin import exported, net
from lyntin.modules import modutils
import zlib

# The mc_state class represents the current compression state of a connection.
#
# Reading / writing:
#
#   Reading from the server connection must go through the decompressor at
#   all times. The decompressor handles both negotiation and decompression
#   transparently - it receives input directly from the server, then provides
#   the main client code with decompressed data, hiding all protocol details.
#
#   When data is received from the mud server, call receive, passing it
#   the data read.  It is VITAL that ALL data read is passed to the
#   decompressor - including data with embedded NULs!
#
#   After receive has been called, call pending() to see if any decompressed
#   data is available. It returns the number of bytes pending.
#
#   If there is pending data waiting, call get to retrieve it.
#   Your client can now process this data as if it had been directly read
#   from the server.
#
#   Regularly call response. If non-NULL, you need to write the returned
#   string to the mud server. This is needed when the decompressor
#   is negotiating compression with the server. When called, response
#   clears any pending string, so be sure to save its return value!
#
# Status information:
#
#   stats returns two values: the number of compressed bytes read, and the
#   number of bytes that they decompressed to.
#
#   compressing returns non-0 if the connection is currently using compression.


# Telnet values we're interested in
IAC  = chr(255)
DONT = chr(254)
DO   = chr(253)
WONT = chr(252)
WILL = chr(251)
SB   = chr(250)
SE   = chr(240)

TELOPT_COMPRESS  = chr(85)
TELOPT_COMPRESS2 = chr(86)

# We say DO COMPRESS2 to WILL COMPRESS2, then DONT COMPRESS to WILL COMPRESS
# -or-
# We say DO COMPRESS to WILL COMPRESS if it arrives before any COMPRESS2.
#
# Later the server sends IAC SB COMPRESS IAC SE (v2) or IAC SB COMPRESS WILL
# SE (v1), and immediately following that, begins compressing.
#
# Compression ends when the zlib stream ends, which causes the remaining
# uncompressed data to be in the state.dec.unused_data, from which we copy
# it back into the inbuf to feed into the telnet state machine.

will_v1 = "%s%s%s"     % (IAC, WILL, TELOPT_COMPRESS)
do_v1   = "%s%s%s"     % (IAC, DO,   TELOPT_COMPRESS)
dont_v1 = "%s%s%s"     % (IAC, DONT, TELOPT_COMPRESS)
on_v1   = "%s%s%s%s%s" % (IAC, SB,   TELOPT_COMPRESS, WILL, SE)

will_v2 = "%s%s%s"     % (IAC, WILL, TELOPT_COMPRESS2)
do_v2   = "%s%s%s"     % (IAC, DO,   TELOPT_COMPRESS2)
on_v2   = "%s%s%s%s%s" % (IAC, SB,   TELOPT_COMPRESS2, IAC, SE)

def write_socket(ses, data):
    sock = ses._socket.write(data, convert=0)


# state object 

class mc_state:
    def __init__(self):
        self.stream = 0
        self.inbuf = ""
        self.outbuf = ""
        self.comp = 0
        self.uncomp = 0
        self.resp_v1 = 0
        self.resp_v2 = 0
        self.got_v2 = 0
        self.dec = None

    def receive(self, data):
        """
        Perform decompression and negotiation on some received data.
        """
        # First, we append everything we got onto the end of self.inbuf
        self.inbuf = self.inbuf + data

        # Now, as long as there's any data left, we want to feed it through
        # the decompressor (if it's compressed), react to any telnet negotiation
        # that we recognize, save any incomplete telnet negotiation sequences,
        # and pass anything else intact

        while len(self.inbuf) > 0:
            if self.stream:
                # We're compressed, so let's decompress it
                decbytes = self.dec.decompress(self.inbuf)
                self.uncomp = self.uncomp + len(decbytes)
                self.outbuf = self.outbuf + decbytes
                if len(self.dec.unused_data) > 0:
                    # We ran out of compressed data and have some uncompressed
                    # data behind it.  Save that uncompressed stuff in case there's
                    # more telnet stuff in it. Nuke the decompression object so
                    # we know we're in an uncompressed state.
                    self.comp = self.comp + len(self.inbuf) - len(self.dec.unused_data)
                    self.inbuf = self.dec.unused_data
                    self.stream = 0
                    del self.dec
                    self.dec = None
                else:
                    self.comp = self.comp + len(self.inbuf)
                    self.inbuf = ""
            else:
                # We're not compressed.  Check for telnet stuff
                iacindex = self.inbuf.find(IAC)
                if (iacindex >= 0):
                    # Found something interesting.
                    # First give everything before it to the outbuf.
                    self.outbuf = self.outbuf + self.inbuf[:iacindex]
                    self.inbuf = self.inbuf[iacindex:]
                    if len(self.inbuf) == 1:
                        # We have to save this lone IAC byte until later
                        # because we don't yet know what sequence it starts
                        break;
                    if self.inbuf[1] == IAC:
                        # Pass this on unaltered.
                        self.outbuf = self.outbuf + self.inbuf[0:2]
                        self.inbuf = self.inbuf[2:]
                    elif self.inbuf[1] == WILL:
                        if len(self.inbuf) == 2:
                            # We have to save these two bytes because we don't
                            # yet know which option the server wants to negotiate
                            break;
                        if self.inbuf[2] == TELOPT_COMPRESS:
                            # Reject if we already got WILL COMPRESS2
                            # Accept if we didn't
                            if self.resp_v2:
                                self.resp_v1 = -1
                            else:
                                self.resp_v1 = 1
                        elif self.inbuf[2] == TELOPT_COMPRESS2:
                            # Note that we got it and want to answer it.
                            self.resp_v2 = 1
                            self.got_v2 = 1
                        else:
                            # Not for us.  Pass it on.
                            self.outbuf = self.outbuf + self.inbuf[0:3]
                        # Remove it from the datastream
                        self.inbuf = self.inbuf[3:]
                    elif self.inbuf[1] == SB:
                        seindex = self.inbuf.find(SE)
                        if seindex < 0:
                            # Protect ourselves against endless streaming from
                            # an evil mud.  No sub-negotiation deserves to be
                            # longer than 1 kilobyte.
                            if len(self.inbuf) > 1024:
                                self.outbuf = self.outbuf + self.inbuf[0:2]
                                self.inbuf = self.inbuf[2:]
                            else:
                                # Save for later.  This code hasn't ended yet.
                                break;
                        if self.inbuf[0:seindex + 1] == on_v1:
                            # Success!  We have gotten the sequence
                            # that turns on compression with version 1
                            self.stream = 1
                            self.dec = zlib.decompressobj(15)
                            # Now chop off the magic sequence so we have
                            # a nice string to feed to the decompressor.
                            self.inbuf = self.inbuf[5:]
                        elif self.inbuf[0:seindex + 1] == on_v2:
                            # Success!  We have gotten the sequence
                            # that turns on compression with version 2
                            self.stream = 1
                            self.dec = zlib.decompressobj(15)
                            # Now chop off the magic sequence so we have
                            # a nice string to feed to the decompressor.
                            self.inbuf = self.inbuf[5:]
                        else:
                            # It wasn't either of our sub-negotiation sequences.
                            # Pass it on to the client.
                            self.outbuf = self.outbuf + self.inbuf[0:seindex+1]
                            self.inbuf = self.inbuf[seindex+1:]
                    else:
                        if len(self.inbuf) < 3:
                            # save for later
                            break
                        # send the telnet option to someone else
                        self.outbuf = self.outbuf + self.inbuf[:3]
                        self.inbuf = self.inbuf[3:]
                else:
                    # No telnet stuff left in the string, and it's not even
                    # compressed.  Pass the whole thing to the outbuf
                    self.outbuf = self.outbuf + self.inbuf
                    self.inbuf = ""


    def pending(self):
        """
        Return the number of pending decompressed bytes that can currently
        be read by mudcompress_get
        """
        return len(self.outbuf)

    def get(self):
        """
        Read decompressed data from the decompressor.
        """
        s = self.outbuf
        self.outbuf = ""
        return s

    def stats(self):
        """
        Return the number of compressed bytes read and the number of bytes they
        expanded to, for this decompressor.
        """
        return (self.comp, self.uncomp)

    def response(self, ses):
        """
        Check for a negotiation response. If this returns None, no output is
        needed. If it returns anything else, that string should be sent to the
        mud server.  Calling this function clears the pending string (so be sure
        to save the result).
        """
        if self.resp_v1 == 1:
            self.resp_v1 = 0
            write_socket(ses, do_v1)
            return

        if self.resp_v1 == -1:
            self.resp_v1 = 0
            write_socket(ses, dont_v1)
            return

        if self.resp_v2:
            self.resp_v2 = 0
            write_socket(ses, do_v2)
            return

    def compressing(self):
        """
        Return true (non-0) if this decompressor has successfully negotiated
        compression and is currently performing decompression.
        """
        return self.stream


# -------------------------------------------------------------
# end of material from the mccpDecompress module
# -------------------------------------------------------------

from lyntin import manager

class MCCPManager(manager.Manager):
    def __init__(self):
        self._states = {}

    def connectHandler(self, args):
        ses = args["session"]
        self._states[ses] = mc_state()

    def netReadHandler(self, args):
        ses = args["session"]
        data = args["dataadj"]

        if not data:
            return data

        mcstate = self._states[ses]

        mcstate.receive(data)
        mcstate.response(ses)
        data = mcstate.get()

        return data

commands_dict = {}

def mccpstats_cmd(ses, args, input):
    """
    Tells you the compressed and decompressed stats for the session
    thus far.

    category: mccp
    """
    global mccpm

    if not mccpm._states.has_key(ses):
        exported.write_message("mccpstats: MCCP is not enabled for this session.")
        return

    mcstate = mccpm._states[ses]
    if mcstate.got_v2:
        comp = "(V2)"
    else:
        comp = "(V1)"

    exported.write_message("mccpstats: MCCP is enabled for this session.  %s" % comp)
    c, d = mcstate.stats()
    exported.write_message("mccpstats: %d bytes -> %d decompressed" % (c, d))
    return

commands_dict["mccpstats"] = (mccpstats_cmd, "")


mccpm = None

def load():
    global mccpm

    mccpm = MCCPManager()

    # register with connect hook to build an mc_state object per session
    exported.hook_register("connect_hook", mccpm.connectHandler)

    # register with net_read_data_filter hook to decompress incoming data
    exported.hook_register("net_read_data_filter", mccpm.netReadHandler)

    modutils.load_commands(commands_dict)

def unload():
    global mccpm

    exported.hook_unregister("connect_hook", mccpm.connectHandler)
    exported.hook_unregister("net_read_data_filter", mccpm.netReadHandler)
    modutils.unload_commands(commands_dict)