summaryrefslogtreecommitdiff
path: root/lib/bencode/bencode.py
blob: 4acf7885081a4020697c77895e069a85fea40849 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
# The contents of this file are subject to the BitTorrent Open Source License
# Version 1.1 (the License).  You may not copy or use this file, in either
# source code or executable form, except in compliance with the License.  You
# may obtain a copy of the License at http://www.bittorrent.com/license/.
#
# Software distributed under the License is distributed on an AS IS basis,
# WITHOUT WARRANTY OF ANY KIND, either express or implied.  See the License
# for the specific language governing rights and limitations under the
# License.

# Written by Petru Paler
# Modified by Julien (jvoisin) Voisin

'''
    A quick (and also nice) lib to bencode/bdecode torrent files
'''


import types


class BTFailure(Exception):
    '''Custom Exception'''
    pass


class Bencached(object):
    '''Custom type : cached string'''
    __slots__ = ['bencoded']

    def __init__(self, string):
        self.bencoded = string


def decode_int(x, f):
    '''decode an int'''
    f += 1
    newf = x.index('e', f)
    n = int(x[f:newf])
    if x[f] == '-':
        if x[f + 1] == '0':
            raise ValueError
    elif x[f] == '0' and newf != f + 1:
        raise ValueError
    return (n, newf + 1)


def decode_string(x, f):
    '''decode a string'''
    colon = x.index(':', f)
    n = int(x[f:colon])
    if x[f] == '0' and colon != f + 1:
        raise ValueError
    colon += 1
    return (x[colon:colon + n], colon + n)


def decode_list(x, f):
    '''decode a list'''
    result = []
    f += 1
    while x[f] != 'e':
        v, f = DECODE_FUNC[x[f]](x, f)
        result.append(v)
    return (result, f + 1)


def decode_dict(x, f):
    '''decode a dict'''
    result = {}
    f += 1
    while x[f] != 'e':
        k, f = decode_string(x, f)
        result[k], f = DECODE_FUNC[x[f]](x, f)
    return (result, f + 1)


def encode_bool(x, r):
    '''bencode a boolean'''
    if x:
        encode_int(1, r)
    else:
        encode_int(0, r)


def encode_int(x, r):
    '''bencode an integer/float'''
    r.extend(('i', str(x), 'e'))


def encode_list(x, r):
    '''bencode a list/tuple'''
    r.append('l')
    [ENCODE_FUNC[type(item)](item, r) for item in x]
    r.append('e')


def encode_dict(x, result):
    '''bencode a dict'''
    result.append('d')
    ilist = x.items()
    ilist.sort()
    for k, v in ilist:
        result.extend((str(len(k)), ':', k))
        ENCODE_FUNC[type(v)](v, result)
    result.append('e')


DECODE_FUNC = {}
DECODE_FUNC.update(dict([(str(x), decode_string) for x in xrange(9)]))
DECODE_FUNC['l'] = decode_list
DECODE_FUNC['d'] = decode_dict
DECODE_FUNC['i'] = decode_int


ENCODE_FUNC = {}
ENCODE_FUNC[Bencached] = lambda x, r: r.append(x.bencoded)
ENCODE_FUNC[types.IntType] = encode_int
ENCODE_FUNC[types.LongType] = encode_int
ENCODE_FUNC[types.StringType] = lambda x, r: r.extend((str(len(x)), ':', x))
ENCODE_FUNC[types.ListType] = encode_list
ENCODE_FUNC[types.TupleType] = encode_list
ENCODE_FUNC[types.DictType] = encode_dict
ENCODE_FUNC[types.BooleanType] = encode_bool


def bencode(string):
    '''bencode $string'''
    table = []
    ENCODE_FUNC[type(string)](string, table)
    return ''.join(table)


def bdecode(string):
    '''decode $string'''
    try:
        result, lenght = DECODE_FUNC[string[0]](string, 0)
    except (IndexError, KeyError, ValueError):
        raise BTFailure('Not a valid bencoded string')
    if lenght != len(string):
        raise BTFailure('Invalid bencoded value (data after valid prefix)')
    return result