/brz/remove-bazaar

# Copyright (C) 2007 Canonical Ltd

#

# This program is free software; you can redistribute it and/or modify

# it under the terms of the GNU General Public License as published by

# the Free Software Foundation; either version 2 of the License, or

# (at your option) any later version.

#

# This program is distributed in the hope that it will be useful,

# but WITHOUT ANY WARRANTY; without even the implied warranty of

# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the

# GNU General Public License for more details.

#

# You should have received a copy of the GNU General Public License

# along with this program; if not, write to the Free Software

# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA

"""Pyrex implementation for bencode coder/decoder"""

cdef extern from "Python.h":

    ctypedef int  Py_ssize_t

    object Py_BuildValue(char *format, ...)

    int PyInt_CheckExact(object o)

    int PyLong_CheckExact(object o)

    int PyString_CheckExact(object o)

    int PyTuple_CheckExact(object o)

    int PyList_CheckExact(object o)

    int PyDict_CheckExact(object o)

    int PyBool_Check(object o)

    object PyString_FromStringAndSize(char *v, Py_ssize_t len)

    int PyString_AsStringAndSize(object o, char **buffer, Py_ssize_t *length)

    long PyInt_GetMax()

    object PyLong_FromString(char *str, char **pend, int base)

cdef extern from "stddef.h":

    ctypedef unsigned int size_t

cdef extern from "stdlib.h":

    void free(void *memblock)

    void *malloc(size_t size)

    void *realloc(void *memblock, size_t size)

cdef extern from "string.h":

    void *memcpy(void *dest, void *src, size_t count)

cdef extern from "_bencode_c.h":

    int snprintf(char* buffer, size_t nsize, char* fmt, ...)

class NotEnoughMemory(Exception):

    """Memory allocation error"""

    pass

cdef enum:  # Codes for used characters

    MINUS   = 0x2D      # ord(-)

    CHAR_0  = 0x30      # ord(0)

    CHAR_1  = 0x31      # ord(1)

    CHAR_2  = 0x32      # ord(2)

    CHAR_3  = 0x33      # ord(3)

    CHAR_4  = 0x34      # ord(4)

    CHAR_5  = 0x35      # ord(5)

    CHAR_6  = 0x36      # ord(6)

    CHAR_7  = 0x37      # ord(7)

    CHAR_8  = 0x38      # ord(8)

    CHAR_9  = 0x39      # ord(9)

    COLON   = 0x3A      # ord(:)

    SMALL_D = 0x64      # ord(d)

    SMALL_E = 0x65      # ord(e)

    SMALL_I = 0x69      # ord(i)

    SMALL_L = 0x6c      # ord(l)

