spoon.ber.stream

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