#######################################################################
# Copyright (C) 2008-2011 by Carnegie Mellon University.
#
# @OPENSOURCE_HEADER_START@
#
# Use of the SILK system and related source code is subject to the terms
# of the following licenses:
#
# GNU Public License (GPL) Rights pursuant to Version 2, June 1991
# Government Purpose License Rights (GPLR) pursuant to DFARS 252.227.7013
#
# NO WARRANTY
#
# ANY INFORMATION, MATERIALS, SERVICES, INTELLECTUAL PROPERTY OR OTHER
# PROPERTY OR RIGHTS GRANTED OR PROVIDED BY CARNEGIE MELLON UNIVERSITY
# PURSUANT TO THIS LICENSE (HEREINAFTER THE "DELIVERABLES") ARE ON AN
# "AS-IS" BASIS. CARNEGIE MELLON UNIVERSITY MAKES NO WARRANTIES OF ANY
# KIND, EITHER EXPRESS OR IMPLIED AS TO ANY MATTER INCLUDING, BUT NOT
# LIMITED TO, WARRANTY OF FITNESS FOR A PARTICULAR PURPOSE,
# MERCHANTABILITY, INFORMATIONAL CONTENT, NONINFRINGEMENT, OR ERROR-FREE
# OPERATION. CARNEGIE MELLON UNIVERSITY SHALL NOT BE LIABLE FOR INDIRECT,
# SPECIAL OR CONSEQUENTIAL DAMAGES, SUCH AS LOSS OF PROFITS OR INABILITY
# TO USE SAID INTELLECTUAL PROPERTY, UNDER THIS LICENSE, REGARDLESS OF
# WHETHER SUCH PARTY WAS AWARE OF THE POSSIBILITY OF SUCH DAMAGES.
# LICENSEE AGREES THAT IT WILL NOT MAKE ANY WARRANTY ON BEHALF OF
# CARNEGIE MELLON UNIVERSITY, EXPRESS OR IMPLIED, TO ANY PERSON
# CONCERNING THE APPLICATION OF OR THE RESULTS TO BE OBTAINED WITH THE
# DELIVERABLES UNDER THIS LICENSE.
#
# Licensee hereby agrees to defend, indemnify, and hold harmless Carnegie
# Mellon University, its trustees, officers, employees, and agents from
# all claims or demands made against them (and any related losses,
# expenses, or attorney's fees) arising out of, or relating to Licensee's
# and/or its sub licensees' negligent use or willful misuse of or
# negligent conduct or willful misconduct regarding the Software,
# facilities, or other rights or assistance granted by Carnegie Mellon
# University under this License, including, but not limited to, any
# claims of product liability, personal injury, death, damage to
# property, or violation of any laws or regulations.
#
# Carnegie Mellon University Software Engineering Institute authored
# documents are sponsored by the U.S. Department of Defense under
# Contract FA8721-05-C-0003. Carnegie Mellon University retains
# copyrights in all material produced under this contract. The U.S.
# Government retains a non-exclusive, royalty-free license to publish or
# reproduce these documents, or allow others to do so, for U.S.
# Government purposes only pursuant to the copyright license under the
# contract clause at 252.227.7013.
#
# @OPENSOURCE_HEADER_END@
#
#######################################################################

#######################################################################
# $SiLK: plugin.py 16752 2011-02-24 17:41:51Z mthomas $
#######################################################################

##########################################################################
# Plugin-related code

import warnings

__all__ = ['register_plugin_field', 'register_field', 'register_filter',
           'register_ipv4_field', 'register_ip_field', 'register_int_field',
           'register_enum_field', 'register_int_sum_aggregator',
           'register_int_min_aggregator', 'register_int_max_aggregator',
           'register_switch']

# Old to new name mappings for backwards compatibility
_deprecated_names = {
    'field_len':  'column_width',
    'text_len':   'column_width',
    'bin_len':    'bin_bytes'}

def _check_type(names, kwds):
    for (kwd, val) in kwds.iteritems():
        if kwd not in names:
            raise ValueError("Illegal argument name '%s'" % kwd)
        name_spec = names[kwd]
        if isinstance(name_spec, type):
            if not isinstance(val, name_spec):
                raise TypeError("Argument '%s' must be of type %s" %
                                (kwd, name_spec))
        elif not callable(val):
            raise TypeError("Argument '%s' must be a function of %d arguments"
                            % (kwd, name_spec))


