示例#1
0
class BaseModInput(smi.Script):
    '''
    This is a modular input wrapper, which provides some helper
    functions to read the paramters from setup pages and the arguments
    from input definition
    '''
    LogLevelMapping = {'debug': logging.DEBUG,
                       'info': logging.INFO,
                       'warning': logging.WARNING,
                       'error': logging.ERROR,
                       'critical': logging.CRITICAL}

    def __init__(self, app_namespace, input_name, use_single_instance=False):
        super(BaseModInput, self).__init__()
        self.use_single_instance = use_single_instance
        self._canceled = False
        self.input_type = input_name
        self.input_stanzas = {}
        self.context_meta = {}
        self.namespace = app_namespace
        # redirect all the logging to one file
        Logs.set_context(namespace=app_namespace,
                         root_logger_log_file=input_name)
        self.logger = logging.getLogger()
        self.logger.setLevel(logging.INFO)
        self.rest_helper = TARestHelper(self.logger)
        # check point
        self.ckpt = None
        self.setup_util = None

    @property
    def app(self):
        return self.get_app_name()

    @property
    def global_setup_util(self):
        """
        This is a private API used in AoB code internally. It is not allowed to be used in user's code.

        :return: setup util instance to read global configurations
        """
        return self.setup_util

    def get_app_name(self):
        """Get TA name.

        :return: the name of TA this modular input is in
        """
        raise NotImplemented

    def get_scheme(self):
        """Get basic scheme, with use_single_instance field set.

        :return: a basic input scheme
        """
        scheme = smi.Scheme(self.input_type)
        scheme.use_single_instance = self.use_single_instance
        return scheme

    def stream_events(self, inputs, ew):
        """The method called to stream events into Splunk.

        This method overrides method in splunklib modular input.
        It pre-processes the input args and call collect_events to stream events.

        :param inputs: An ``InputDefinition`` object.
        :param ew: An object with methods to write events and log messages to Splunk.
        """
        # the input metadata is like
        # {
        #     'server_uri': 'https://127.0.0.1:8089',
        #     'server_host': 'localhost',
        #     'checkpoint_dir': '...',
        #     'session_key': 'ceAvf3z^hZHYxe7wjTyTNo6_0ZRpf5cvWPdtSg'
        # }
        self.context_meta = inputs.metadata
        # init setup util
        uri = inputs.metadata["server_uri"]
        session_key = inputs.metadata['session_key']
        self.setup_util = Setup_Util(uri, session_key, self.logger)

        input_definition = smi.input_definition.InputDefinition()
        input_definition.metadata = copy.deepcopy(inputs.metadata)
        input_definition.inputs = copy.deepcopy(inputs.inputs)
        try:
            self.parse_input_args(input_definition)
        except Exception as e:
            import traceback
            self.log_error(traceback.format_exc(e))
            print >> sys.stderr, traceback.format_exc(e)
            self.input_stanzas = {}
        if not self.input_stanzas:
            # if no stanza found. Just return
            return
        try:
            self.set_log_level(self.log_level)
        except:
            self.log_debug('set log level fails.')
        try:
            self.collect_events(ew)
        except Exception as e:
            import traceback
            self.log_error('Get error when collecting events.\n' + traceback.format_exc(e))
            print >> sys.stderr, traceback.format_exc(e)
            raise RuntimeError(str(e))

    def collect_events(self, event_writer):
        """Collect events and stream to Splunk using event writer provided.

        Note: This method is originally collect_events(self, inputs, event_writer).

        :param event_writer: An object with methods to write events and log messages to Splunk.
        """
        raise NotImplemented()

    def parse_input_args(self, inputs):
        """Parse input arguments, either from os environment when testing or from global configuration.

        :param inputs: An ``InputDefinition`` object.
        :return:
        """
        if os.environ.get(AOB_TEST_FLAG, 'false') == 'true':
            self._parse_input_args_from_env(inputs)
        else:
            self._parse_input_args_from_global_config(inputs)
        if not self.use_single_instance:
            assert len(self.input_stanzas) == 1

    def _parse_input_args_from_global_config(self, inputs):
        """Parse input arguments from global configuration.

        :param inputs:
        """
        dirname = os.path.dirname
        config_path = os.path.join(dirname(dirname(dirname(dirname(__file__)))), 'appserver', 'static', 'js', 'build',
                                   'globalConfig.json')
        with open(config_path) as f:
            schema_json = ''.join([l for l in f])
        global_schema = GlobalConfigSchema(json.loads(schema_json))

        uri = inputs.metadata["server_uri"]
        session_key = inputs.metadata['session_key']
        global_config = GlobalConfig(uri, session_key, global_schema)
        ucc_inputs = global_config.inputs.load(input_type=self.input_type)
        all_stanzas = ucc_inputs.get(self.input_type, {})
        if not all_stanzas:
            # for single instance input. There might be no input stanza.
            # Only the default stanza. In this case, modinput should exit.
            self.log_warning("No stanza found for input type: " + self.input_type)
            sys.exit(0)

        account_fields = self.get_account_fields()
        checkbox_fields = self.get_checkbox_fields()
        self.input_stanzas = {}
        for stanza in all_stanzas:
            full_stanza_name = '{}://{}'.format(self.input_type, stanza.get('name'))
            if full_stanza_name in inputs.inputs:
                if stanza.get('disabled', False):
                    raise RuntimeError("Running disabled data input!")
                stanza_params = {}
                for k, v in stanza.iteritems():
                    if k in checkbox_fields:
                        stanza_params[k] = sutils.is_true(v)
                    elif k in account_fields:
                        stanza_params[k] = copy.deepcopy(v)
                    else:
                        stanza_params[k] = v
                self.input_stanzas[stanza.get('name')] = stanza_params

    def _parse_input_args_from_env(self, inputs):
        """Parse input arguments from os environment. This is used for testing inputs.

        :param inputs:
        """
        data_inputs_options = json.loads(os.environ.get(DATA_INPUTS_OPTIONS, '[]'))
        account_fields = self.get_account_fields()
        checkbox_fields = self.get_checkbox_fields()
        self.input_stanzas = {}
        while len(inputs.inputs) > 0:
            input_stanza, stanza_args = inputs.inputs.popitem()
            kind_and_name = input_stanza.split("://")
            if len(kind_and_name) == 2:
                stanza_params = {}
                for arg_name, arg_value in stanza_args.iteritems():
                    try:
                        arg_value_trans = json.loads(arg_value)
                    except ValueError:
                        arg_value_trans = arg_value
                    stanza_params[arg_name] = arg_value_trans
                    if arg_name in account_fields:
                        stanza_params[arg_name] = self.get_user_credential_by_id(arg_value_trans)
                    elif arg_name in checkbox_fields:
                        stanza_params[arg_name] = sutils.is_true(arg_value_trans)
                self.input_stanzas[kind_and_name[1]] = stanza_params

    def get_account_fields(self):
        """Get the names of account variables.

        Should be implemented in subclass.

        :return: a list of variable names
        """
        raise NotImplemented

    def get_checkbox_fields(self):
        """Get the names of checkbox variables.

        Should be implemented in subclass.

        :return: a list of variable names
        """
        raise NotImplemented

    def get_global_checkbox_fields(self):
        """Get the names of checkbox global parameters.

        :return: a list of global variable names
        """
        raise NotImplemented

    # Global setting related functions.
    # Global settings consist of log setting, proxy, account(user_credential) and customized settings.
    @property
    def log_level(self):
        return self.get_log_level()

    def get_log_level(self):
        """Get the log level configured in global configuration.

        :return: log level set in global configuration or "INFO" by default.
        """
        return self.setup_util.get_log_level()

    def set_log_level(self, level):
        """Set the log level this python process uses.

        :param level: log level in `string`. Accept "DEBUG", "INFO", "WARNING", "ERROR" and "CRITICAL".
        """
        if isinstance(level, basestring):
            level = level.lower()
            if level in self.LogLevelMapping:
                level = self.LogLevelMapping[level]
            else:
                level = logging.INFO
        self.logger.setLevel(level)

    def log(self, msg):
        """Log msg using logging level in global configuration.

        :param msg: log `string`
        """
        self.logger.log(level=self.log_level, msg=msg)

    def log_debug(self, msg):
        """Log msg using logging.DEBUG level.

        :param msg: log `string`
        """
        self.logger.debug(msg)

    def log_info(self, msg):
        """Log msg using logging.INFO level.

        :param msg: log `string`
        """
        self.logger.info(msg)

    def log_warning(self, msg):
        """Log msg using logging.WARNING level.

        :param msg: log `string`
        """
        self.logger.warning(msg)

    def log_error(self, msg):
        """Log msg using logging.ERROR level.

        :param msg: log `string`
        """
        self.logger.error(msg)

    def log_critical(self, msg):
        """Log msg using logging.CRITICAL level.

        :param msg: log `string`
        """
        self.logger.critical(msg)

    @property
    def proxy(self):
        return self.get_proxy()

    def get_proxy(self):
        """Get proxy settings in global configuration.

        Proxy settings include fields "proxy_url", "proxy_port", "proxy_username", "proxy_password", "proxy_type" and "proxy_rdns".

        :return: a `dict` containing proxy parameters or empty `dict` if proxy is not set.
        """
        return self.setup_util.get_proxy_settings()

    def get_user_credential_by_username(self, username):
        """Get global credential information based on username.

        Credential settings include fields "name"(account id), "username" and "password".

        :param username: `string`
        :return: if credential with username exists, return a `dict`, else None.
        """
        return self.setup_util.get_credential_by_username(username)

    def get_user_credential_by_id(self, account_id):
        """Get global credential information based on account id.

        Credential settings include fields "name"(account id), "username" and "password".

        :param account_id: `string`
        :return: if credential with account_id exists, return a `dict`, else None.
        """
        return self.setup_util.get_credential_by_id(account_id)

    def get_global_setting(self, var_name):
        """Get customized setting value configured in global configuration.

        :param var_name: `string`
        :return: customized global configuration value or None
        """
        var_value = self.setup_util.get_customized_setting(var_name)
        if var_value is not None and var_name in self.get_global_checkbox_fields():
            var_value = sutils.is_true(var_value)
        return var_value

    # Functions to help create events.
    def new_event(self, data, time=None, host=None, index=None, source=None, sourcetype=None, done=True, unbroken=True):
        """Create a Splunk event object.

        :param data: ``string``, the event's text.
        :param time: ``float``, time in seconds, including up to 3 decimal places to represent milliseconds.
        :param host: ``string``, the event's host, ex: localhost.
        :param index: ``string``, the index this event is specified to write to, or None if default index.
        :param source: ``string``, the source of this event, or None to have Splunk guess.
        :param sourcetype: ``string``, source type currently set on this event, or None to have Splunk guess.
        :param done: ``boolean``, is this a complete ``Event``? False if an ``Event`` fragment.
        :param unbroken: ``boolean``, Is this event completely encapsulated in this ``Event`` object?
        :return: ``Event`` object
        """
        return smi.Event(data=data, time=time, host=host, index=index,
                         source=source, sourcetype=sourcetype, done=done, unbroken=unbroken)

    # Basic get functions. To get params in input stanza.
    def get_input_type(self):
        """Get input type.

        :return: the modular input type
        """
        return self.input_type

    def get_input_stanza(self, input_stanza_name=None):
        """Get input stanzas.

        If stanza name is None, return a dict with stanza name as key and params as values.
        Else return a dict with param name as key and param value as value.

        :param input_stanza_name: None or `string`
        :return: `dict`
        """
        if input_stanza_name:
            return self.input_stanzas.get(input_stanza_name, None)
        return self.input_stanzas

    def get_input_stanza_names(self):
        """Get all stanza names this modular input instance is given.

        For multi instance mode, a single string value will be returned.
        For single instance mode, stanza names will be returned in a list.

        :return: `string` or `list`
        """
        if self.input_stanzas:
            names = self.input_stanzas.keys()
            if self.use_single_instance:
                return names
            else:
                assert len(names) == 1
                return names[0]
        return None

    def get_arg(self, arg_name, input_stanza_name=None):
        """Get the input argument.

        If input_stanza_name is not provided:
            For single instance mode, return a dict <input_name, arg_value>.
            For multi instance mode, return a single value or None.
        If input_stanza_name is provided, return a single value or None.

        :param arg_name: `string`, argument name
        :param input_stanza_name: None or `string`, a stanza name
        :return: `dict` or `string` or None
        """
        if input_stanza_name is None:
            args_dict = {k: args[
                arg_name] for k, args in self.input_stanzas.iteritems() if arg_name in args}
            if self.use_single_instance:
                return args_dict
            else:
                if len(args_dict) == 1:
                    return args_dict.values()[0]
                return None
        else:
            return self.input_stanzas.get(input_stanza_name, {}).get(arg_name, None)

    def get_output_index(self, input_stanza_name=None):
        """Get output Splunk index.

        :param input_stanza_name: `string`
        :return: `string` output index
        """
        return self.get_arg('index', input_stanza_name)

    def get_sourcetype(self, input_stanza_name=None):
        """Get sourcetype to index.

        :param input_stanza_name: `string`
        :return: the sourcetype to index to
        """
        return self.get_arg('sourcetype', input_stanza_name)

    # HTTP request helper
    def send_http_request(self, url, method, parameters=None, payload=None, headers=None, cookies=None, verify=True,
                          cert=None, timeout=None, use_proxy=True):
        """Send http request and get response.

        :param url: URL for the new Request object.
        :param method: method for the new Request object. Can be "GET", "POST", "PUT", "DELETE"
        :param parameters: (optional) Dictionary or bytes to be sent in the query string for the Request.
        :param payload: (optional) Dictionary, bytes, or file-like object to send in the body of the Request.
        :param headers: (optional) Dictionary of HTTP Headers to send with the Request.
        :param cookies: (optional) Dict or CookieJar object to send with the Request.
        :param verify: (optional) whether the SSL cert will be verified. A CA_BUNDLE path can also be provided.
        :param cert: (optional) if String, path to ssl client cert file (.pem). If Tuple, ('cert', 'key') pair.
        :param timeout: (optional) How long to wait for the server to send data before giving up, as a float,
            or a (connect timeout, read timeout) tuple. Default to (10.0, 5.0).
        :param use_proxy: (optional) whether to use proxy. If set to True, proxy in global setting will be used.
        :return: Response
        """
        return self.rest_helper.send_http_request(url=url, method=method, parameters=parameters, payload=payload,
                                                  headers=headers, cookies=cookies, verify=verify, cert=cert,
                                                  timeout=timeout,
                                                  proxy_uri=self._get_proxy_uri() if use_proxy else None)

    def _get_proxy_uri(self):
        uri = None
        proxy = self.get_proxy()
        if proxy and proxy.get('proxy_url') and proxy.get('proxy_type'):
            uri = proxy['proxy_url']
            if proxy.get('proxy_port'):
                uri = '{0}:{1}'.format(uri, proxy.get('proxy_port'))
            if proxy.get('proxy_username') and proxy.get('proxy_password'):
                uri = '{0}://{1}:{2}@{3}/'.format(proxy['proxy_type'], proxy[
                    'proxy_username'], proxy['proxy_password'], uri)
            else:
                uri = '{0}://{1}'.format(proxy['proxy_type'], uri)
        return uri

    # Checkpointing related functions
    def _init_ckpt(self):
        if self.ckpt is None:
            if 'AOB_TEST' in os.environ:
                ckpt_dir = self.context_meta.get('checkpoint_dir', tempfile.mkdtemp())
                if not os.path.exists(ckpt_dir):
                    os.makedirs(ckpt_dir)
                self.ckpt = checkpointer.FileCheckpointer(ckpt_dir)
            else:
                if 'server_uri' not in self.context_meta:
                    raise ValueError('server_uri not found in input meta.')
                if 'session_key' not in self.context_meta:
                    raise ValueError('session_key not found in input meta.')
                dscheme, dhost, dport = sutils.extract_http_scheme_host_port(self.context_meta[
                                                                                 'server_uri'])
                self.ckpt = checkpointer.KVStoreCheckpointer(self.app + "_checkpointer",
                                                             self.context_meta['session_key'], self.app,
                                                             scheme=dscheme, host=dhost, port=dport)

    def get_check_point(self, key):
        """Get checkpoint.

        :param key: `string`
        :return: Checkpoint state if exists else None.
        """
        if self.ckpt is None:
            self._init_ckpt()
        return self.ckpt.get(key)

    def save_check_point(self, key, state):
        """Update checkpoint.

        :param key: Checkpoint key. `string`
        :param state: Checkpoint state.
        """
        if self.ckpt is None:
            self._init_ckpt()
        self.ckpt.update(key, state)

    def batch_save_check_point(self, states):
        """Batch update checkpoint.

        :param states: a `dict` states with chekpoint key as key and checkpoint state as value.
        """
        if self.ckpt is None:
            self._init_ckpt()
        self.ckpt.batch_update(states)

    def delete_check_point(self, key):
        """Delete checkpoint.

        :param key: Checkpoint key. `string`
        """
        if self.ckpt is None:
            self._init_ckpt()
        self.ckpt.delete(key)
