Running Elastalert on AWS Lambda (and optionally using dynamic values in rules at runtime)

[briandefiant]

I wanted to share some code that we're using to run Elastalert2 in an AWS Lambda container function. The main benefits of this for us are that it gives us less server infrastructure to maintain, and we can include dynamic values in our rules at runtime (via a custom rules loader that loads rules from S3 and renders them using Jinja2 templating).

There are some existing issues and solutions out there for running on Lambda:

• Anyone tried to run Elastalert on AWS Lambda? Yelp/elastalert#496
• https://github.com/beezz/elastalert-lambda
• https://git.zero-downtime.net/ZeroDownTime/elastalert-lambda

None of these work since Elastalert started using APScheduler's BackgroundScheduler, which returns immediately after starting. This causes the Lambda function to exit before any rules have a chance to run (also detailed in this comment on an issue about running Elastalert once and then exiting).

I worked around this by extending the ElastAlerter class and overriding the start method. The new method schedules all jobs to run immediately, and only exits once all jobs have run (it uses an APScheduler listener function attached to the scheduler to determine when all jobs are complete). handle_pending_alerts is then called manually (instead of being added as a job to the scheduler). Here's an incomplete example (just the class) from my Lambda function code.

Also note that I overrode the handle_error method because I wanted to make the Lambda exit and report a failure if anything went wrong (with the exception of read timeouts from Elasticsearch, which I was seeing somewhat often with Amazon OpenSearch Service).

from apscheduler.events import EVENT_JOB_ERROR, EVENT_JOB_EXECUTED
from datetime import datetime
from elastalert.elastalert import ElastAlerter
from elastalert.util import elastalert_logger


class LambdaElastAlerter(ElastAlerter):
    def apscheduler_listener(self, event):
        if event.exception:
            elastalert_logger.error(
                "Listener exiting Elastalert due to exception from job ID \"{}\": {}".format(
                    event.job_id,
                    event.exception
                )
            )
            os._exit(1)
        else:
            elastalert_logger.info(
                "Listener marking job ID \"{}\" as successful and pausing".format(
                    event.job_id
                )
            )
            self.all_jobs[event.job_id]["success"] = True
            self.scheduler.pause_job(event.job_id)

    def start(self):
        # Wait for ES to become responsive
        self.wait_until_responsive(timeout=self.args.timeout)
        # Attach listener to scheduler so we can mark a job as succeeded once it's been executed
        self.scheduler.add_listener(self.apscheduler_listener, EVENT_JOB_ERROR | EVENT_JOB_EXECUTED)
        # Get all jobs (rules), set them to run now, and initially set their success value to false
        self.all_jobs = {}
        for job in self.scheduler.get_jobs():
            job.modify(next_run_time=datetime.now())
            self.all_jobs[job.id] = {"name": job.name, "success": False}
        # Start the scheduler
        elastalert_logger.info("Starting up")
        self.scheduler.start()
        self.running = True
        while self.running:
            # Check if all jobs have succeeded
            if all(v["success"] for v in iter(self.all_jobs.values())):
                elastalert_logger.info("All jobs complete!")
                # Get info about each job that ran
                for job in self.scheduler.get_jobs():
                    self.all_jobs[job.id].update({
                        "aggregate_alert_time": str(job.args[0]["aggregate_alert_time"]),
                        "processed_hits": str(job.args[0]["processed_hits"]),
                        "initial_starttime": str(job.args[0]["initial_starttime"]),
                        "starttime": str(job.args[0]["starttime"]),
                        "original_starttime": str(job.args[0]["original_starttime"]),
                    })
                elastalert_logger.info("Handling pending alerts")
                self.handle_pending_alerts()
                self.running = False
            else:
                elastalert_logger.info("All jobs have not run. Sleeping for 1 second.")
                time.sleep(1)
        return {
            "num_jobs": len(self.all_jobs),
            "job_info": self.all_jobs
        }

    def handle_error(self, message, data=None):
        # Logs message at error level and writes message, data and traceback to Elasticsearch
        elastalert_logger.error(message)
        body = {'message': message}
        tb = traceback.format_exc()
        body['traceback'] = tb.strip().split('\n')
        if data:
            body['data'] = data
        # If message is a read timeout, don't try to write to ES and don't exit
        if "ReadTimeout" not in message:
            self.writeback('elastalert_error', body)
            # Added for LambdaElastAlerter class so Lambda reports an error and exits if anything
            # fails
            elastalert_logger.error("Exiting on error")
            os._exit(1)

