Building A WAF Query Agent To Analyse WAF Logs Using GPT-4o-mini

I’ve recently been working with AWS WAF logs, and I must say, analysing them can be quite the challenge. The logs are verbose, contain nested JSON structures, and extracting meaningful insights often requires complex queries. After struggling with traditional tools, including AWS Athena, I decided to build my own solution that lets me query WAF logs using plain English (or any other language, actually!). Here’s how I did it.

The Challenge

Have you ever tried to answer questions like “Show me all IPs that sent more than 1000 requests with a user-agent containing ‘Python’ in the last hour”? With traditional tools, you’d need to write complex queries or scripts. Even worse, if you want to modify the query slightly, you often need to rewrite significant portions of the code.

Enter the WAF Query Agent

I built an agent that leverages GPT-4o-mini to translate natural language queries into Python code. But here’s the clever bit: instead of sending massive amounts of log data to the AI for analysis (which would be slow and expensive), we only send the query and the data schema. The AI Agent generates the Python code needed to answer that query, and then we run this code locally on our logs.

Here’s how it works:

1. User types a query in plain English.
2. The query and WAF log schema (not the actual data) go to GPT-4o-mini.
3. GPT-4o-mini generates specific Python code to answer that query.
4. Our agent runs this code locally against the actual log files.
5. Results are formatted and displayed.

This architecture has several advantages:

• Natural Language Queries: No need to learn complex query syntax. Just ask questions in plain English, or Spanish, or…
• Flexibility: Want to modify your query? Just ask differently. The agent figures out the new code.
• Intelligent Analysis: The agent understands context and can handle complex analytical queries.
• Cost-effective: We’re only sending tiny prompts to the AI, not gigabytes of log data. GPT-4o-mini is good enough to interpret the query and generate the Python code, and very cheap for that.
• Private: Your log data stays local; only the schema is shared.
• Fast: Processing happens locally, not waiting for API responses.
• Scalable: Can handle any size of log files as processing is local.

Here’s a simple example:

query = "How many unique IPs made requests from countries other than the UK?"
result = engine.execute_query(query)

Behind the scenes, the agent generates and executes Python code like this:

=== Generated Code ===
# Initialize a set to store unique IPs
unique_ips = set()
# Iterate through the logs to find unique IPs from countries other than the UK
for log in logs:
    if 'httpRequest' in log and 'clientIp' in log['httpRequest'] and 'country' in log['httpRequest']:
        country = log['httpRequest']['country']
        ip = log['httpRequest']['clientIp']
        # Check if the country is not the UK
        if country != 'UK':
            unique_ips.add(ip)
# Set result to the count of unique IPs
result = len(unique_ips)

Current Implementation and Future Improvements

Currently, the agent requires downloaded WAF logs from S3 to your local machine before analysis. These logs are read into memory for processing. While this approach works well for moderate amounts of data, it does present some limitations when dealing with very large log files or when requiring real-time analysis. If you are dealing with an incident that spans many hours or days, my current architecture may be inadequate.

So, I’m actively exploring several architectural improvements to make this more dynamic:

• Streaming processing to handle logs without full memory loading.
• Direct S3 integration with selective downloading.
• Intelligent caching and indexing strategies.
• Chunked processing for large datasets.

The goal is to maintain the current performance and flexibility while handling larger datasets more efficiently.

I also want to try it out with other cheaper models, such as the AWS Nova models that AWS has just recently launched this month (December 2024).

Clever Code Highlights

One of the most interesting parts of this agent is how it safely executes the generated code. Here’s a snippet of how we create a sandboxed environment:

