Source Code for Module pyrobase.bencode

  1  # -*- coding: utf-8 -*- 
  2  # pylint: disable=too-few-public-methods 
  3  """ Bencode support. 
  4   
  5      Copyright (c) 2009-2017 The PyroScope Project <pyroscope.project@gmail.com> 
  6   
  7      See http://en.wikipedia.org/wiki/Bencode 
  8  """ 
  9  # This program is free software; you can redistribute it and/or modify 
 10  # it under the terms of the GNU General Public License as published by 
 11  # the Free Software Foundation; either version 2 of the License, or 
 12  # (at your option) any later version. 
 13  # 
 14  # This program is distributed in the hope that it will be useful, 
 15  # but WITHOUT ANY WARRANTY; without even the implied warranty of 
 16  # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the 
 17  # GNU General Public License for more details. 
 18  # 
 19  # You should have received a copy of the GNU General Public License along 
 20  # with this program; if not, write to the Free Software Foundation, Inc., 
 21  # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA. 
 22   
 23 -class BencodeError(ValueError): 
 24      """ Error during decoding or encoding. 
 25      """ 
 26   
 27   
 28 -class Decoder(object): 
 29      """ Decode a string or stream to an object. 
 30      """ 
 31   
 32 -    def __init__(self, data, char_encoding=None): 
 33          """ Initialize encoder. 
 34          """ 
 35          self.data = data 
 36          self.offset = 0 
 37          self.char_encoding = char_encoding 
 38   
 39   
 40 -    def decode(self, check_trailer=False): # pylint: disable=I0011,R0912 
 41          """ Decode data in C{self.data} and return deserialized object. 
 42   
 43              @param check_trailer: Raise error if trailing junk is found in data? 
 44              @raise BencodeError: Invalid data. 
 45          """ 
 46          try: 
 47              kind = self.data[self.offset] 
 48          except IndexError: 
 49              raise BencodeError("Unexpected end of data at offset %d/%d" % ( 
 50                  self.offset, len(self.data), 
 51              )) 
 52   
 53          if kind.isdigit(): 
 54              # String 
 55              try: 
 56                  end = self.data.find(':', self.offset) 
 57                  length = int(self.data[self.offset:end], 10) 
 58              except (ValueError, TypeError): 
 59                  raise BencodeError("Bad string length at offset %d (%r...)" % ( 
 60                      self.offset, self.data[self.offset:self.offset+32] 
 61                  )) 
 62   
 63              self.offset = end+length+1 
 64              obj = self.data[end+1:self.offset] 
 65   
 66              if self.char_encoding: 
 67                  try: 
 68                      obj = obj.decode(self.char_encoding) 
 69                  except UnicodeError: 
 70                      # deliver non-decodable string (byte arrays) as-is 
 71                      pass 
 72          elif kind == 'i': 
 73              # Integer 
 74              try: 
 75                  end = self.data.find('e', self.offset+1) 
 76                  obj = int(self.data[self.offset+1:end], 10) 
 77              except (ValueError, TypeError): 
 78                  raise BencodeError("Bad integer at offset %d (%r...)" % ( 
 79                      self.offset, self.data[self.offset:self.offset+32] 
 80                  )) 
 81              self.offset = end+1 
 82          elif kind == 'l': 
 83              # List 
 84              self.offset += 1 
 85              obj = []  # pylint: disable=redefined-variable-type 
 86              while self.data[self.offset:self.offset+1] != 'e': 
 87                  obj.append(self.decode()) 
 88              self.offset += 1 
 89          elif kind == 'd': 
 90              # Dict 
 91              self.offset += 1 
 92              obj = {} 
 93              while self.data[self.offset:self.offset+1] != 'e': 
 94                  key = self.decode() 
 95                  obj[key] = self.decode() 
 96              self.offset += 1 
 97          else: 
 98              raise BencodeError("Format error at offset %d (%r...)" % ( 
 99                  self.offset, self.data[self.offset:self.offset+32] 
100              )) 
101   
102          if check_trailer and self.offset != len(self.data): 
103              raise BencodeError("Trailing data at offset %d (%r...)" % ( 
104                  self.offset, self.data[self.offset:self.offset+32] 
105              )) 
106   
107          return obj 
108   
109   
110 -class Encoder(object): 
111      """ Encode a given object to a string or stream. 
112      """ 
113   
114 -    def __init__(self): 
115          """ Initialize encoder. 
116          """ 
117          self.result = [] 
118   
119   
120 -    def encode(self, obj): 
121          """ Add the given object to the result. 
122          """ 
123          if isinstance(obj, (int, long, bool)): 
124              self.result.append("i%de" % obj) 
125          elif isinstance(obj, basestring): 
126              self.result.extend([str(len(obj)), ':', str(obj)]) 
127          elif hasattr(obj, "__bencode__"): 
128              self.encode(obj.__bencode__()) 
129          elif hasattr(obj, "items"): 
130              # Dictionary 
131              self.result.append('d') 
132              for key, val in sorted(obj.items()): 
133                  key = str(key) 
134                  self.result.extend([str(len(key)), ':', key]) 
135                  self.encode(val) 
136              self.result.append('e') 
137          else: 
138              # Treat as iterable 
139              try: 
140                  items = iter(obj) 
141              except TypeError, exc: 
142                  raise BencodeError("Unsupported non-iterable object %r of type %s (%s)" % ( 
143                      obj, type(obj), exc 
144                  )) 
145              else: 
146                  self.result.append('l') 
147                  for item in items: 
148                      self.encode(item) 
149                  self.result.append('e') 
150   
151          return self.result 
152   
153   
154 -def bdecode(data, char_encoding=None): 
155      """ Decode a string or stream to an object. 
156      """ 
157      return Decoder(data, char_encoding).decode(check_trailer=True) 
158   
159   
160 -def bencode(obj): 
161      """ Encode a given object to a string. 
162      """ 
163      return ''.join(Encoder().encode(obj)) 
164   
165   
166 -def bread(stream): 
167      """ Decode a file or stream to an object. 
168      """ 
169      if hasattr(stream, "read"): 
170          return bdecode(stream.read()) 
171      else: 
172          handle = open(stream, "rb") 
173          try: 
174              return bdecode(handle.read()) 
175          finally: 
176              handle.close() 
177   
178   
179 -def bwrite(stream, obj): 
180      """ Encode a given object to a file or stream. 
181      """ 
182      handle = None 
183      if not hasattr(stream, "write"): 
184          stream = handle = open(stream, "wb") 
185      try: 
186          stream.write(bencode(obj)) 
187      finally: 
188          if handle: 
189              handle.close() 
190