bencode.py

   1 """
   2 A library for streaming and unstreaming of simple objects, designed
   3 for speed, compactness, and ease of implementation.
   4
   5 The basic functions are bencode and bdecode. bencode takes an object
   6 and returns a string, bdecode takes a string and returns an object.
   7 bdecode raises a ValueError if you give it an invalid string.
   8
   9 The objects passed in may be nested dicts, lists, ints, strings,
  10 and None. For example, all of the following may be bencoded -
  11
  12 {'a': [0, 1], 'b': None}
  13
  14 [None, ['a', 2, ['c', None]]]
  15
  16 {'spam': (2,3,4)}
  17
  18 {'name': 'Cronus', 'spouse': 'Rhea', 'children': ['Hades', 'Poseidon']}
  19
  20 In general bdecode(bencode(spam)) == spam, but tuples and lists are
  21 encoded the same, so bdecode(bencode((0, 1))) is [0, 1] rather
  22 than (0, 1). Longs and ints are also encoded the same way, so
  23 bdecode(bencode(4)) is a long.
  24
  25 dict keys are required to be strings, to avoid a mess of potential
  26 implementation incompatibilities. bencode is intended to be used
  27 for protocols which are going to be re-implemented many times, so
  28 it's very conservative in that regard.
  29
  30 Which type is encoded is determined by the first character, 'i', 'n',
  31 'd', 'l' and any digit. They indicate integer, null, dict, list, and
  32 string, respectively.
  33
  34 Strings are length-prefixed in base 10, followed by a colon.
  35
  36 bencode('spam') == '4:spam'
  37
  38 Nulls are indicated by a single 'n'.
  39
  40 bencode(None) == 'n'
  41
  42 integers are encoded base 10 and terminated with an 'e'.
  43
  44 bencode(3) == 'i3e'
  45 bencode(-20) == 'i-20e'
  46
  47 Lists are encoded in list order, terminated by an 'e' -
  48
  49 bencode(['abc', 'd']) == 'l3:abc1:de'
  50 bencode([2, 'f']) == 'li2e1:fe'
  51
  52 Dicts are encoded by containing alternating keys and values,
  53 with the keys in sorted order, terminated by an 'e'. For example -
  54
  55 bencode({'spam': 'eggs'}) == 'd4:spam4:eggse'
  56 bencode({'ab': 2, 'a': None}) == 'd1:an2:abi2ee'
  57
  58 Truncated strings come first, so in sort order 'a' comes before 'abc'.
  59
  60 If a function is passed to bencode, it's called and it's return value
  61 is included as a raw string, for example -
  62
  63 bdecode(bencode(lambda: None)) == None
  64 """
  65
  66 # This file is licensed under the GNU Lesser General Public License v2.1.
  67 # originally written for Mojo Nation by Bryce Wilcox, Bram Cohen, and Greg P. Smith
  68 # since then, almost completely rewritten by Bram Cohen
  69
  70 from types import *
  71 from cStringIO import StringIO
  72 import re
  73
  74 def bencode(data):
  75     """
  76     encodes objects as strings, see module documentation for more info
  77     """
  78     result = StringIO()
  79     bwrite(data, result)
  80     return result.getvalue()
  81
  82 def bwrite(data, result):
  83     encoder = encoders.get(type(data))
  84     assert encoder is not None, 'unsupported data type: ' + `type(data)`
  85     encoder(data, result)
  86
  87 encoders = {}
  88
  89 def encode_int(data, result):
  90     result.write('i' + str(data) + 'e')
  91
  92 encoders[IntType] = encode_int
  93 encoders[LongType] = encode_int
  94
  95 def encode_list(data, result):
  96     result.write('l')
  97     for i in data:
  98         bwrite(i, result)
  99     result.write('e')
 100
 101 encoders[TupleType] = encode_list
 102 encoders[ListType] = encode_list
 103
 104 def encode_string(data, result):
 105     result.write(str(len(data)) + ':' + data)
 106
 107 encoders[StringType] = encode_string
 108
 109 def encode_dict(data, result):
 110     result.write('d')
 111     keys = data.keys()
 112     keys.sort()
 113     for key in keys:
 114         assert type(key) is StringType, 'bencoded dictionary key must be a string'
 115         bwrite(key, result)
 116         bwrite(data[key], result)
 117     result.write('e')
 118
 119 encoders[DictType] = encode_dict
 120
 121 encoders[NoneType] = lambda data, result: result.write('n')
 122
 123 encoders[FunctionType] = lambda data, result: result.write(data())
 124 encoders[MethodType] = encoders[FunctionType]
 125
 126 def bdecode(s):
 127     """
 128     Does the opposite of bencode. Raises a ValueError if there's a problem.
 129     """
 130     try:
 131         result, index = bread(s, 0)
 132         if index != len(s):
 133             raise ValueError('left over stuff at end')
 134         return result
 135     except IndexError, e:
 136         raise ValueError(str(e))
 137     except KeyError, e:
 138         raise ValueError(str(e))
 139
 140 def bread(s, index):
 141     return decoders[s[index]](s, index)
 142
 143 decoders = {}
 144
 145 _bre = re.compile(r'(0|[1-9][0-9]*):')
 146
 147 def decode_raw_string(s, index):
 148     x = _bre.match(s, index)
 149     if x is None:
 150         raise ValueError('invalid integer encoding')
 151     endindex = x.end() + long(s[index:x.end() - 1])
 152     if endindex > len(s):
 153         raise ValueError('length encoding indicated premature end of string')
 154     return s[x.end(): endindex], endindex
 155
 156 for c in '0123456789':
 157     decoders[c] = decode_raw_string
 158
 159 _int_re = re.compile(r'i(0|-?[1-9][0-9]*)e')
 160
 161 def decode_int(s, index):
 162     x = _int_re.match(s, index)
 163     if x is None:
 164         raise ValueError('invalid integer encoding')
 165     return long(s[index + 1:x.end() - 1]), x.end()
 166
 167 decoders['i'] = decode_int
 168
 169 decoders['n'] = lambda s, index: (None, index + 1)
 170
 171 def decode_list(s, index):
 172     result = []
 173     index += 1
 174     while s[index] != 'e':
 175         next, index = bread(s, index)
 176         result.append(next)
 177     return result, index + 1
 178
 179 decoders['l'] = decode_list
 180
 181 def decode_dict(s, index):
 182     result = {}
 183     index += 1
 184     prevkey = None
 185     while s[index] != 'e':
 186         key, index = decode_raw_string(s, index)
 187         if key <= prevkey:
 188             raise ValueError("out of order keys")
 189         prevkey = key
 190         value, index = bread(s, index)
 191         result[key] = value
 192     return result, index + 1
 193
 194 decoders['d'] = decode_dict
 195
 196 def test_decode_raw_string():
 197     assert decode_raw_string('1:a', 0) == ('a', 3)
 198     assert decode_raw_string('0:', 0) == ('', 2)
 199     assert decode_raw_string('10:aaaaaaaaaaaaaaaaaaaaaaaaa', 0) == ('aaaaaaaaaa', 13)
 200     assert decode_raw_string('10:', 1) == ('', 3)
 201     try:
 202         decode_raw_string('01:a', 0)
 203         assert 0, 'failed'
 204     except ValueError:
 205         pass
 206     try:
 207         decode_raw_string('--1:a', 0)
 208         assert 0, 'failed'
 209     except ValueError:
 210         pass
 211     try:
 212         decode_raw_string('h', 0)
 213         assert 0, 'failed'
 214     except ValueError:
 215         pass
 216     try:
 217         decode_raw_string('h:', 0)
 218         assert 0, 'failed'
 219     except ValueError:
 220         pass
 221     try:
 222         decode_raw_string('1', 0)
 223         assert 0, 'failed'
 224     except ValueError:
 225         pass
 226     try:
 227         decode_raw_string('', 0)
 228         assert 0, 'failed'
 229     except ValueError:
 230         pass
 231     try:
 232         decode_raw_string('5:a', 0)
 233         assert 0, 'failed'
 234     except ValueError:
 235         pass
 236
 237 def test_dict_enforces_order():
 238     bdecode('d1:an1:bne')
 239     try:
 240         bdecode('d1:bn1:ane')
 241         assert 0, 'failed'
 242     except ValueError:
 243         pass
 244
 245 def test_dict_forbids_non_string_key():
 246     try:
 247         bdecode('di3ene')
 248         assert 0, 'failed'
 249     except ValueError:
 250         pass
 251
 252 def test_dict_forbids_key_repeat():
 253     try:
 254         bdecode('d1:an1:ane')
 255         assert 0, 'failed'
 256     except ValueError:
 257         pass
 258
 259 def test_empty_dict():
 260     assert bdecode('de') == {}
 261
 262 def test_ValueError_in_decode_unknown():
 263     try:
 264         bdecode('x')
 265         assert 0, 'flunked'
 266     except ValueError:
 267         pass
 268
 269 def test_encode_and_decode_none():
 270     assert bdecode(bencode(None)) == None
 271
 272 def test_encode_and_decode_long():
 273     assert bdecode(bencode(-23452422452342L)) == -23452422452342L
 274
 275 def test_encode_and_decode_int():
 276     assert bdecode(bencode(2)) == 2
 277
 278 def test_decode_noncanonical_int():
 279     try:
 280         bdecode('i03e')
 281         assert 0
 282     except ValueError:
 283         pass
 284     try:
 285         bdecode('i3 e')
 286         assert 0
 287     except ValueError:
 288         pass
 289     try:
 290         bdecode('i 3e')
 291         assert 0
 292     except ValueError:
 293         pass
 294     try:
 295         bdecode('i-0e')
 296         assert 0
 297     except ValueError:
 298         pass
 299
 300 def test_encode_and_decode_dict():
 301     x = {'42': 3}
 302     assert bdecode(bencode(x)) == x
 303
 304 def test_encode_and_decode_list():
 305     assert bdecode(bencode([])) == []
 306
 307 def test_encode_and_decode_tuple():
 308     assert bdecode(bencode(())) == []
 309
 310 def test_encode_and_decode_empty_dict():
 311     assert bdecode(bencode({})) == {}
 312
 313 def test_encode_and_decode_complex_object():
 314     spam = [[], 0, -3, -345234523543245234523L, {}, 'spam', None, {'a': [3]}, {}]
 315     assert bencode(bdecode(bencode(spam))) == bencode(spam)
 316     assert bdecode(bencode(spam)) == spam
 317
 318 def test_unfinished_list():
 319     try:
 320         bdecode('ln')
 321         assert 0
 322     except ValueError:
 323         pass
 324
 325 def test_unfinished_dict():
 326     try:
 327         bdecode('d')
 328         assert 0
 329     except ValueError:
 330         pass
 331     try:
 332         bdecode('d1:a')
 333         assert 0
 334     except ValueError:
 335         pass