This is a minimal example of a Lambda handler that uses this class. Because of some issues with ever-increasing memory usage between Lambda warm starts that I couldn't track down (I tried a lot of different combinations of deling objects and manually calling gc.collect(), with no luck), it runs Elastalert in a separate process using Python's multiprocessing module to ensure that memory is freed when it's done. This example creates an Elastalert config file as a Python dict and then dumps it to a YAML file at /tmp/config.yml, but you could also just copy your config file into place with a step in the Dockerfile.

import json
import logging
import multiprocessing as mp

logger = logging.getLogger()
logger.setLevel("INFO")


def run_elastalert(args, send_end):
    # Create an instance of our custom LambdaElastAlerter class and start it. It will schedule all
    # jobs (i.e., rules) to run immediately, and return a result when they are all complete.
    client = LambdaElastAlerter(args)
    result = client.start()
    send_end.send(result)
    send_end.close()


def lambda_handler(event, context):
    elastalert_config = {
        # 'run_every' shouldn't have any effect when using the LambdaElastAlterter class. It's
        # normally used when using 'use_count_query' or 'use_terms_query' and previous endtime of a
        # rule can't be determined by other methods (like querying Elasticsearch). It's also
        # normally used to schedule next run of a rule, but this isn't relevant since the
        # LambdaElastAlerter class pauses jobs after they've run once.
        "run_every": {
            "minutes": 10
        },
        "buffer_time": {
            "minutes": 15
        },
        "alert_time_limit": {
            "days": 2
        },
        "use_ssl": True,
        "verify_certs": True,
        "disable_rules_on_error": False,
        "writeback_index": "elastalert_status",
        "writeback_alias": "elastalert_alerts",
        "aws_region": os.environ["AWS_DEFAULT_REGION"],
        # Allow infinite grace time just in case there's a delay between creating the
        # LambdaElastAlerter client and starting it
        "misfire_grace_time": None
    }

    # Set args to pass to Elastalert. We override the start() function of the ElastAlerter class
    # and don't call handle_config_change() to update rules, so pass --pin_rules to disable rule
    # change logic that might be elsewhere.
    elastalert_args = ["--config", "/tmp/config.yml", "--pin_rules"]
    if ELASTALERT_DEBUG == "True":
        elastalert_args.extend(["--debug"])
    else:
        elastalert_args.extend(["--verbose"])

    # Run in separate process to prevent ever-increasing memory usage
    # https://aws.amazon.com/blogs/compute/parallel-processing-in-python-with-aws-lambda/
    logger.info("Running Elastalert using multiprocessing.Process()")
    recv_end, send_end = mp.Pipe(False)
    process = mp.Process(target=run_elastalert, args=(elastalert_args, send_end,))
    process.start()
    process.join()
    result = recv_end.recv()

    # Log and return results
    logger.info(json.dumps(result))
    return json.dumps(result)

We're deploying this as a container on Lambda, using the following Dockerfile and requirements.txt.

FROM public.ecr.aws/lambda/python:3.10

COPY requirements.txt ./
RUN  pip3 install -r requirements.txt --target "$LAMBDA_TASK_ROOT"

RUN mkdir "${LAMBDA_TASK_ROOT}/elastalert_modules/"
RUN touch "${LAMBDA_TASK_ROOT}/elastalert_modules/__init__.py"

COPY lambda_function.py "$LAMBDA_TASK_ROOT"

CMD [ "lambda_function.lambda_handler" ]

elastalert2~=2.10

If all you want to do is run Elastalert on Lambda, you could stop here (and maybe add a COPY step to the Dockerfile to copy your rule files). But we're also using the following custom rules loader to download rules from S3.

import boto3
import hashlib
import jinja2
import json
import os
import yaml

from elastalert.loaders import RulesLoader
from elastalert.util import elastalert_logger

