ydb/contrib/python/Twisted/py2/twisted/spread/banana.py

# -*- test-case-name: twisted.spread.test.test_banana -*-
# Copyright (c) Twisted Matrix Laboratories.
# See LICENSE for details.

"""
Banana -- s-exp based protocol.

Future Plans: This module is almost entirely stable.  The same caveat applies
to it as applies to L{twisted.spread.jelly}, however.  Read its future plans
for more details.

@author: Glyph Lefkowitz
"""

from __future__ import absolute_import, division

import copy, struct
from io import BytesIO

from twisted.internet import protocol
from twisted.persisted import styles
from twisted.python import log
from twisted.python.compat import iterbytes, long, _bytesChr as chr
from twisted.python.reflect import fullyQualifiedName

class BananaError(Exception):
    pass

def int2b128(integer, stream):
    if integer == 0:
        stream(chr(0))
        return
    assert integer > 0, "can only encode positive integers"
    while integer:
        stream(chr(integer & 0x7f))
        integer = integer >> 7


def b1282int(st):
    """
    Convert an integer represented as a base 128 string into an L{int} or
    L{long}.

    @param st: The integer encoded in a byte string.
    @type st: L{bytes}

    @return: The integer value extracted from the byte string.
    @rtype: L{int} or L{long}
    """
    e = 1
    i = 0
    for char in iterbytes(st):
        n = ord(char)
        i += (n * e)
        e <<= 7
    return i


# delimiter characters.
LIST     = chr(0x80)
INT      = chr(0x81)
STRING   = chr(0x82)
NEG      = chr(0x83)
FLOAT    = chr(0x84)
# "optional" -- these might be refused by a low-level implementation.
LONGINT  = chr(0x85)
LONGNEG  = chr(0x86)
# really optional; this is part of the 'pb' vocabulary
VOCAB    = chr(0x87)

HIGH_BIT_SET = chr(0x80)

def setPrefixLimit(limit):
    """
    Set the limit on the prefix length for all Banana connections
    established after this call.

    The prefix length limit determines how many bytes of prefix a banana
    decoder will allow before rejecting a potential object as too large.

    @type limit: L{int}
    @param limit: The number of bytes of prefix for banana to allow when
    decoding.
    """
    global _PREFIX_LIMIT
    _PREFIX_LIMIT = limit

_PREFIX_LIMIT = None
setPrefixLimit(64)

SIZE_LIMIT = 640 * 1024   # 640k is all you'll ever need :-)

