#
# python-ipfix (c) 2013 Brian Trammell.
#
# Many thanks to the mPlane consortium (http://www.ict-mplane.eu) for
# its material support of this effort.
#
# This program is free software: you can redistribute it and/or modify it under
# the terms of the GNU Lesser General Public License as published by the Free
# Software Foundation, either version 3 of the License, or (at your option) any
# later version.
#
# This program is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public License for more
# details.
#
# You should have received a copy of the GNU General Public License along with
# this program.  If not, see <http://www.gnu.org/licenses/>.
#

"""
Interface to read IPFIX Messages from a stream.

"""

from . import message

[docs]class MessageStreamReader(object):
    """
    Reads records from a stream of IPFIX messages.

    Uses an :class:`ipfix.message.MessageBuffer` internally, and continually reads
    messages from the given stream into the buffer, iterating over records,
    until the end of the stream. Use :func:`from_stream` to get an instance.

    Suitable for reading from IPFIX files (see :rfc:`5655`) as well as from
    UDP or TCP sockets with :class:`socketserver.StreamRequestHandler`.
    When opening a stream from a file, use mode='rb'.

    """
    def __init__(self, stream):
        self.stream = stream
        self.msg = message.MessageBuffer()
        self.msgcount = 0

[docs]    def namedict_iterator(self):
        """
        Iterate over all records in the stream, as dicts mapping IE names
        to values.

        :returns: a name dictionary iterator

        """
        try:
            while(True):
                self.msg.read_message(self.stream)
                for name in self.msg.namedict_iterator():
                    yield name
                    self.msgcount += 1
        except EOFError:
            return

[docs]    def tuple_iterator(self, ielist):
        """
        Iterate over all records in the stream containing all the IEs in
        the given ielist. Records are returned as tuples in ielist order.

        :param ielist: an instance of :class:`ipfix.ie.InformationElementList`
                       listing IEs to return as a tuple
        :returns: a tuple iterator for tuples in ielist order
        """
        try:
            while(True):
                self.msg.read_message(self.stream)
                for tuple_ in self.msg.tuple_iterator(ielist):
                    yield tuple_
                    self.msgcount += 1
        except EOFError:
            return

[docs]def from_stream(stream):
    """
    Get a MessageStreamReader for a given stream

    :param stream: stream to read
    :return: a :class:`MessageStreamReader` wrapped around the stream.

    """
    return MessageStreamReader(stream)