apt_dht_Khashmir/bencode.py

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