instrumentation/next-share/BaseLib/Core/DecentralizedTracking/kadtracker/bencode.py

   1 # Copyright (C) 2009 Raul Jimenez
   2 # Released under GNU LGPL 2.1
   3 # See LICENSE.txt for more information
   4
   5 import cStringIO
   6 import logging
   7
   8 logger = logging.getLogger('dht')
   9
  10 class LoggingException(Exception):
  11
  12     def __init__(self, msg):
  13         logger.info('%s: %s' % (self.__class__, msg))
  14
  15
  16 class EncodeError(LoggingException):
  17     """Raised by encoder when invalid input."""
  18
  19 class DecodeError(LoggingException):
  20     """Raised by decoder when invalid bencode input."""
  21     def __init__(self, msg, bencoded):
  22         LoggingException.__init__(self, '\nBencoded: '.join((msg, bencoded)))
  23
  24 class RecursionDepthError(DecodeError):
  25     """Raised when the bencoded recursivity is too deep.
  26
  27     This check prevents from using too much recursivity when an
  28     accidentally/maliciously constructed bencoded string looks like
  29     'llllllllllllllllllllllllllllllllllll'.
  30
  31     """
  32
  33
  34 def encode(data):
  35     output = cStringIO.StringIO()
  36     encode_f = _get_encode_f(data)
  37     encode_f(data, output)
  38     result = output.getvalue()
  39     output.close()
  40     return result
  41
  42 def decode(bencoded, max_depth=4):
  43     if not bencoded:
  44         raise DecodeError('Empty bencoded string', bencoded)
  45     try:
  46         decode_f = _get_decode_f(bencoded, 0)
  47         data, next_pos, = decode_f(bencoded, 0, max_depth)
  48     except (KeyError, IndexError, ValueError):
  49         raise DecodeError('UNEXPECTED>>>>>>>>>>>>', bencoded)
  50     else:
  51         if next_pos != len(bencoded):
  52             raise DecodeError('Extra characters after valid bencode.', bencoded)
  53     return data
  54
  55
  56 def _encode_str(data, output):
  57     """Encode a string object
  58
  59     The result format is:
  60     <string length encoded in base ten ASCII>:<string data>
  61
  62     """
  63     output.write('%d:%s' % (len(data), data))
  64
  65 def _encode_int(data, output):
  66     """Encode an integer (or long) object
  67
  68     The result format is:
  69     i<integer encoded in base ten ASCII>e
  70
  71     """
  72     output.write('i%de' % data)
  73
  74 def _encode_list(data, output):
  75     """Encode a list object
  76
  77     The result format is:
  78     l<bencoded element>...<bencoded element>e
  79
  80     """
  81     output.write('l')
  82     for item in data:
  83         encode_f = _get_encode_f(item)
  84         encode_f(item, output)
  85     output.write('e')
  86
  87 def _encode_dict(data, output):
  88     """Encode a dict object
  89
  90     The result format is:
  91     d<bencoded key><bencoded value>...<bencoded key><bencoded value>e
  92     Keys must be string and will be encoded in lexicographical order
  93
  94     """
  95     output.write('d')
  96     keys = data.keys()
  97     keys.sort()
  98     for key in keys:
  99         if type(key) != str: # key must be a string)
 100             raise EncodeError, 'Found a non-string key. Data: %r' % data
 101         value = data[key]
 102         _encode_fs[str](key, output)
 103         encode_f = _get_encode_f(value)
 104         encode_f(value, output)
 105     output.write('e')
 106
 107
 108
 109
 110 def _decode_str(bencoded, pos, _):
 111     """
 112
 113
 114     """
 115     str_len, str_begin = _get_int(bencoded, pos, ':')
 116     str_end = str_begin + str_len
 117     return (bencoded[str_begin:str_end], str_end)
 118
 119 def _decode_int(bencoded, pos, _):
 120     """
 121
 122
 123     """
 124     return  _get_int(bencoded, pos + 1, 'e') # +1 to skip 'i'
 125
 126 def _decode_list(bencoded, pos, max_depth):
 127     """
 128
 129
 130     """
 131     if max_depth == 0:
 132         raise RecursionDepthError('maximum recursion depth exceeded', bencoded)
 133
 134     result = []
 135     next_pos = pos + 1 # skip 'l'
 136     while bencoded[next_pos] != 'e':
 137         decode_f = _get_decode_f(bencoded, next_pos)
 138         item, next_pos = decode_f(bencoded,
 139                                   next_pos, max_depth - 1)
 140         result.append(item)
 141     return result, next_pos + 1 # correct for 'e'
 142
 143 def _decode_dict(bencoded, pos, max_depth):
 144     """
 145
 146     """
 147     if max_depth == 0:
 148         raise RecursionDepthError, 'maximum recursion depth exceeded'
 149
 150     result = {}
 151     next_pos = pos + 1 # skip 'd'
 152     while bencoded[next_pos] != 'e':
 153         # Decode key
 154         decode_f = _get_decode_f(bencoded, next_pos)
 155         if decode_f != _decode_str:
 156             raise DecodeError('Keys must be string. Found: <%s>' % (
 157                     bencoded[next_pos]),
 158                               bencoded)
 159         key, next_pos = decode_f(bencoded,
 160                                  next_pos, max_depth - 1)
 161         # Decode value
 162         decode_f = _get_decode_f(bencoded, next_pos)
 163         value, next_pos = decode_f(bencoded,
 164                                    next_pos, max_depth - 1)
 165         result[key] = value
 166     return result, next_pos + 1 # skip 'e'
 167
 168
 169
 170 def _get_encode_f(value):
 171     try:
 172         return _encode_fs[type(value)]
 173     except (KeyError), e:
 174         raise EncodeError, 'Invalid type: <%r>' % e
 175
 176 def _get_int(bencoded, pos, char):
 177     try:
 178         end = bencoded.index(char, pos)
 179     except (ValueError):
 180         raise DecodeError('Character %s not found.', bencoded)
 181     try:
 182         result = int(bencoded[pos:end])
 183     except (ValueError), e:
 184         raise DecodeError('Not an integer: %r' %e, bencoded)
 185     return (result, end + 1) # +1 to skip final character
 186
 187 def _get_decode_f(bencoded, pos):
 188     try:
 189         return _decode_fs[bencoded[pos]]
 190     except (KeyError), e:
 191         raise DecodeError('Caracter in position %d raised %r' % (pos, e),
 192                           bencoded)
 193
 194
 195 _encode_fs = {str : _encode_str,
 196               int :  _encode_int,
 197               long : _encode_int,
 198               tuple : _encode_list,
 199               list :  _encode_list,
 200               dict : _encode_dict
 201               }
 202
 203 _decode_fs = {'i' : _decode_int,
 204              'l' : _decode_list,
 205              'd' : _decode_dict}
 206 for i in xrange(10):
 207     _decode_fs[str(i)] = _decode_str
 208
 209
 210