def _register_fieldlike(name, data, names, kwds, deprecated = False):
    if not isinstance(name, str):
        raise TypeError("Name must be a string")

    for val in data:
        if val['name'] == name:
            raise ValueError("A field '%s' has already been registered" % name)

    if deprecated:
        # For backwards compatability
        for (oldname, newname) in _deprecated_names.iteritems():
            if oldname in kwds and newname not in kwds:
                kwds[newname] = kwds[oldname]
                del kwds[oldname]

    _check_type(names, kwds)

    if 'rec_to_bin' in kwds and 'add_rec_to_bin' in kwds:
        raise ValueError("Cannot have both rec_to_bin and add_rec_to_bin")

    kwds['name'] = name
    data.append(kwds)


# The order for these fields must be the same as the order of the
# field_index_t in silkpython.c
_plugin_name_list = ['name', 'description', 'initialize', 'column_width',
                     'rec_to_text', 'bin_bytes', 'rec_to_bin',
                     'bin_to_text', 'add_rec_to_bin', 'bin_merge',
                     'bin_compare', 'initial_value']

def _get_generic_data(data, attr_list):
    retval = []
    for item in data:
        if '__used__' not in item:
            retval.append(tuple([item.get(x, None) for x in attr_list]))
            item['__used__'] = True
    return retval

_plugin_data = []
_plugin_names = {'description': str,
                 'initialize' : 0,
                 'column_width': int,
                 'rec_to_text': 1,
                 'bin_bytes': int,
                 'rec_to_bin': 1,
                 'bin_to_text': 1,
                 'add_rec_to_bin': 2,
                 'bin_merge' : 2,
                 'bin_compare' : 2,
                 'initial_value' : str}

def _get_field_data():
    return _get_generic_data(_plugin_data, _plugin_name_list)

def register_plugin_field(name, **kwds):
    deprecated_field_warning = ("register_plugin_field is deprecated "
                                "in favor of register_field")
    warnings.warn(deprecated_field_warning, DeprecationWarning)
    return _register_fieldlike(name, _plugin_data,
                               _plugin_names, kwds, deprecated=True)

def register_field(name, **kwds):
    return _register_fieldlike(name, _plugin_data, _plugin_names, kwds)

# The order for these fields must be the same as the order of the
# filter_index_t in silkpython.c
_filter_name_list = ['filter', 'initialize', 'finalize']

_filter_data = []
_filter_names = {'filter': 1,
                 'initialize' : 0,
                 'finalize' : 0}

def _get_filter_data():
    return _get_generic_data(_filter_data, _filter_name_list)

def register_filter(filter, **kwds):
    kwds['filter'] = filter
    _check_type(_filter_names, kwds)
    _filter_data.append(kwds)

# The order for these fields must be the same as the order of the
# switch_index_t in silkpython.c
_cmd_line_name_list = ['name', 'handler', 'arg', 'help']

_cmd_line_args = []
_cmd_line_names = {'name': str,
                   'handler': 1,
                   'arg': bool,
                   'help': str}

def _get_cmd_line_args():
    return _get_generic_data(_cmd_line_args, _cmd_line_name_list);

def register_switch(name, **kwds):
    kwds['name'] = name
    kwds.setdefault('arg', True)
    kwds.setdefault('help', 'No help for this switch')
    _check_type(_cmd_line_names, kwds)
    _cmd_line_args.append(kwds)

import os
import sys
import traceback

_programname = None

def _no_traceback_exception(t, value, tb):
    global _programname
    t = str(t).split('.')[-1]
    if _programname:
        printname = _programname
    else:
        try:
            printname = sys.argv[0]
        except:
            printname = "python"
    sys.stderr.write(printname + ": PySiLK plugin: " + t)
    if value and str(value):
        sys.stderr.write(": " + str(value))
    sys.stderr.write("\n")
    if os.getenv("SILK_PYTHON_TRACEBACK"):
        traceback.print_tb(tb)

