apt_dht_Khashmir/bencode.py

   1
   2 """Functions for bencoding and bdecoding data.
   3
   4 @type decode_func: C{dictionary} of C{function}
   5 @var decode_func: a dictionary of function calls to be made, based on data,
   6     the keys are the first character of the data and the value is the
   7     function to use to decode that data
   8 @type bencached_marker: C{list}
   9 @var bencached_marker: mutable type to ensure class origination
  10 @type encode_func: C{dictionary} of C{function}
  11 @var encode_func: a dictionary of function calls to be made, based on data,
  12     the keys are the type of the data and the value is the
  13     function to use to encode that data
  14 @type BencachedType: C{type}
  15 @var BencachedType: the L{Bencached} type
  16 """
  17
  18 from types import IntType, LongType, StringType, ListType, TupleType, DictType, BooleanType
  19 try:
  20     from types import UnicodeType
  21 except ImportError:
  22     UnicodeType = None
  23 from cStringIO import StringIO
  24
  25 from twisted.python import log
  26
  27 def decode_int(x, f):
  28     """Bdecode an integer.
  29
  30     @type x: C{string}
  31     @param x: the data to decode
  32     @type f: C{int}
  33     @param f: the offset in the data to start at
  34     @rtype: C{int}, C{int}
  35     @return: the bdecoded integer, and the offset to read next
  36     @raise ValueError: if the data is improperly encoded
  37
  38     """
  39
  40     f += 1
  41     newf = x.index('e', f)
  42     try:
  43         n = int(x[f:newf])
  44     except:
  45         n = long(x[f:newf])
  46     if x[f] == '-':
  47         if x[f + 1] == '0':
  48             raise ValueError
  49     elif x[f] == '0' and newf != f+1:
  50         raise ValueError
  51     return (n, newf+1)
  52
  53 def decode_string(x, f):
  54     """Bdecode a string.
  55
  56     @type x: C{string}
  57     @param x: the data to decode
  58     @type f: C{int}
  59     @param f: the offset in the data to start at
  60     @rtype: C{string}, C{int}
  61     @return: the bdecoded string, and the offset to read next
  62     @raise ValueError: if the data is improperly encoded
  63
  64     """
  65
  66     colon = x.index(':', f)
  67     try:
  68         n = int(x[f:colon])
  69     except (OverflowError, ValueError):
  70         n = long(x[f:colon])
  71     if x[f] == '0' and colon != f+1:
  72         raise ValueError
  73     colon += 1
  74     return (x[colon:colon+n], colon+n)
  75
  76 def decode_unicode(x, f):
  77     """Bdecode a unicode string.
  78
  79     @type x: C{string}
  80     @param x: the data to decode
  81     @type f: C{int}
  82     @param f: the offset in the data to start at
  83     @rtype: C{int}, C{int}
  84     @return: the bdecoded unicode string, and the offset to read next
  85
  86     """
  87
  88     s, f = decode_string(x, f+1)
  89     return (s.decode('UTF-8'),f)
  90
  91 def decode_list(x, f):
  92     """Bdecode a list.
  93
  94     @type x: C{string}
  95     @param x: the data to decode
  96     @type f: C{int}
  97     @param f: the offset in the data to start at
  98     @rtype: C{list}, C{int}
  99     @return: the bdecoded list, and the offset to read next
 100
 101     """
 102
 103     r, f = [], f+1
 104     while x[f] != 'e':
 105         v, f = decode_func[x[f]](x, f)
 106         r.append(v)
 107     return (r, f + 1)
 108
 109 def decode_dict(x, f):
 110     """Bdecode a dictionary.
 111
 112     @type x: C{string}
 113     @param x: the data to decode
 114     @type f: C{int}
 115     @param f: the offset in the data to start at
 116     @rtype: C{dictionary}, C{int}
 117     @return: the bdecoded dictionary, and the offset to read next
 118     @raise ValueError: if the data is improperly encoded
 119
 120     """
 121
 122     r, f = {}, f+1
 123     lastkey = None
 124     while x[f] != 'e':
 125         k, f = decode_string(x, f)
 126         if lastkey >= k:
 127             raise ValueError
 128         lastkey = k
 129         r[k], f = decode_func[x[f]](x, f)
 130     return (r, f + 1)
 131
 132 decode_func = {}
 133 decode_func['l'] = decode_list
 134 decode_func['d'] = decode_dict
 135 decode_func['i'] = decode_int
 136 decode_func['0'] = decode_string
 137 decode_func['1'] = decode_string
 138 decode_func['2'] = decode_string
 139 decode_func['3'] = decode_string
 140 decode_func['4'] = decode_string
 141 decode_func['5'] = decode_string
 142 decode_func['6'] = decode_string
 143 decode_func['7'] = decode_string
 144 decode_func['8'] = decode_string
 145 decode_func['9'] = decode_string
 146 decode_func['u'] = decode_unicode
 147
 148 def bdecode(x, sloppy = 0):
 149     """Bdecode a string of data.
 150
 151     @type x: C{string}
 152     @param x: the data to decode
 153     @type sloppy: C{boolean}
 154     @param sloppy: whether to allow errors in the decoding
 155     @rtype: unknown
 156     @return: the bdecoded data
 157     @raise ValueError: if the data is improperly encoded
 158
 159     """
 160
 161     try:
 162         r, l = decode_func[x[0]](x, 0)
 163 #    except (IndexError, KeyError):
 164     except (IndexError, KeyError, ValueError):
 165         log.err()
 166         raise ValueError, "bad bencoded data"
 167     if not sloppy and l != len(x):
 168         raise ValueError, "bad bencoded data"
 169     return r
 170
 171 def test_bdecode():
 172     """A test routine for the bdecoding functions."""
 173     try:
 174         bdecode('0:0:')
 175         assert 0
 176     except ValueError:
 177         pass
 178     try:
 179         bdecode('ie')
 180         assert 0
 181     except ValueError:
 182         pass
 183     try:
 184         bdecode('i341foo382e')
 185         assert 0
 186     except ValueError:
 187         pass
 188     assert bdecode('i4e') == 4L
 189     assert bdecode('i0e') == 0L
 190     assert bdecode('i123456789e') == 123456789L
 191     assert bdecode('i-10e') == -10L
 192     try:
 193         bdecode('i-0e')
 194         assert 0
 195     except ValueError:
 196         pass
 197     try:
 198         bdecode('i123')
 199         assert 0
 200     except ValueError:
 201         pass
 202     try:
 203         bdecode('')
 204         assert 0
 205     except ValueError:
 206         pass
 207     try:
 208         bdecode('i6easd')
 209         assert 0
 210     except ValueError:
 211         pass
 212     try:
 213         bdecode('35208734823ljdahflajhdf')
 214         assert 0
 215     except ValueError:
 216         pass
 217     try:
 218         bdecode('2:abfdjslhfld')
 219         assert 0
 220     except ValueError:
 221         pass
 222     assert bdecode('0:') == ''
 223     assert bdecode('3:abc') == 'abc'
 224     assert bdecode('10:1234567890') == '1234567890'
 225     try:
 226         bdecode('02:xy')
 227         assert 0
 228     except ValueError:
 229         pass
 230     try:
 231         bdecode('l')
 232         assert 0
 233     except ValueError:
 234         pass
 235     assert bdecode('le') == []
 236     try:
 237         bdecode('leanfdldjfh')
 238         assert 0
 239     except ValueError:
 240         pass
 241     assert bdecode('l0:0:0:e') == ['', '', '']
 242     try:
 243         bdecode('relwjhrlewjh')
 244         assert 0
 245     except ValueError:
 246         pass
 247     assert bdecode('li1ei2ei3ee') == [1, 2, 3]
 248     assert bdecode('l3:asd2:xye') == ['asd', 'xy']
 249     assert bdecode('ll5:Alice3:Bobeli2ei3eee') == [['Alice', 'Bob'], [2, 3]]
 250     try:
 251         bdecode('d')
 252         assert 0
 253     except ValueError:
 254         pass
 255     try:
 256         bdecode('defoobar')
 257         assert 0
 258     except ValueError:
 259         pass
 260     assert bdecode('de') == {}
 261     assert bdecode('d3:agei25e4:eyes4:bluee') == {'age': 25, 'eyes': 'blue'}
 262     assert bdecode('d8:spam.mp3d6:author5:Alice6:lengthi100000eee') == {'spam.mp3': {'author': 'Alice', 'length': 100000}}
 263     try:
 264         bdecode('d3:fooe')
 265         assert 0
 266     except ValueError:
 267         pass
 268     try:
 269         bdecode('di1e0:e')
 270         assert 0
 271     except ValueError:
 272         pass
 273     try:
 274         bdecode('d1:b0:1:a0:e')
 275         assert 0
 276     except ValueError:
 277         pass
 278     try:
 279         bdecode('d1:a0:1:a0:e')
 280         assert 0
 281     except ValueError:
 282         pass
 283     try:
 284         bdecode('i03e')
 285         assert 0
 286     except ValueError:
 287         pass
 288     try:
 289         bdecode('l01:ae')
 290         assert 0
 291     except ValueError:
 292         pass
 293     try:
 294         bdecode('9999:x')
 295         assert 0
 296     except ValueError:
 297         pass
 298     try:
 299         bdecode('l0:')
 300         assert 0
 301     except ValueError:
 302         pass
 303     try:
 304         bdecode('d0:0:')
 305         assert 0
 306     except ValueError:
 307         pass
 308     try:
 309         bdecode('d0:')
 310         assert 0
 311     except ValueError:
 312         pass
 313
 314 bencached_marker = []
 315
 316 class Bencached:
 317     """Dummy data structure for storing bencoded data in memory.
 318
 319     @type marker: C{list}
 320     @ivar marker: mutable type to make sure the data was encoded by this class
 321     @type bencoded: C{string}
 322     @ivar bencoded: the bencoded data stored in a string
 323
 324     """
 325
 326     def __init__(self, s):
 327         """
 328
 329         @type s: C{string}
 330         @param s: the new bencoded data to store
 331
 332         """
 333
 334         self.marker = bencached_marker
 335         self.bencoded = s
 336
 337 BencachedType = type(Bencached('')) # insufficient, but good as a filter
 338
 339 def encode_bencached(x,r):
 340     """Bencode L{Bencached} data.
 341
 342     @type x: L{Bencached}
 343     @param x: the data to encode
 344     @type r: C{list}
 345     @param r: the currently bencoded data, to which the bencoding of x
 346         will be appended
 347
 348     """
 349
 350     assert x.marker == bencached_marker
 351     r.append(x.bencoded)
 352
 353 def encode_int(x,r):
 354     """Bencode an integer.
 355
 356     @type x: C{int}
 357     @param x: the data to encode
 358     @type r: C{list}
 359     @param r: the currently bencoded data, to which the bencoding of x
 360         will be appended
 361
 362     """
 363
 364     r.extend(('i',str(x),'e'))
 365
 366 def encode_bool(x,r):
 367     """Bencode a boolean.
 368
 369     @type x: C{boolean}
 370     @param x: the data to encode
 371     @type r: C{list}
 372     @param r: the currently bencoded data, to which the bencoding of x
 373         will be appended
 374
 375     """
 376
 377     encode_int(int(x),r)
 378
 379 def encode_string(x,r):
 380     """Bencode a string.
 381
 382     @type x: C{string}
 383     @param x: the data to encode
 384     @type r: C{list}
 385     @param r: the currently bencoded data, to which the bencoding of x
 386         will be appended
 387
 388     """
 389
 390     r.extend((str(len(x)),':',x))
 391
 392 def encode_unicode(x,r):
 393     """Bencode a unicode string.
 394
 395     @type x: C{unicode}
 396     @param x: the data to encode
 397     @type r: C{list}
 398     @param r: the currently bencoded data, to which the bencoding of x
 399         will be appended
 400
 401     """
 402
 403     #r.append('u')
 404     encode_string(x.encode('UTF-8'),r)
 405
 406 def encode_list(x,r):
 407     """Bencode a list.
 408
 409     @type x: C{list}
 410     @param x: the data to encode
 411     @type r: C{list}
 412     @param r: the currently bencoded data, to which the bencoding of x
 413         will be appended
 414
 415     """
 416
 417     r.append('l')
 418     for e in x:
 419         encode_func[type(e)](e, r)
 420     r.append('e')
 421
 422 def encode_dict(x,r):
 423     """Bencode a dictionary.
 424
 425     @type x: C{dictionary}
 426     @param x: the data to encode
 427     @type r: C{list}
 428     @param r: the currently bencoded data, to which the bencoding of x
 429         will be appended
 430
 431     """
 432
 433     r.append('d')
 434     ilist = x.items()
 435     ilist.sort()
 436     for k,v in ilist:
 437         r.extend((str(len(k)),':',k))
 438         encode_func[type(v)](v, r)
 439     r.append('e')
 440
 441 encode_func = {}
 442 encode_func[BencachedType] = encode_bencached
 443 encode_func[IntType] = encode_int
 444 encode_func[LongType] = encode_int
 445 encode_func[StringType] = encode_string
 446 encode_func[ListType] = encode_list
 447 encode_func[TupleType] = encode_list
 448 encode_func[DictType] = encode_dict
 449 if BooleanType:
 450     encode_func[BooleanType] = encode_bool
 451 if UnicodeType:
 452     encode_func[UnicodeType] = encode_unicode
 453
 454 def bencode(x):
 455     """Bencode some data.
 456
 457     @type x: unknown
 458     @param x: the data to encode
 459     @rtype: string
 460     @return: the bencoded data
 461     @raise ValueError: if the data contains a type that cannot be encoded
 462
 463     """
 464     r = []
 465     try:
 466         encode_func[type(x)](x, r)
 467     except:
 468         log.err()
 469         assert 0
 470     return ''.join(r)
 471
 472 def test_bencode():
 473     """A test routine for the bencoding functions."""
 474     assert bencode(4) == 'i4e'
 475     assert bencode(0) == 'i0e'
 476     assert bencode(-10) == 'i-10e'
 477     assert bencode(12345678901234567890L) == 'i12345678901234567890e'
 478     assert bencode('') == '0:'
 479     assert bencode('abc') == '3:abc'
 480     assert bencode('1234567890') == '10:1234567890'
 481     assert bencode([]) == 'le'
 482     assert bencode([1, 2, 3]) == 'li1ei2ei3ee'
 483     assert bencode([['Alice', 'Bob'], [2, 3]]) == 'll5:Alice3:Bobeli2ei3eee'
 484     assert bencode({}) == 'de'
 485     assert bencode({'age': 25, 'eyes': 'blue'}) == 'd3:agei25e4:eyes4:bluee'
 486     assert bencode({'spam.mp3': {'author': 'Alice', 'length': 100000}}) == 'd8:spam.mp3d6:author5:Alice6:lengthi100000eee'
 487     try:
 488         bencode({1: 'foo'})
 489         assert 0
 490     except AssertionError:
 491         pass
 492
 493
 494 try:
 495     import psyco
 496     psyco.bind(bdecode)
 497     psyco.bind(bencode)
 498 except ImportError:
 499     pass