# See:
#   https://github.com/jertel/elastalert2/blob/master/elastalert/loaders.py#L537
#   https://elastalert2.readthedocs.io/en/latest/recipes/adding_loaders.html


class S3RulesLoader(RulesLoader):
    # Required global (config.yaml) configuration options for the loader
    required_globals = frozenset(['rules_bucket'])

    def __init__(self, conf):
        super(S3RulesLoader, self).__init__(conf)
        self.session = boto3.Session()
        self.s3_client = self.session.client("s3")
        self.bucket = conf["rules_bucket"]
        if "jinja_enabled" in conf and conf["jinja_enabled"]:
            self.jinja_enabled = True
            self.jinja_context = conf["jinja_context"]
        elastalert_logger.info(f"Initialized S3 rules loader with bucket {self.bucket}")

    def get_names(self, conf, use_rule=None):
        if use_rule and os.path.isfile(use_rule):
            return [use_rule]
        object_keys = []
        objects_paginator = self.s3_client.get_paginator("list_objects_v2")
        for page in objects_paginator.paginate(Bucket=self.bucket):
            for object in page["Contents"]:
                if "Key" in object and self.is_yaml(object["Key"]):
                    object_keys.append(object["Key"])
        elastalert_logger.info(f"Found object keys {json.dumps(object_keys)}")
        return object_keys

    def get_hashes(self, conf, use_rule=None):
        rule_objects = self.get_names(conf, use_rule)
        rule_mod_times = {}
        for rule_object in rule_objects:
            rule_mod_times[rule_object] = self.get_rule_file_hash(rule_object)
        return rule_mod_times

    def get_yaml(self, object_key):
        object = self.s3_client.get_object(Bucket=self.bucket, Key=object_key)
        if self.jinja_enabled:
            environment = jinja2.Environment()
            template = environment.from_string(object["Body"].read().decode())
            rule_yaml = template.render(self.jinja_context)
        else:
            rule_yaml = object["Body"].read().decode()
        return yaml.load(rule_yaml, Loader=yaml.FullLoader)

    def get_rule_file_hash(self, object_key):
        object = self.s3_client.get_object(Bucket=self.bucket, Key=object_key)
        object_hash = hashlib.sha1(object["Body"].read()).digest()
        for import_rule_object in self.import_rules.get(object_key, []):
            object_hash += self.get_rule_file_hash(import_rule_object)
        return object_hash

    def is_yaml(self, object_key):
        if self.jinja_enabled:
            return (
                object_key.endswith(".yaml")
                or object_key.endswith(".yml")
                or object_key.endswith(".j2")
            )
        else:
            return (
                object_key.endswith(".yaml")
                or object_key.endswith(".yml")
            )

If you save this as custom_loaders.py, you can copy it into the Lambda image by adding this to the Dockerfile:

COPY custom_loaders.py "${LAMBDA_TASK_ROOT}/elastalert_modules/"

To use the S3 rule loader (and also enable Jinja templating) add the following lines to the Elastalert config defined in the Lambda code.

        # Use custom rules loader to get rules from S3, and enable Jinja rendering of rule files
        "rules_loader": "elastalert_modules.custom_loaders.S3RulesLoader",
        "rules_bucket": RULES_BUCKET,
        "jinja_enabled": True,
        "jinja_context": { "some_key": "some_value" },

You could then write an Elastalert rule with a filter like this, which will query for some_value:

filter:
- bool:
    must:
      - match_phrase:
          eventSource:
            query: {{ some_key }}

This is the main reason we define the Elastalert config in Python, so that we can dynamically extend the Jinja context that is passed to the rules loader.

I hope this is helpful for anyone else looking to run Elastalert on Lambda!

[jertel]

This is very interesting. Do you have any interest in submitting a PR with the changes to this project, in a way such that it doesn't conflict with the main usage format? Perhaps place the new python code into an elastalert.lambda directory, and add a startup parameter to tell it to use "lambda" mode. The new code would need unit test coverage, and the documentation would need updated.

Looking at the Dockerfile, it looks like it might be possible to have a single ElastAlert 2 Docker image that can handle both run modes (original and lamba). Unless that base image you are using is a Lambda requirement.