示例#2
0
class ModularAlertBase(ModularAction):
    def __init__(self, ta_name, alert_name):
        self._alert_name = alert_name
        # self._logger_name = "modalert_" + alert_name
        self._logger_name = alert_name + "_modalert"
        self._logger = get_logger(self._logger_name)
        super(ModularAlertBase, self).__init__(
            sys.stdin.read(), self._logger, alert_name)
        self.setup_util_module = None
        self.setup_util = None
        self.result_handle = None
        self.ta_name = ta_name
        self.splunk_uri = self.settings.get('server_uri')
        self.setup_util = Setup_Util(self.splunk_uri, self.session_key, self._logger)

        self.rest_helper = TARestHelper(self._logger)

    def log_error(self, msg):
        self.message(msg, 'failure', level=logging.ERROR)

    def log_info(self, msg):
        self.message(msg, 'success', level=logging.INFO)

    def log_debug(self, msg):
        self.message(msg, None, level=logging.DEBUG)

    def log_warn(self, msg):
        self.message(msg, None, level=logging.WARN)

    def set_log_level(self, level):
        self._logger.setLevel(level)

    def get_param(self, param_name):
        return self.configuration.get(param_name)

    def get_global_setting(self, var_name):
        return self.setup_util.get_customized_setting(var_name)

    def get_user_credential(self, username):
        '''
        if the username exists, return
        {
            "username": username,
            "password": credential
        }
        '''
        return self.setup_util.get_credential_by_username(username)

    @property
    def log_level(self):
        return self.get_log_level()

    @property
    def proxy(self):
        return self.get_proxy()

    def get_log_level(self):
        return self.setup_util.get_log_level()

    def get_proxy(self):
        ''' if the proxy setting is set. return a dict like
        {
        proxy_url: ... ,
        proxy_port: ... ,
        proxy_username: ... ,
        proxy_password: ... ,
        proxy_type: ... ,
        proxy_rdns: ...
        }
        '''
        return self.setup_util.get_proxy_settings()

    def _get_proxy_uri(self):
        uri = None
        proxy = self.get_proxy()
        if proxy and proxy.get('proxy_url') and proxy.get('proxy_type'):
            uri = proxy['proxy_url']
            if proxy.get('proxy_port'):
                uri = '{0}:{1}'.format(uri, proxy.get('proxy_port'))
            if proxy.get('proxy_username') and proxy.get('proxy_password'):
                uri = '{0}://{1}:{2}@{3}/'.format(proxy['proxy_type'], proxy[
                                                  'proxy_username'], proxy['proxy_password'], uri)
            else:
                uri = '{0}://{1}'.format(proxy['proxy_type'], uri)
        return uri

    def send_http_request(self, url, method, parameters=None, payload=None, headers=None, cookies=None, verify=True, cert=None, timeout=None, use_proxy=True):
        return self.rest_helper.send_http_request(url=url, method=method, parameters=parameters, payload=payload,
                                                  headers=headers, cookies=cookies, verify=verify, cert=cert,
                                                  timeout=timeout,
                                                  proxy_uri=self._get_proxy_uri() if use_proxy else None)

    def build_http_connection(self, config, timeout=120,
                              disable_ssl_validation=False):
        from httplib2 import (socks, ProxyInfo, Http)
        """
        :config: dict like, proxy and account information are in the following
                format {
                    "username": xx,
                    "password": yy,
                    "proxy_url": zz,
                    "proxy_port": aa,
                    "proxy_username": bb,
                    "proxy_password": cc,
                    "proxy_type": http,http_no_tunnel,sock4,sock5,
                    "proxy_rdns": 0 or 1,
                }
        :return: Http2.Http object
        """
        if not config:
            config = {}

        proxy_type_to_code = {
            "http": socks.PROXY_TYPE_HTTP,
            "http_no_tunnel": socks.PROXY_TYPE_HTTP_NO_TUNNEL,
            "socks4": socks.PROXY_TYPE_SOCKS4,
            "socks5": socks.PROXY_TYPE_SOCKS5,
        }
        if config.get("proxy_type") in proxy_type_to_code:
            proxy_type = proxy_type_to_code[config["proxy_type"]]
        else:
            proxy_type = socks.PROXY_TYPE_HTTP

        rdns = config.get("proxy_rdns")

        proxy_info = None
        if config.get("proxy_url") and config.get("proxy_port"):
            if config.get("proxy_username") and config.get("proxy_password"):
                proxy_info = ProxyInfo(proxy_type=proxy_type,
                                       proxy_host=config["proxy_url"],
                                       proxy_port=int(config["proxy_port"]),
                                       proxy_user=config["proxy_username"],
                                       proxy_pass=config["proxy_password"],
                                       proxy_rdns=rdns)
            else:
                proxy_info = ProxyInfo(proxy_type=proxy_type,
                                       proxy_host=config["proxy_url"],
                                       proxy_port=int(config["proxy_port"]),
                                       proxy_rdns=rdns)
        if proxy_info:
            http = Http(proxy_info=proxy_info, timeout=timeout,
                        disable_ssl_certificate_validation=disable_ssl_validation)
        else:
            http = Http(timeout=timeout,
                        disable_ssl_certificate_validation=disable_ssl_validation)

        if config.get("username") and config.get("password"):
            http.add_credentials(config["username"], config["password"])
        return http

    def process_event(self, *args, **kwargs):
        raise NotImplemented()

    def pre_handle(self, num, result):
        result.setdefault('rid', str(num))
        self.update(result)
        return result

    def get_events(self):
        try:
            try:
                self.result_handle = gzip.open(self.results_file, 'rt')
            except ValueError: # Workaround for Python 2.7 on Windows
                self.result_handle = gzip.open(self.results_file, 'r')
            return (self.pre_handle(num, result) for num, result in enumerate(csv.DictReader(self.result_handle)))
        except IOError:
            msg = "Error: {}."
            self.log_error(msg.format("No search result. Cannot send alert action."))
            sys.exit(2)

    def prepare_meta_for_cam(self):
        try:
            try:
                rf = gzip.open(self.results_file, 'rt')
            except ValueError: # Workaround for Python 2.7 on Windows
                rf = gzip.open(self.results_file, 'r')
            for num, result in enumerate(csv.DictReader(rf)):
                result.setdefault('rid', str(num))
                self.update(result)
                self.invoke()
                break
        finally:
            if rf:
                rf.close()

    def run(self, argv):
        status = 0
        if len(argv) < 2 or argv[1] != "--execute":
            msg = 'Error: argv="{}", expected="--execute"'.format(argv)
            print(msg, file=sys.stderr)
            sys.exit(1)

        # prepare meta first for permission lack error handling: TAB-2455
        self.prepare_meta_for_cam()
        try:
            level = self.get_log_level()
            if level:
                self._logger.setLevel(level)
        except Exception as e:
            if e and '403' in str(e):
                self.log_error('User does not have permissions')
            else:
                self.log_error('Unable to set log level')
            sys.exit(2)

        try:
            status = self.process_event()
        except IOError:
            msg = "Error: {}."
            self.log_error(msg.format("No search result. Cannot send alert action."))
            sys.exit(2)
        except Exception as e:
            msg = "Unexpected error: {}."
            if e:
                self.log_error(msg.format(str(e)))
            else:
                import traceback
                self.log_error(msg.format(traceback.format_exc()))
            sys.exit(2)
        finally:
            if self.result_handle:
                self.result_handle.close()

        return status