Package spoon :: Package ber :: Module stream
[hide private]
[frames] | no frames]

Source Code for Module spoon.ber.stream

  1  # 
  2  # Copyright (C) 2003-2006, Robey Pointer <robey@lag.net> 
  3  # 
  4  # Permission is hereby granted, free of charge, to any person obtaining 
  5  # a copy of this software and associated documentation files (the 
  6  # "Software"), to deal in the Software without restriction, including 
  7  # without limitation the rights to use, copy, modify, merge, publish, 
  8  # distribute, sublicense, and/or sell copies of the Software, and to 
  9  # permit persons to whom the Software is furnished to do so, subject to 
 10  # the following conditions: 
 11  #  
 12  # The above copyright notice and this permission notice shall be included 
 13  # in all copies or substantial portions of the Software. 
 14  #  
 15  # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS 
 16  # OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 
 17  # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. 
 18  # IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY 
 19  # CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, 
 20  # TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE 
 21  # SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 
 22  # 
 23   
 24  from cStringIO import StringIO 
 25   
 26  from common import BERException, deflate_long, inflate_long 
 27  from common import UNIVERSAL, APPLICATION, CONTEXT, PRIVATE 
 28  from tag import Tag 
 29   
 30   
 31  UNIVERSAL_BOOL = 1 
 32  UNIVERSAL_INT = 2 
 33  UNIVERSAL_BYTES = 4 
 34  UNIVERSAL_NULL = 5 
 35  UNIVERSAL_UTF8 = 12 
 36  UNIVERSAL_LIST = 16 
 37   
 38  EOF_TYPE = Tag(UNIVERSAL, 0, 0) 
 39  NULL_TYPE = Tag(UNIVERSAL, UNIVERSAL_NULL) 
 40  BOOL_TYPE = Tag(UNIVERSAL, UNIVERSAL_BOOL) 
 41  INT_TYPE = Tag(UNIVERSAL, UNIVERSAL_INT) 
 42  BYTES_TYPE = Tag(UNIVERSAL, UNIVERSAL_BYTES) 
 43  UTF8_TYPE = Tag(UNIVERSAL, UNIVERSAL_UTF8) 
 44  LIST_TYPE = Tag(UNIVERSAL, UNIVERSAL_LIST, container=True) 
 45   
 46   