def _init_silkpython_plugin(name):
    global _programname
    _programname = name
    sys.excepthook = _no_traceback_exception


## Register field convenience functions

import struct
import silk
struct_fmts = ['!B', '!H', '!L', '!Q']
struct_map = {}

class convert(object):
    def __init__(self, fmt):
        self.fmt = fmt
    def pack(self, num):
        return struct.pack(self.fmt, num)
    def unpack(self, s):
        return struct.unpack(self.fmt, s)[0]

class convert_from_larger(object):
    def __init__(self, fmt, size):
        self.fmt = fmt
        self.size = -size
        self.pad = '\0' * (struct.calcsize(fmt) - size)
    def pack(self, num):
        return struct.pack(self.fmt, num)[self.size:]
    def unpack(self, s):
        return struct.unpack(self.fmt, self.pad + s)[0]

class convert_128(object):
    def __init__(self, conv64):
        self.fmt = '!2' + conv64.fmt[-1]
    def pack(self, num):
        return struct.pack(self.fmt, num >> 64, num & 0xffffffffffffffff)
    def unpack(self, s):
        (a, b) = struct.unpack(self.fmt, s)
        return (a << 64) | b

def setup_struct_map():
    current = 1
    for fmt in struct_fmts:
        size = struct.calcsize(fmt)
        struct_map[size] = convert(fmt)
        for i in xrange(current, size):
            struct_map[i] = convert_from_larger(fmt, i)
        current = size + 1
    if 16 not in struct_map and isinstance(struct_map.get(8), convert):
        struct_map[16] = convert_128(struct_map[8])

setup_struct_map()

def Int2String(num, nbytes):
    bytes = ['\0']*nbytes
    n, count = num, (nbytes - 1)
    while n:
        bytes[count] = chr(n & 0xff)
        n >>= 8
        count -= 1
    str = ''.join(bytes)
    return str

def int_to_bin(num, nbytes):
    fn =  struct_map.get(nbytes)
    if fn:
        return fn.pack(num)
    return Int2String(num, nbytes)

def String2Int(str):
    num = reduce(lambda x, y: (x << 8) | y, map(ord, str))
    return num

def bin_to_int(str, nbytes):
    fn = struct_map.get(nbytes)
    if fn:
        return fn.unpack(str)
    return String2Int(str)

def bytes_for_val(num):
    n = num
    count = 0
    while n:
        n >>= 8
        count += 1
    return count

def register_ipv4_field(name, ipv4_function, width=15):
    def call_fn(rec):
        res = ipv4_function(rec)
        if not isinstance(res, silk.IPv4Addr):
            err = "%s expects only IPv4 addresses" % name
            raise TypeError, err
        return res
    def rec_to_bin(rec):
        return int_to_bin(int(call_fn(rec)), 4)
    def rec_to_text(rec):
        return str(call_fn(rec))
    def bin_to_text(bin):
        return str(silk.IPv4Addr(bin_to_int(bin, 4)))
    register_field(name,
                   rec_to_text = rec_to_text,
                   rec_to_bin = rec_to_bin,
                   bin_to_text = bin_to_text,
                   column_width = width,
                   bin_bytes = 4)

def _register_ip_field(name, ip_function, width=39):
    def rec_to_bin(rec):
        ip = ip_function(rec).to_ipv6()
        return int_to_bin(int(ip_function(rec)), 16)
    def rec_to_text(rec):
        return str(ip_function(rec))
    def bin_to_text(bin):
        ipv6 = silk.IPv6Addr(bin_to_int(bin, 16))
        ip = ipv6.to_ipv4()
        ip = ip or ipv6
        return str(ip)
    register_field(name,
                   rec_to_text = rec_to_text,
                   rec_to_bin = rec_to_bin,
                   bin_to_text = bin_to_text,
                   column_width = width,
                   bin_bytes = 16)

if silk.ipv6_enabled():
    register_ip_field = _register_ip_field
else:
    register_ip_field = register_ipv4_field

class int_conv(object):
    def __init__(self, fn, min, max):
        range = max - min + 1
        self.min = min
        self.bytes = bytes_for_val(range)
        self.fn = fn
    def rec_to_bin(self, rec):
        return int_to_bin(self.fn(rec) - self.min, self.bytes)
    def rec_to_text(self, rec):
        return str(self.fn(rec))
    def bin_to_text(self, bin):
        return str(bin_to_int(bin, self.bytes))