def execute_code(self, code: str) -> Dict[str, Any]:
    """Execute generated code safely."""
    stdout = StringIO()
    
    try:
        # Create safe execution environment with limited builtins
        safe_builtins = {
            'print': print,
            'len': len,
            'dict': dict,
            'list': list,
            'set': set,
            'int': int,
            'float': float,
            'str': str,
            'bool': bool,
            'sum': sum,
            'min': min,
            'max': max,
            'sorted': sorted,
            'map': map,
        }
        
        # Set up isolated execution context
        globals_dict = {'__builtins__': safe_builtins}
        locals_dict = {
            'logs': self.data_loader.data,
            'result': None
        }
        
        # Execute code and capture output
        with redirect_stdout(stdout):
            exec(code, globals_dict, locals_dict)
        
        return {
            'success': True,
            'result': locals_dict.get('result'),
            'output': stdout.getvalue()
        }

This code ensures that the generated Python can’t access system resources or import potentially dangerous modules.

Command-Line Interface Features

While the natural language processing is the star of the show, I’ve also put some effort into making the tool pleasant to use day-to-day. Here are a couple of the interface features that make it particularly user-friendly:

Command History Navigation

The agent maintains a persistent command history between sessions. Using the prompt_toolkit library, I’ve implemented a robust history system:

from prompt_toolkit import PromptSession
from prompt_toolkit.history import History

class MemoryHistory(History):
    def __init__(self, memory):
        super().__init__()
        self.memory = memory

    def load_history_strings(self) -> Iterable[str]:
        """Load history from memory in reverse chronological order."""
        return self.memory.get_last_queries()  # Newest first

This means you can:

• Use Up/Down arrows to cycle through previous queries.
• Access history even after restarting the tool.
• Search through history using Ctrl+R.
• Never lose your carefully crafted queries.

Smart Pagination

For queries that return large result sets, I’ve implemented an effective pagination system. I quickly realised that in some cases the output would be thousands and thousands of IPs for example, so pagination was kind of necessary.

@classmethod
def format_and_paginate(cls, result, prompt_session=None):
    """Format result and handle pagination."""
    # Convert result to list of lines
    lines = cls._convert_to_lines(result)
    
    current_line = 0
    while current_line < len(lines):
        # Display current page (20 lines)
        end_line = min(current_line + cls.ITEMS_PER_PAGE, len(lines))
        print("\n".join(lines[current_line:end_line]))
        
        # Prompt for continuation if more lines exist
        if end_line < len(lines):
            response = cls._get_continuation_response(prompt_session)
            if response in ('n', 'no'):
                break
            current_line = end_line

This ensures that:

• Results are displayed in manageable chunks.
• You don’t get overwhelmed with output.
• You can easily review large datasets.

Conclusion

Gone are the days of writing complex queries or scripts for WAF log analysis. With this agent, you can focus on what you want to know, not how to ask for it. It’s fast, it’s flexible, and most importantly, it makes log analysis accessible to everyone on the team, not just the coding experts.

But this is just the beginning. Imagine this tool as part of a multi-AI agents deployment, each specialised in different aspects of web application security. Here’s how it could work:

1. A Detection Agent could analyse patterns and alert you: “I’m detecting a distributed credential stuffing attack from multiple IP addresses, primarily originating from <area of the world>”.
2. Our WAF Query Agent could then automatically generate and execute the necessary queries:
   • “Show me all IPs making failed login attempts in the last hour”
   • “What’s the rate of requests per IP?”
   • “Are there common user-agent patterns?”
   • “Show me the geographic distribution of these requests”
3. A Defence Agent could then use this information to craft precise WAF rules:
   • Creating rate-limiting rules for the identified IP ranges.
   • Implementing user-agent blacklists.
   • Setting up geographic restrictions.
   • Suggesting additional security headers or challenge rules.

The beauty of this system would be its ability to adapt and learn. Each agent could specialise in its domain while working together to provide comprehensive security coverage. The WAF Query Agent would be the crucial bridge between detection and response, providing the detailed evidence needed for targeted, real-time mitigation.

Disclaimer: The perspectives shared here are my own and do not necessarily represent those of my employer. I use GenAI as a tool to help me compose and structure my articles.