47 -class CountingFile (object):
48 - def __init__(self, fd):
49 self._fd = fd 50 self.count = 0
51
52 - def read(self, n):
53 out = self._fd.read(n) 54 self.count += len(out) 55 return out
56
57 - def __getattr__(self, name):
58 return getattr(self._fd, name)
59 60
61 -class BERStream (object):
62 """ 63 FIXME docs 64 """ 65 66 _encoder_table = {} 67 _zencoder_table = {} 68 _decoder_table = {} 69
70 - def __init__(self, fd, size=None):
71 self._fd = fd 72 self._bytes_read = 0 73 self._size = size 74 self._hit_eof = False 75 self._advance_tag = None
76
77 - def _next_tag(self):
78 tag = Tag.from_stream(self._fd) 79 if tag is None: 80 raise BERException('Stream overrun') 81 self._bytes_read += len(tag) 82 if (self._size is not None) and (self._bytes_read > self._size): 83 raise BERException('Stream overrun') 84 return tag
85
86 - def has_next(self):
87 """ 88 Return C{True} if there is still more data left in this stream, and 89 a future call to L{next} should succeed. 90 91 @return: C{True} if the stream has more data to read 92 """ 93 if self._size is not None: 94 return self._bytes_read < self._size 95 if self._hit_eof: 96 return False 97 98 # indefinite length -- read the next tag 99 if self._advance_tag is not None: 100 # i've already read the next tag :) 101 return not self._hit_eof 102 self._advance_tag = self._next_tag() 103 if self._advance_tag == EOF_TYPE: 104 self._hit_eof = True 105 return False 106 return True
107
108 - def next(self):
109 """ 110 Return the next item from this stream. Objects are decoded using the 111 codecs registered via the L{decoder} operator, though simple types 112 (None, bool, int, long, str, unicode, list, and tuple) have default 113 decoders. 114 115 @return: the next object from the stream 116 @rtype: object 117 @raise BERException: if you're already at the end of the stream 118 """ 119 120 if not self.has_next(): 121 # end of stream 122 raise BERException('End of stream') 123 124 if self._advance_tag is not None: 125 # indefinite length mode 126 tag = self._advance_tag 127 self._advance_tag = None 128 else: 129 tag = self._next_tag() 130 131 decoder = self._decoder_table.get(tag, None) 132 if decoder is None: 133 raise BERException('Can\'t decode object of type %r' % (tag,)) 134 135 # insert a counting fd so we can track how many bytes were actually read 136 count_fd = CountingFile(self._fd) 137 obj = decoder(count_fd, tag) 138 if tag.size is not None: 139 if count_fd.count > tag.size: 140 raise BERException('Overrun in decoder for type %r' % (tag,)) 141 if count_fd.count < tag.size: 142 raise BERException('Underrun in decoder for type %r' % (tag,)) 143 self._bytes_read += tag.size 144 else: 145 # that was an indefinite-length sequence: trust that the correct number of bytes were read 146 self._bytes_read += count_fd.count 147 return obj
148
149 - def add(self, item, compress=False):
150 """ 151 Write an object into the stream. Simple types (None, bool, int, long, 152 str, unicode, list, and tuple) are handled by default encoders. Other 153 encoders may be added with the L{encoder} decorator. 154 155 @param item: object to add 156 @type item: object 157 """ 158 encoder = BERStream._get_encoder(item, compress) 159 if encoder is None: 160 raise BERException('Can\'t encode object of type %r' % (type(item),)) 161 encoder(self._fd, item)
162
163 - def _add_eof(self):
164 EOF_TYPE.write(self._fd)
165 166 @staticmethod
167 - def _get_encoder(item, compress=False):
168 # walk the mro, so a codec can be defined for a superclass 169 for cls in type(item).__mro__: 170 if compress: 171 encoder = BERStream._zencoder_table.get(cls, None) 172 if encoder is None: 173 # Don't compress if we can't find something to compress with 174 encoder = BERStream._encoder_table.get(cls, None) 175 else: 176 encoder = BERStream._encoder_table.get(cls, None) 177 if encoder is not None: 178 return encoder 179 return None
180 181 @staticmethod
182 - def can_encode(item):
183 return (BERStream._get_encoder(item) is not None)
184 -class zencoder (object):
185 # decorator used to mark functions for encoding compressed types 186 # compressed functions should use different types than their non 187 # compressed counterparts, and thus a zdecoder isn't necessary.
188 - def __init__(self, *types):
189 self._types = types
190
191 - def __call__(self, f):
192 for t in self._types: 193 BERStream._zencoder_table[t] = f 194 return f
195
196 -class encoder (object):
197 # decorator 198
199 - def __init__(self, *types):
200 self._types = types
201
202 - def __call__(self, f):
203 for t in self._types: 204 BERStream._encoder_table[t] = f 205 return f
206 207
208 -class decoder (object):
209 # decorator 210
211 - def __init__(self, ttype):
212 self._type = ttype
213
214 - def __call__(self, f):
215 BERStream._decoder_table[self._type] = f 216 return f
217 218
219 -def encode_container(fd, tag, items):
220 """ 221 Encode a list of items into a container with the given tag and write it 222 to a stream. The list is written using indefinite-length encoding, so 223 no extra copying occurs. 224 225 @param fd: the file object to write into 226 @type fd: file 227 @param tag: the tag to use for this list 228 @type tag: L{Tag} 229 @param items: a list of items to put into the container 230 @type items: list or iterable 231 """ 232 Tag.from_tag(tag, None).write(fd) 233 b = BERStream(fd) 234 for x in items: 235 b.add(x) 236 b._add_eof()
237 238
239 -def decode_container(fd, tag):
240 """ 241 Decode a container into a list of items. 242 243 @param fd: the file object to read from 244 @type fd: file 245 @param tag: the tag from the container 246 @type tag: L{Tag} 247 @return: a list of decoded objects 248 @rtype: list 249 """ 250 b = BERStream(fd, tag.size) 251 out = [] 252 while b.has_next(): 253 out.append(b.next()) 254 return out
255 256 257 @encoder(type(None))
258 -def encode_null(fd, item):
259 Tag.from_tag(NULL_TYPE, 0).write(fd)
260 261 @decoder(NULL_TYPE)
262 -def decode_null(fd, tag):
263 return None
264 265 @encoder(bool)
266 -def encode_bool(fd, item):
267 Tag.from_tag(BOOL_TYPE, 1).write(fd) 268 if item: 269 fd.write('\xff') 270 else: 271 fd.write('\x00')
272 273 274 @decoder(BOOL_TYPE)
275 -def decode_bool(fd, tag):
276 if tag.size != 1: 277 raise BERException('unexpected size of boolean (%d)' % (tag.size,)) 278 data = fd.read(1) 279 if len(data) < 1: 280 raise BERException('abrupt end of stream') 281 if ord(data) == 0: 282 return False 283 return True
284 285 @encoder(int, long)
286 -def encode_int(fd, item):
287 x = deflate_long(item) 288 Tag.from_tag(INT_TYPE, len(x)).write(fd) 289 fd.write(x)
290 291 @decoder(INT_TYPE)
292 -def decode_int(fd, tag):
293 return inflate_long(fd.read(tag.size))
294 295 @encoder(str)
296 -def encode_str(fd, item):
297 Tag.from_tag(BYTES_TYPE, len(item)).write(fd) 298 fd.write(item)
299 300 @decoder(BYTES_TYPE)
301 -def decode_str(fd, tag):
302 return fd.read(tag.size)
303 304 @encoder(unicode)
305 -def encode_unicode(fd, item):
306 x = item.encode('utf-8') 307 Tag.from_tag(UTF8_TYPE, len(x)).write(fd) 308 fd.write(x)
309 310 @decoder(UTF8_TYPE)
311 -def decode_unicode(fd, tag):
312 return fd.read(tag.size).decode('utf-8')
313 314 @encoder(list, tuple)
315 -def encode_list(fd, items):
316 encode_container(fd, LIST_TYPE, items)
317 318 @decoder(LIST_TYPE)
319 -def decode_list(fd, tag):
320 return decode_container(fd, tag)
321