AWS Elastic Load Balancing

import gzip
import io
import json
import shlex
import urllib.parse

import boto3
import requests

BETTER_STACK_SOURCE_TOKEN = '$SOURCE_TOKEN'
BETTER_STACK_URL = 'https://$INGESTING_HOST'

s3 = boto3.client('s3')

def process_log_line(line, fields):
    res = shlex.split(line, posix=False)
    return dict(zip(fields, res))

def get_load_balancer_type(key):
    if key.endswith(".gz"):
        key = key[:-3]

    file_parts = key.split(".")
    if len(file_parts) == 2:
        # Only one dot separating file name and extension - i.e. no dot in file name
        return "classic"

    loadbalancer_type = file_parts[0].split("_")[-1]
    return loadbalancer_type if loadbalancer_type in ["app", "net"] else "classic"

def get_log_fields(load_balancer_type):
    if load_balancer_type == "app":
        return ["type", "dt", "elb", "client_port", "target_port", "request_processing_time", "target_processing_time", "response_processing_time", "elb_status_code", "target_status_code", "received_bytes", "sent_bytes", "request", "user_agent", "ssl_cipher", "ssl_protocol", "target_group_arn", "trace_id", "domain_name", "chosen_cert_arn", "matched_rule_priority", "request_creation_time", "actions_executed", "redirect_url", "error_reason", "target_port_list", "target_status_code_list", "classification", "classification_reason", "conn_trace_id", "transformed_host", "transformed_uri", "request_transform_status", "ip_address"]

    if load_balancer_type == "net":
        return ["type", "version", "dt", "elb", "listener", "client_port", "destination_port", "connection_time", "tls_handshake_time", "received_bytes", "sent_bytes", "incoming_tls_alert", "chosen_cert_arn", "chosen_cert_serial", "tls_cipher", "tls_protocol_version", "tls_named_group", "domain_name", "alpn_fe_protocol", "alpn_be_protocol", "alpn_client_preference_list", "tls_connection_creation_time"]

    # use "classic" elb as a default
    return ["dt", "elb", "client_port", "backend_port", "request_processing_time", "backend_processing_time", "response_processing_time", "elb_status_code", "backend_status_code", "received_bytes", "sent_bytes", "request", "user_agent", "ssl_cipher", "ssl_protocol"]

def read_log_lines(bucket, key):
    body = s3.get_object(Bucket=bucket, Key=key)['Body']

    if key.endswith('.log.gz'):
        body = gzip.GzipFile(fileobj=body)
    elif not key.endswith('.log'):
        raise ValueError(f"Unexpected file extension for file {key} in bucket {bucket}. Expected .log or .log.gz file.")

    # Stream the file line by line, so large files don't have to fit into the Lambda's memory
    for line in io.TextIOWrapper(body, encoding='utf-8'):
        line = line.strip()
        if line:
            yield line

def upload_in_batches(session, logs):
    batch, batch_bytes = [], 2  # 2 bytes for the surrounding [ and ]

    for log in logs:
        encoded = json.dumps(log)  # ASCII only, so its length is its size in bytes
        # Better Stack accepts up to 10 MiB per request
        if batch and batch_bytes + len(encoded) + 1 > 10 * 1024 * 1024:
            upload_batch(session, batch)
            batch, batch_bytes = [], 2
        batch.append(encoded)
        batch_bytes += len(encoded) + 1  # 1 byte for the separating comma

    if batch:
        upload_batch(session, batch)

def upload_batch(session, batch):
    body = gzip.compress(f"[{','.join(batch)}]".encode('utf-8'))
    response = session.post(BETTER_STACK_URL, data=body)
    response.raise_for_status()

    print(f"Sent {len(batch)} logs to {BETTER_STACK_URL}. Got response code: {response.status_code}")

def lambda_handler(event, context):
    if '$' in BETTER_STACK_SOURCE_TOKEN or '$' in BETTER_STACK_URL:
        print("Error: BETTER_STACK_SOURCE_TOKEN or BETTER_STACK_URL at the top of the Lambda code still contains a placeholder. Set them to your Better Stack source token and ingesting host, then deploy again. No logs were sent.")
        return

    session = requests.Session()
    session.headers.update({
        'Authorization': f"Bearer {BETTER_STACK_SOURCE_TOKEN}",
        'Content-Type': 'application/json',
        'Content-Encoding': 'gzip'
    })

    for record in event['Records']:
        bucket = record['s3']['bucket']['name']
        key = urllib.parse.unquote_plus(record['s3']['object']['key'], encoding='utf-8')
        print(f"Processing {key} from bucket {bucket}")

        log_fields = get_log_fields(get_load_balancer_type(key))
        logs = (process_log_line(line, log_fields) for line in read_log_lines(bucket, key))
        upload_in_batches(session, logs)