# Plugin for the prelude IDS, depending on the prelude python bindings.

try:
    from prelude import *
except ImportError, e:
    raise ImportError, 'You need to install libprelude with python bindings'
from gozerbot.generic import rlog
from gozerbot.persistconfig import PersistConfig
from gozerbot import config
import sys, time

cfg = PersistConfig()
cfg.define('manager-host',   '127.0.0.1')
cfg.define('sensor-name',    'gozerbot')
cfg.define('sensor-model',   'gozerbot')
cfg.define('sensor-version', config.ver.replace(' ', '-')) 
cfg.define('sensor-class',   'gozerbot')

# update version
if cfg.get('sensor-version').startswith('GOZERBOT-'):
    cfg.set('sensor-version', config.ver.replace(' ', '-'))

class Prelude(object):
    def __init__(self):
        # initiation
        prelude_init(1, [cfg.get('sensor-name')])
        self.client = prelude_client_new(cfg.get('sensor-name'))
        self.analyzer = prelude_client_get_analyzer(self.client)
        # registration
        idmef_analyzer_set_name(self.analyzer, cfg.get('sensor-name'))
        idmef_analyzer_set_model(self.analyzer, cfg.get('sensor-model'))
        idmef_analyzer_set_version(self.analyzer, cfg.get('sensor-version'))
        idmef_analyzer_set_class(self.analyzer, cfg.get('sensor-class'))
        # startup
        try:
            prelude_client_start(self.client)
        except PreludeError, e:
            if prelude_client_is_setup_needed(e.errno):
                import os
                need_register = 'Error starting client, did you register this sensor? If not, use: prelude-admin register gozerbot "idmef:w" <manager host> --uid %d --gid %d' % \
                    (os.geteuid(), os.getegid())
                rlog(10, 'prelude', need_register)
                raise Exception, need_register
            else:
                raise e
        else:
            prelude_client_set_flags(self.client, PRELUDE_CLIENT_FLAGS_ASYNC_SEND|PRELUDE_CLIENT_FLAGS_ASYNC_TIMER)
            
    def __destroy__(self):
        prelude_client_destroy(self.client, PRELUDE_CLIENT_EXIT_STATUS_SUCCESS)

    def _add_obj(self, m, o, v):
        try: 
            path = idmef_path_new(o)
        except Exception, e:
            rlog(10, 'prelude', "%s: %s" % (o, e))
            return None
        try: 
            val = idmef_value_new_from_path(path, v)
        except Exception, e:
            rlog(10, 'prelude', "could not create value for object %s: %s" % (o, e))
            return None
        try: 
            idmef_path_set(path, m, val)
        except Exception, e:
            rlog(10, 'prelude', "error setting path %s: %s" % (o, e))
            return None
        idmef_value_destroy(val)
        idmef_path_destroy(path)
 
    def report(self, msg, ievent, severity='low'):
        rlog(10, 'prelude', '%s: %s' % (severity.upper(), msg))
        try:
            msg = msg.encode('ascii', 'replace')
        except:
            msg = str(msg)
        try:
            m = idmef_message_new()
            t = idmef_time_new_from_time(long(time.time()))
            a = idmef_message_new_alert(m)
            idmef_alert_set_create_time(a, t)
            self._add_obj(m, "alert.classification.text", msg)
            if ievent:
                self._add_obj(m, "alert.source(0).node.address(0).address", ievent.userhost)
                self._add_obj(m, "alert.target(0).node.address(0).address", ievent.target)
            self._add_obj(m, "alert.assessment.impact.severity", severity)
            idmef_alert_set_analyzer(idmef_message_get_alert(m), idmef_analyzer_ref(self.analyzer), -1)
            self._add_obj(m, "alert.analyzer(0).process.name", sys.argv[0])
            prelude_client_send_idmef(self.client, m)
            idmef_message_destroy(m)
        except Exception, e:
            rlog(10, 'prelude', 'FAILED TO SEND MESSAGE!')
    def info(self, msg, ievent):
        self.report(msg, ievent, 'low')

    def warning(self, msg, ievent):
        self.report(msg, ievent, 'medium')

    def critical(self, msg, ievent):
        self.report(msg, ievent, 'high')
    error = critical

def wrap(obj, func, message='', ievent=None):
    fobj = getattr(obj, func, False)
    if not fobj:
        rlog(10, 'prelude', 'can not wrap %s.%s: no such attribute' % (obj.__class__.__name__, func))
        return
    elif hasattr(obj, '__wrapped__') and obj.__wrapped__.has_key(func):
        rlog(10, 'prelude', 'can not wrap %s.%s: function already wrapped' % (obj.__class__.__name__, func))
        return
    else:
        rlog(10, 'prelude', 'wrapping %s.%s' % (obj.__class__.__name__, func))
        def wrapper(*a, **k):
            result = fobj(*a, **k)
            alerts = message(result, *a, **k)
            if alerts:
                ids.warning(alerts, ievent)
            return result
    # store function in wrapped dict
    if not hasattr(obj, '__wrapped__'):
        obj.__wrapped__ = {}
    obj.__wrapped__[func] = fobj
    # replace function
    obj.__setattr__(func, wrapper)

def unwrap(obj):
    if not hasattr(obj, '__wrapped__'):
        rlog(10, 'prelude', 'can not unwrap %s: not wrapped' % (obj.__class__.__name__,))
        return
    for func in obj.__wrapped__.keys():
        rlog(10, 'prelude', 'unwrapping %s.%s' % (obj.__class__.__name__, func))
        obj.__setattr__(func, obj.__wrapped__[func])
    obj.__delattr__('__wrapped__')

ids = None

def init():
    global ids
    ids = Prelude()
    rlog(10, 'prelude', 'wrapping function calls...')
    try:
        from gozerbot.users import users as _users
        wrap(_users, 'allowed', lambda r, *a, **k: not r and 'userhost %s with not allowed (needs %s)' % (a[0], str(a[1])))
    except ImportError, e:
        rlog(10, 'prelude', 'import error: %s' % str(e))
    rlog(10, 'prelude', '... done')
    return 1

def shutdown():
    global ids
    if ids:
        try:
            from gozerbot.users import users as _users
            unwrap(_users)
        except ImportError, e:
            rlog(10, 'prelude', 'import error: %s' % str(e))

        ids.__destroy__()
    else:
        rlog(10, 'prelude', 'no prelude instance found, not unloading')
    return 1