cdef class Decoder:

    """Bencode decoder"""

    cdef readonly object __s

    cdef readonly char *tail

    cdef readonly int   size

    cdef readonly long   _MAXINT

    cdef readonly int    _MAXN

    cdef readonly object _longint

    def __init__(self, s):

        """Initialize decoder engine.

        @param  s:  Python string.

        """

        cdef Py_ssize_t k

        cdef char *pstr

        if not PyString_CheckExact(s):

            raise TypeError

        PyString_AsStringAndSize(s, &pstr, &k)

        if pstr == NULL:

            raise ValueError

        self.__s = s

        self.tail = pstr

        self.size = <int>k

        self._MAXINT = PyInt_GetMax()

        self._MAXN = len(str(self._MAXINT))

        self._longint = long(0)

    def __repr__(self):

        return 'Decoder(%s)' % repr(self.__s)

    def decode(self):

        result = self.decode_object()

        if self.size != 0:

            raise ValueError('junk in stream')

        return result

    def decode_object(self):

        cdef char ch

        if 0 == self.size:

            raise ValueError('stream underflow')

        ch = self.tail[0]

        if ch == SMALL_I:

            self._update_tail(1)

            return self._decode_int()

        elif CHAR_0 <= ch <= CHAR_9:

            return self._decode_string()

        elif ch == SMALL_L:

            self._update_tail(1)

            return self._decode_list()

        elif ch == SMALL_D:

            self._update_tail(1)

            return self._decode_dict()

        raise ValueError('unknown object')

    cdef void _update_tail(self, int n):

        """Update tail pointer and resulting size by n characters"""

        self.size = self.size - n

        self.tail = &self.tail[n]

    cdef object _decode_int(self):

        cdef int result

        result = self._decode_int_until(SMALL_E)

        if result != self._MAXINT:

            return result

        else:

            return self._longint

    cdef int _decode_int_until(self, char stop_char) except? -1:

        """Decode int from stream until stop_char encountered"""

        cdef int result

        cdef int i, n

        cdef int sign

        cdef char digit

        cdef char *longstr

        for n from 0 <= n < self.size:

            if self.tail[n] == stop_char:

                break

        else:

            raise ValueError

        sign = 0

        if MINUS == self.tail[0]:

            sign = 1

        if n-sign == 0:

            raise ValueError    # ie / i-e

        if self.tail[sign] == CHAR_0:   # special check for zero

            if sign:

                raise ValueError    # i-0e

            if n > 1:

                raise ValueError    # i00e / i01e

            self._update_tail(n+1)

            return 0

        if n-sign < self._MAXN:

            # plain int

            result = 0

            for i from sign <= i < n:

                digit = self.tail[i]

                if CHAR_0 <= digit <= CHAR_9:

                    result = result * 10 + (digit - CHAR_0)

                else:

                    raise ValueError

            if sign:

                result = -result

            self._update_tail(n+1)

        else:

            # long int

            result = self._MAXINT

            longstr = <char*>malloc(n+1)

            if NULL == longstr:

                raise NotEnoughMemory

            memcpy(longstr, self.tail, n)

            longstr[n] = 0

            self._longint = PyLong_FromString(longstr, NULL, 10)

            free(longstr)

            self._update_tail(n+1)

        return result

    cdef object _decode_string(self):

        cdef int n

        n = self._decode_int_until(COLON)

        if n == 0:

            return ''

        if n == self._MAXINT:

            # strings longer than 1GB is not supported

            raise ValueError('too long string')

        if n > self.size:

            raise ValueError('stream underflow')

        result = PyString_FromStringAndSize(self.tail, n)

        self._update_tail(n)

        return result

    cdef object _decode_list(self):

        result = []

        while self.size > 0:

            if self.tail[0] == SMALL_E:

                self._update_tail(1)

                return result

            else:

                result.append(self.decode_object())

        raise ValueError('malformed list')

    cdef object _decode_dict(self):

        cdef char ch

        result = {}

        lastkey = None

        while self.size > 0:

            ch = self.tail[0]

            if ch == SMALL_E:

                self._update_tail(1)

                return result

            elif CHAR_0 <= ch <= CHAR_9:

                # keys should be strings only

                key = self._decode_string()

                if lastkey >= key:

                    raise ValueError('dict keys disordered')

                else:

                    lastkey = key

                value = self.decode_object()

                result[key] = value

            else:

                raise ValueError('keys in dict should be strings only')

        raise ValueError('malformed dict')

def bdecode(object s):

    """Decode string x to Python object"""

    return Decoder(s).decode()

class Bencached(object):

    __slots__ = ['bencoded']

    def __init__(self, s):

        self.bencoded = s

cdef enum:

    INITSIZE = 1024     # initial size for encoder buffer