class Banana(protocol.Protocol, styles.Ephemeral):
    """
    L{Banana} implements the I{Banana} s-expression protocol, client and
    server.

    @ivar knownDialects: These are the profiles supported by this Banana
        implementation.
    @type knownDialects: L{list} of L{bytes}
    """

    # The specification calls these profiles but this implementation calls them
    # dialects instead.
    knownDialects = [b"pb", b"none"]

    prefixLimit = None
    sizeLimit = SIZE_LIMIT

    def setPrefixLimit(self, limit):
        """
        Set the prefix limit for decoding done by this protocol instance.

        @see: L{setPrefixLimit}
        """
        self.prefixLimit = limit
        self._smallestLongInt = -2 ** (limit * 7) + 1
        self._smallestInt = -2 ** 31
        self._largestInt = 2 ** 31 - 1
        self._largestLongInt = 2 ** (limit * 7) - 1


    def connectionReady(self):
        """Surrogate for connectionMade
        Called after protocol negotiation.
        """


    def _selectDialect(self, dialect):
        self.currentDialect = dialect
        self.connectionReady()


    def callExpressionReceived(self, obj):
        if self.currentDialect:
            self.expressionReceived(obj)
        else:
            # this is the first message we've received
            if self.isClient:
                # if I'm a client I have to respond
                for serverVer in obj:
                    if serverVer in self.knownDialects:
                        self.sendEncoded(serverVer)
                        self._selectDialect(serverVer)
                        break
                else:
                    # I can't speak any of those dialects.
                    log.msg("The client doesn't speak any of the protocols "
                            "offered by the server: disconnecting.")
                    self.transport.loseConnection()
            else:
                if obj in self.knownDialects:
                    self._selectDialect(obj)
                else:
                    # the client just selected a protocol that I did not suggest.
                    log.msg("The client selected a protocol the server didn't "
                            "suggest and doesn't know: disconnecting.")
                    self.transport.loseConnection()


    def connectionMade(self):
        self.setPrefixLimit(_PREFIX_LIMIT)
        self.currentDialect = None
        if not self.isClient:
            self.sendEncoded(self.knownDialects)


    def gotItem(self, item):
        l = self.listStack
        if l:
            l[-1][1].append(item)
        else:
            self.callExpressionReceived(item)

    buffer = b''

    def dataReceived(self, chunk):
        buffer = self.buffer + chunk
        listStack = self.listStack
        gotItem = self.gotItem
        while buffer:
            assert self.buffer != buffer, "This ain't right: %s %s" % (repr(self.buffer), repr(buffer))
            self.buffer = buffer
            pos = 0
            for ch in iterbytes(buffer):
                if ch >= HIGH_BIT_SET:
                    break
                pos = pos + 1
            else:
                if pos > self.prefixLimit:
                    raise BananaError("Security precaution: more than %d bytes of prefix" % (self.prefixLimit,))
                return
            num = buffer[:pos]
            typebyte = buffer[pos:pos+1]
            rest = buffer[pos+1:]
            if len(num) > self.prefixLimit:
                raise BananaError("Security precaution: longer than %d bytes worth of prefix" % (self.prefixLimit,))
            if typebyte == LIST:
                num = b1282int(num)
                if num > SIZE_LIMIT:
                    raise BananaError("Security precaution: List too long.")
                listStack.append((num, []))
                buffer = rest
            elif typebyte == STRING:
                num = b1282int(num)
                if num > SIZE_LIMIT:
                    raise BananaError("Security precaution: String too long.")
                if len(rest) >= num:
                    buffer = rest[num:]
                    gotItem(rest[:num])
                else:
                    return
            elif typebyte == INT:
                buffer = rest
                num = b1282int(num)
                gotItem(num)
            elif typebyte == LONGINT:
                buffer = rest
                num = b1282int(num)
                gotItem(num)
            elif typebyte == LONGNEG:
                buffer = rest
                num = b1282int(num)
                gotItem(-num)
            elif typebyte == NEG:
                buffer = rest
                num = -b1282int(num)
                gotItem(num)
            elif typebyte == VOCAB:
                buffer = rest
                num = b1282int(num)
                item = self.incomingVocabulary[num]
                if self.currentDialect == b'pb':
                    # the sender issues VOCAB only for dialect pb
                    gotItem(item)
                else:
                    raise NotImplementedError(
                        "Invalid item for pb protocol {0!r}".format(item))
            elif typebyte == FLOAT:
                if len(rest) >= 8:
                    buffer = rest[8:]
                    gotItem(struct.unpack("!d", rest[:8])[0])
                else:
                    return
            else:
                raise NotImplementedError(("Invalid Type Byte %r" % (typebyte,)))
            while listStack and (len(listStack[-1][1]) == listStack[-1][0]):
                item = listStack.pop()[1]
                gotItem(item)
        self.buffer = b''


    def expressionReceived(self, lst):
        """Called when an expression (list, string, or int) is received.
        """
        raise NotImplementedError()


    outgoingVocabulary = {
        # Jelly Data Types
        b'None'           :  1,
        b'class'          :  2,
        b'dereference'    :  3,
        b'reference'      :  4,
        b'dictionary'     :  5,
        b'function'       :  6,
        b'instance'       :  7,
        b'list'           :  8,
        b'module'         :  9,
        b'persistent'     : 10,
        b'tuple'          : 11,
        b'unpersistable'  : 12,

        # PB Data Types
        b'copy'           : 13,
        b'cache'          : 14,
        b'cached'         : 15,
        b'remote'         : 16,
        b'local'          : 17,
        b'lcache'         : 18,

        # PB Protocol Messages
        b'version'        : 19,
        b'login'          : 20,
        b'password'       : 21,
        b'challenge'      : 22,
        b'logged_in'      : 23,
        b'not_logged_in'  : 24,
        b'cachemessage'   : 25,
        b'message'        : 26,
        b'answer'         : 27,
        b'error'          : 28,
        b'decref'         : 29,
        b'decache'        : 30,
        b'uncache'        : 31,
        }

    incomingVocabulary = {}
    for k, v in outgoingVocabulary.items():
        incomingVocabulary[v] = k


    def __init__(self, isClient=1):
        self.listStack = []
        self.outgoingSymbols = copy.copy(self.outgoingVocabulary)
        self.outgoingSymbolCount = 0
        self.isClient = isClient


    def sendEncoded(self, obj):
        """
        Send the encoded representation of the given object:

        @param obj: An object to encode and send.

        @raise BananaError: If the given object is not an instance of one of
            the types supported by Banana.

        @return: L{None}
        """
        encodeStream = BytesIO()
        self._encode(obj, encodeStream.write)
        value = encodeStream.getvalue()
        self.transport.write(value)


    def _encode(self, obj, write):
        if isinstance(obj, (list, tuple)):
            if len(obj) > SIZE_LIMIT:
                raise BananaError(
                    "list/tuple is too long to send (%d)" % (len(obj),))
            int2b128(len(obj), write)
            write(LIST)
            for elem in obj:
                self._encode(elem, write)
        elif isinstance(obj, (int, long)):
            if obj < self._smallestLongInt or obj > self._largestLongInt:
                raise BananaError(
                    "int/long is too large to send (%d)" % (obj,))
            if obj < self._smallestInt:
                int2b128(-obj, write)
                write(LONGNEG)
            elif obj < 0:
                int2b128(-obj, write)
                write(NEG)
            elif obj <= self._largestInt:
                int2b128(obj, write)
                write(INT)
            else:
                int2b128(obj, write)
                write(LONGINT)
        elif isinstance(obj, float):
            write(FLOAT)
            write(struct.pack("!d", obj))
        elif isinstance(obj, bytes):
            # TODO: an API for extending banana...
            if self.currentDialect == b"pb" and obj in self.outgoingSymbols:
                symbolID = self.outgoingSymbols[obj]
                int2b128(symbolID, write)
                write(VOCAB)
            else:
                if len(obj) > SIZE_LIMIT:
                    raise BananaError(
                        "byte string is too long to send (%d)" % (len(obj),))
                int2b128(len(obj), write)
                write(STRING)
                write(obj)
        else:
            raise BananaError("Banana cannot send {0} objects: {1!r}".format(
                fullyQualifiedName(type(obj)), obj))


# For use from the interactive interpreter
_i = Banana()
_i.connectionMade()
_i._selectDialect(b"none")


def encode(lst):
    """Encode a list s-expression."""
    encodeStream = BytesIO()
    _i.transport = encodeStream
    _i.sendEncoded(lst)
    return encodeStream.getvalue()


def decode(st):
    """
    Decode a banana-encoded string.
    """
    l = []
    _i.expressionReceived = l.append
    try:
        _i.dataReceived(st)
    finally:
        _i.buffer = b''
        del _i.expressionReceived
    return l[0]