#
# python-ipfix (c) 2013-2014 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/>.
#

from . import ie, template, message
from .template import IpfixEncodeError, IpfixDecodeError
from datetime import datetime, timedelta
from ipaddress import ip_address
import base64
import io

_stored_test_message = base64.b64decode(b' ')

_test_strings = ["alfa", "bravo", "charlie", "delta", "echo", "foxtrot", "grüezi"]

def mktest_record(sequence):
    return { 'sourceIPv4Address': ip_address(0x7f000000 + sequence),
             'flowStartMilliseconds': datetime(2009, 2, 20, tzinfo=None) + timedelta(0, sequence / 1000),
             'testString': _test_strings[sequence % len(_test_strings)],
             'octetDeltaCount': sequence % 33,
             'packetDeltaCount': sequence % 27 }

def mktest_template(tid=257):
    ie.use_iana_default()
    ie.for_spec("testString(35566/32766)<string>")
    return template.from_ielist(tid, 
           ie.spec_list(["sourceIPv4Address",
                         "flowStartMilliseconds",
                         "testString",
                         "octetDeltaCount[4]",
                         "packetDeltaCount"]))

def mktest_message(rec_count=128, odid=8304, tid=257):
    """
    Make a test message with a certain number of records. 
    The test message contains a template with numeric, datetime, 
    address, and string Information Elements, using both normal 
    and reduced-length encoding.
    """

    msg = message.MessageBuffer()
    msg.begin_export(odid)
    msg.add_template(mktest_template(tid))
    msg.export_ensure_set(tid)
    msg.set_export_time(datetime(2009, 2, 20, 19, 18, 17, tzinfo=None))

    for seq in range(rec_count):
        msg.export_namedict(mktest_record(seq))

    return msg

def fuzzy_datetime_compare(a, b):
    # FIXME This is a hack. We should find a way to make datetimes export exactly
    return abs((a-b).total_seconds()) < 0.01

def msg_to_python(msg, varname="_stored_test_message"):
    return varname+' = base64.b64decode('+repr(base64.b64encode(msg.to_bytes()))+')'

def mktest_message_python(rec_count=128, odid=8304, tid=257, varname="_stored_test_message"):
    return msg_to_python(mktest_message(rec_count, odid, tid), varname=varname)

def test_stored_message():
    assert(_stored_test_message == mktest_message().to_bytes())

    msg = message.MessageBuffer()
    msg.from_bytes(_stored_test_message)
    for i, rec in enumerate(msg.namedict_iterator()):
        trec = mktest_record(i)
        assert(rec['packetDeltaCount'] == trec['packetDeltaCount'])
        assert(rec['octetDeltaCount'] == trec['octetDeltaCount'])
        assert(rec['sourceIPv4Address'] == trec['sourceIPv4Address'])
        assert(rec['testString'] == trec['testString'])
        assert(fuzzy_datetime_compare(rec['flowStartMilliseconds'], trec['flowStartMilliseconds']))

def test_message_write_internals():
    # make sure EOM works
    try:
        mktest_message(rec_count=65536)
        assert(False)
    except message.EndOfMessage as m:
        pass

def test_message_read_internals():
    msg = mktest_message()

    # test get_export_time
    assert(fuzzy_datetime_compare(msg.get_export_time(), datetime(2009, 2, 20, 19, 18, 17, tzinfo=None)))


def test_message_read_errors():
    short_read_test_message_hdr = bytearray(_stored_test_message)
    short_read_test_message_hdr = short_read_test_message_hdr[0:12]
    msg = message.MessageBuffer()
    try:
        msg.from_bytes(bytes(short_read_test_message_hdr))
        assert(False)
    except IpfixDecodeError as e:
        pass

    short_read_test_message_body = bytearray(_stored_test_message)
    short_read_test_message_body = short_read_test_message_body[0:33]
    msg = message.MessageBuffer()
    try:
        msg.from_bytes(bytes(short_read_test_message_body))
        assert(False)
    except IpfixDecodeError as e:
        pass

    bad_msg_version_test_message = bytearray(_stored_test_message)
    bad_msg_version_test_message[0] = 1
    bad_msg_version_test_message[1] = 2
    msg = message.MessageBuffer()
    try:
        msg.from_bytes(bytes(bad_msg_version_test_message))
        assert(False)
    except IpfixDecodeError as e:
        pass

    bad_msg_length_test_message = bytearray(_stored_test_message)
    bad_msg_length_test_message[2] = 0
    bad_msg_length_test_message[3] = 17
    msg = message.MessageBuffer()
    try:
        msg.from_bytes(bytes(bad_msg_length_test_message))
        assert(False)
    except IpfixDecodeError as e:
        pass

    bad_set_length_test_message_short = bytearray(_stored_test_message)
    bad_set_length_test_message_short[18] = 0
    bad_set_length_test_message_short[19] = 1
    msg = message.MessageBuffer()
    try:
        msg.from_bytes(bytes(bad_set_length_test_message_short))
        assert(False)
    except IpfixDecodeError as e:
        pass

    bad_set_length_test_message_long = bytearray(_stored_test_message)
    bad_set_length_test_message_long[18] = 255
    bad_set_length_test_message_long[19] = 255
    msg = message.MessageBuffer()
    try:
        msg.from_bytes(bytes(bad_set_length_test_message_long))
        assert(False)
    except IpfixDecodeError as e:
        pass