cdef class Encoder:

    """Bencode encoder"""

    cdef readonly char *buffer

    cdef readonly int   maxsize

    cdef readonly char *tail

    cdef readonly int   size

    def __init__(self, int maxsize=INITSIZE):

        """Initialize encoder engine

        @param  maxsize:    initial size of internal char buffer

        """

        cdef char *p

        self.maxsize = 0

        self.size = 0

        self.tail = NULL

        p = <char*>malloc(maxsize)

        if p == NULL:

            raise NotEnoughMemory('Not enough memory to allocate buffer '

                                  'for encoder')

        self.buffer = p

        self.maxsize = maxsize

        self.tail = p

    def __del__(self):

        free(self.buffer)

        self.buffer = NULL

        self.maxsize = 0

    def __str__(self):

        if self.buffer != NULL and self.size != 0:

            return PyString_FromStringAndSize(self.buffer, self.size)

        else:

            return ''

    cdef int _ensure_buffer(self, int required) except 0:

        """Ensure that tail of CharTail buffer has enough size.

        If buffer is not big enough then function try to

        realloc buffer.

        """

        cdef char *new_buffer

        cdef int   new_size

        if self.size + required < self.maxsize:

            return 1

        new_size = self.maxsize

        while new_size < self.size + required:

            new_size = new_size * 2

        new_buffer = <char*>realloc(self.buffer, <size_t>new_size)

        if new_buffer == NULL:

            raise NotEnoughMemory('Cannot realloc buffer for encoder')

        self.buffer = new_buffer

        self.maxsize = new_size

        self.tail = &new_buffer[self.size]

        return 1

    cdef void _update_tail(self, int n):

        """Update tail pointer and resulting size by n characters"""

        self.size = self.size + n

        self.tail = &self.tail[n]

    cdef int _encode_int(self, int x) except 0:

        """Encode int to bencode string iNNNe

        @param  x:  value to encode

        """

        cdef int n

        self._ensure_buffer(32)

        n = snprintf(self.tail, 32, "i%de", x)

        if n < 0:

            raise NotEnoughMemory('int %d too big to encode' % x)

        self._update_tail(n)

        return 1

    cdef int _encode_long(self, x) except 0:

        return self._append_string(''.join(('i', str(x), 'e')))

    cdef int _append_string(self, s) except 0:

        cdef Py_ssize_t k

        cdef int n

        cdef char *pstr

        PyString_AsStringAndSize(s, &pstr, &k)

        k = (<int>k + 1)

        self._ensure_buffer(<int>k)

        n = snprintf(self.tail, k, '%s', pstr)

        if n < 0:

            raise NotEnoughMemory('string %s too big to append' % s)

        self._update_tail(n)

        return 1

    cdef int _encode_string(self, x) except 0:

        cdef Py_ssize_t k

        cdef int n

        cdef char *pstr

        PyString_AsStringAndSize(x, &pstr, &k)

        self._ensure_buffer(<int>k+32)

        n = snprintf(self.tail, k+32, '%d:%s', <int>k, pstr)

        if n < 0:

            raise NotEnoughMemory('string %s too big to encode' % x)

        self._update_tail(n)

        return 1

    cdef int _encode_list(self, x) except 0:

        self._ensure_buffer(2)

        self.tail[0] = SMALL_L

        self._update_tail(1)

        for i in x:

            self.process(i)

        self.tail[0] = SMALL_E

        self._update_tail(1)

        return 1

    cdef int _encode_dict(self, x) except 0:

        self._ensure_buffer(2)

        self.tail[0] = SMALL_D

        self._update_tail(1)

        keys = x.keys()

        keys.sort()

        for k in keys:

            if not PyString_CheckExact(k):

                raise TypeError('key in dict should be string')

            self._encode_string(k)

            self.process(x[k])

        self.tail[0] = SMALL_E

        self._update_tail(1)

        return 1

    def process(self, object x):

        if PyInt_CheckExact(x):

            self._encode_int(x)

        elif PyLong_CheckExact(x):

            self._encode_long(x)

        elif PyString_CheckExact(x):

            self._encode_string(x)

        elif PyList_CheckExact(x) or PyTuple_CheckExact(x):

            self._encode_list(x)

        elif PyDict_CheckExact(x):

            self._encode_dict(x)

        elif PyBool_Check(x):

            self._encode_int(int(x))

        elif isinstance(x, Bencached):

            self._append_string(x.bencoded)

        else:

            raise TypeError('unsupported type')

def bencode(x):

    """Encode Python object x to string"""

    encoder = Encoder()

    encoder.process(x)

    return str(encoder)