def register_int_field(name, int_function, min, max, width=None):
    if min > max:
        raise ValueError, "minimum value must be less than the maximum value"
    conv = int_conv(int_function, min, max)
    if width is None:
        width = len(str(max))
    register_field(name,
                   rec_to_text = conv.rec_to_text,
                   rec_to_bin = conv.rec_to_bin,
                   bin_to_text = conv.bin_to_text,
                   column_width = width,
                   bin_bytes = conv.bytes)


class enum_conv(object):
    def __init__(self, fn, ordering):
        self.fn = fn
        self.next = 0
        self.map = {}
        self.list = []
        for item in ordering:
            self.map[item] = self.next
            self.next += 1
            self.list.append(item)

    def rec_to_bin(self, rec):
        key = self.fn(rec)
        i = self.map.get(key)
        if i is None:
            if self.next > 0xffffffff:
                raise RuntimeError, "Too many items"
            i = self.next
            self.next += 1
            self.map[key] = i
            self.list.append(key)
        return int_to_bin(i, 4)

    def rec_to_text(self, rec):
        return str(self.fn(rec))

    def bin_to_text(self, bin):
        return str(self.list[bin_to_int(bin, 4)])


def register_enum_field(name, enum_function, width, ordering=[]):
    conv = enum_conv(enum_function, ordering)
    register_field(name,
                   rec_to_text = conv.rec_to_text,
                   rec_to_bin = conv.rec_to_bin,
                   bin_to_text = conv.bin_to_text,
                   column_width = width,
                   bin_bytes = 4)


class int_agg_conv(object):
    def __init__(self, fn, max, name="Aggregate field", agg=None):
        self.bytes = bytes_for_val(max)
        self.fn = fn
        self.maxval = max
        self.name = name
        if agg is None:
            self.agg = self.sum
        else:
            self.agg = agg

    def add_rec_to_bin(self, rec, bin):
        newval = self.agg(bin_to_int(bin, self.bytes), self.fn(rec))
        return int_to_bin(newval, self.bytes)

    def bin_compare(self, b1, b2):
        return cmp(bin_to_int(b1, self.bytes), bin_to_int(b2, self.bytes))

    def bin_merge(self, b1, b2):
        newval = self.agg(bin_to_int(b1, self.bytes),
                          bin_to_int(b2, self.bytes))
        return int_to_bin(newval, self.bytes)

    def bin_to_text(self, bin):
        return str(bin_to_int(bin, self.bytes))

    def bin_initial(self, max=None):
        if max:
            val = self.maxval
        else:
            val = 0
        return int_to_bin(val, self.bytes)

    def sum(self, a, b):
        rv = a + b
        if rv > self.maxval:
            raise OverflowError, ("%s overflowed" % self.name)
        return rv


def register_int_aggregator(name, int_function,
                            max, width, agg=None, init=None):
    conv = int_agg_conv(int_function, max, name=name, agg=agg)
    bytes = bytes_for_val(max)
    if width is None:
        width = len(str(conv.maxval))
    if init is None:
        init = conv.bin_initial()
    elif init == "__maxval__":
        init = conv.bin_initial(max=True)
    elif len(init) != bytes:
        raise ValueError, "Bad initial value"
    register_field(name,
                   bin_to_text = conv.bin_to_text,
                   add_rec_to_bin = conv.add_rec_to_bin,
                   bin_compare = conv.bin_compare,
                   bin_merge = conv.bin_merge,
                   initial_value = init,
                   column_width = width,
                   bin_bytes = bytes)


max64 = (1 << 64) - 1
maxfn = max

def register_int_sum_aggregator(name, int_function, max=max64, width=None):
    register_int_aggregator(name, int_function, max, width)

def register_int_max_aggregator(name, int_function, max=max64, width=None):
    register_int_aggregator(name, int_function, max, width, maxfn)

def register_int_min_aggregator(name, int_function, max=max64, width=None):
    register_int_aggregator(name, int_function, max, width, min,
                            init="__maxval__")
