POC of Airflow DAG Validator

poc_airflow_dag_validator.py

# =============================================================================
# Airflow DAG Validator — PROOF OF CONCEPT
# =============================================================================
# /!\ WARNING: This script is a Proof of Concept (POC) only.
#     It is intended to demonstrate the feasibility of static analysis on
#     Airflow DAG files prior to deployment. It must NOT be used as the sole
#     security mechanism in a production environment.
#
# Purpose:
#   Scans Airflow DAG files submitted by clients before they are deployed on
#   the platform infrastructure. The goal is to detect malicious or suspicious
#   code patterns that could compromise the platform, leak secrets, or execute
#   unauthorized operations.
#
# How it works:
#   1. Discovers all Python files in the DAG submission folder.
#   2. Identifies DAG files by detecting DAG instantiation patterns.
#   3. Inlines all local imports recursively into a single source string,
#      so that code hidden in helper modules is also analyzed.
#   4. Compiles the inlined source to leverage the compiler's constant folding,
#      which resolves simple string concatenations into their final values.
#   5. Extracts all string constants from the compiled bytecode.
#   6. Checks imported non-local modules against a list of suspicious modules.
#   7. Checks extracted string constants against a list of suspicious paths
#      and keywords specific to Kubernetes, OpenShift, and Airflow.
#
# What it detects:
#   - Usage of dangerous Python modules (e.g. subprocess, pickle, socket)
#   - Access to sensitive Kubernetes/OpenShift paths (e.g. /var/run/secrets)
#   - Access to sensitive Airflow configuration files (e.g. airflow.cfg)
#   - Obfuscated strings resolved at compile time via constant folding
#
# Limitations:
#   - Cannot detect strings built dynamically at runtime via variables
#   - Cannot analyze compiled extensions (.so files)
#   - Does not execute the code — purely static analysis
#   - Suspicious module and string lists are not exhaustive
#
# Context:
#   Designed for an Airflow platform running on OpenShift/Kubernetes.
#   Client DAG submissions are expected to contain only local Python files
#   with no third-party dependencies beyond the platform approved list.
# =============================================================================
import ast
import os
import re
import types
from pathlib import Path

DEFAULT_ENCODING = "utf-8"

# KEY is the module name and VALUE is the reason why is it considered risky
SUSPICIOUS_MODULE_NAMES: dict[str, str] = {
    "pickle": "arbitrary code execution on unpickle",
    "ctypes": "call C functions, access raw memory",
    "marshal": "low-level serialization, code execution",
    "subprocess": "spawn system processes",
    "runpy": "execute arbitrary Python files",
    "importlib": "dynamic module loading",
    "socket": "raw TCP/UDP connections",
    "requests": "HTTP calls to external servers",
    "urllib": "HTTP/FTP connections",
    "base64": "encode/decode obfuscated payloads",
    "codecs": "encode/decode with various codecs",
    "binascii": "binary/ASCII payload conversions",
    "zlib": "compress/decompress hidden payloads",
    "gzip": "compress/decompress hidden payloads",
    "bz2": "compress/decompress hidden payloads",
    "lzma": "compress/decompress hidden payloads",
    "shelve": "pickle-based storage",
    "dill": "extended pickle, code execution",
    "jsonpickle": "JSON + pickle deserialization",
    "paramiko": "SSH connections",
    "aiohttp": "async HTTP calls",
    "asyncio": "async network connections",
    "ftplib": "FTP connections",
    "smtplib": "send emails for data exfiltration",
    "telnetlib": "telnet connections",
    "xmlrpc": "remote procedure calls",
    "http": "raw HTTP connections",
    "cffi": "C foreign function interface",
    "os": "system calls, file operations",
    "sys": "sys.modules manipulation, path injection",
    "shutil": "file copy/move/delete operations",
    "pathlib": "file read/write/delete",
    "glob": "file system discovery",
    "tempfile": "create files in temp locations",
    "fileinput": "modify files in place",
    "threading": "background threads to bypass controls",
    "multiprocessing": "spawn new processes",
    "concurrent": "thread/process pools",
    "inspect": "read source code of any object",
    "gc": "access any object in memory",
    "weakref": "access objects bypassing references",
    "traceback": "expose internal stack info",
    "psutil": "process inspection and killing",
    "signal": "send signals to processes",
    "resource": "system resource manipulation",
    "pty": "spawn shell sessions",
    "popen2": "legacy process spawning",
    "commands": "legacy os.system wrapper",
    "code": "interactive interpreter execution",
    "codeop": "dynamic code compilation",
    "cryptography": "encrypt/decrypt to hide payloads",
    "Crypto": "PyCryptodome, encrypt payloads",
    "hashlib": "hide content behind hashes",
    "dotenv": "read .env files with secrets",
    "configparser": "read config files with secrets",
}

# KEY is the string in lower case and VALUE is the reason why is it considered risky
SUSPICIOUS_STRINGS: dict[str, str] = {
    "/var/run/secrets": "Kubernetes secrets mount path",
    "/run/secrets": "Kubernetes secrets mount path",
    "/vault/secrets": "HashiCorp Vault secrets mount path",
    "/var/run/secrets/kubernetes.io": "Kubernetes service account token and CA",
    "/run/secrets/kubernetes.io": "Kubernetes service account credentials",
    "/var/run/secrets/openshift.io": "OpenShift specific secrets",
    "/etc/kubernetes": "Kubernetes cluster configuration",
    "/etc/origin": "OpenShift origin configuration",
    "serviceaccount/token": "Kubernetes service account JWT token",
    "serviceaccount/ca.crt": "Kubernetes cluster CA certificate",
    "serviceaccount/namespace": "Kubernetes namespace identifier",
    "kubeconfig": "Kubernetes cluster access configuration",
    ".kube/config": "Kubernetes CLI configuration with credentials",
    "/opt/airflow/airflow.cfg": "Airflow main configuration file with DB credentials",
    "/opt/airflow/webserver_config.py": "Airflow webserver config with auth settings",
    "airflow.cfg": "Airflow configuration file",
    "/opt/airflow/logs": "Airflow logs, may contain sensitive data",
    "/opt/airflow/dags": "Airflow DAGs folder",
    "/opt/airflow/plugins": "Airflow plugins folder",
    "airflow_db": "Airflow metadata database",
}


def inline_imports(
    filepath: str | Path,
    base_dir: str | Path,
    visited: set[str] | None = None,
    non_local_modules_import: set[str] | None = None,
) -> str:
    """
    Recursively inline all local Python imports of a given file into a single source string.

    For each import statement found in the file, if the imported module resolves to a local
    .py file within base_dir, its source code is inlined in place of the import statement.
    If the module is not local (e.g. stdlib or third-party), it is added to the
    non_local_modules_import set for later analysis.

    Args:
        filepath: Path to the Python file to inline.
        base_dir: Root directory of the DAG submission. Used to resolve local imports.
        visited: Set of already-visited file paths to prevent circular import loops.
                 Initialized automatically on first call.
        non_local_modules_import: Set collecting all non-local module names encountered
                                  during recursion. Initialized automatically on first call.

    Returns:
        A single string containing the fully inlined Python source code.

    Raises:
        SyntaxError: If any file in the import chain contains invalid Python syntax.
    """
    if non_local_modules_import is None:
        non_local_modules_import = set()
    if visited is None:
        visited = set()

    if filepath in visited:
        return ""
    visited.add(str(filepath))

    with open(filepath, mode="r", encoding=DEFAULT_ENCODING) as f:
        source = f.read()

    try:
        tree = ast.parse(source)
    except SyntaxError as e:
        raise SyntaxError(f"Syntax error in file '{filepath}' at line {e.lineno}: {e.msg}") from e

    replacements: dict[int, str] = {}
    last_inlined_module_name: str = ""

    for node in ast.walk(tree):
        if isinstance(node, (ast.Import, ast.ImportFrom)):
            module_name = node.module if isinstance(node, ast.ImportFrom) else node.names[0].name
            if module_name is None:
                continue
            local_path = os.path.join(base_dir, module_name.replace(".", "/") + ".py")
            if os.path.exists(local_path):
                inlined_source = inline_imports(local_path, base_dir, visited, non_local_modules_import)
                replacements[node.lineno] = inlined_source
                last_inlined_module_name = module_name
            else:
                non_local_modules_import.add(module_name)

    lines = source.splitlines()
    result = []
    for i, line in enumerate(lines, start=1):
        if i in replacements:
            module_label = last_inlined_module_name
            result.append(f"\n# ---- inlined from {module_label} ----")
            result.append(replacements[i])
            result.append(f"# ---- end of {module_label} ----\n")
        else:
            result.append(line)

    return "\n".join(result)


def is_dag_file(filepath: str | Path) -> bool:
    """
    Determine whether a Python file is an Airflow DAG file.

    Detection is based on the presence of a DAG instantiation pattern
    (e.g. `with DAG(` or `dag = DAG(`) in the source code.

    Args:
        filepath: Path to the Python file to inspect.

    Returns:
        True if the file contains at least one DAG instantiation, False otherwise.
    """
    dag_detection_regex = r"\s+DAG\s*\("
    with open(filepath, mode="r", encoding=DEFAULT_ENCODING) as f:
        source = f.read()
    return len(re.findall(dag_detection_regex, source)) > 0


def extract_strings(code_object: types.CodeType) -> list[str]:
    """
    Recursively extract all string constants from a compiled Python code object.

    Traverses the code object and all nested code objects (functions, classes,
    lambdas) to collect every string constant found in co_consts. This allows
    detection of string patterns even inside nested scopes.

    The compiler performs constant folding, meaning simple string concatenations
    such as "/vault/" + "secret" are resolved to "/vault/secret" at compile time
    and will appear as a single constant here.

    Args:
        code_object: A compiled Python code object, typically obtained via compile().

    Returns:
        A flat list of all string constants found in the code object and its children.
    """
    strings: list[str] = []
    for const in code_object.co_consts:
        if isinstance(const, str):
            strings.append(const)
        elif isinstance(const, types.CodeType):
            strings.extend(extract_strings(const))
    return strings


if __name__ == "__main__":
    dags_repo_folder = "/tmp/dags"
    warnings_count = 0

    py_files = list(Path(dags_repo_folder).rglob("*.py"))

    for py_file in py_files:
        if is_dag_file(py_file):
            print(f"== Analyse DAG file: {py_file.name}")
            non_local_modules_import_detected: set[str] = set()

            # Inline all local imports into a single source string.
            # Non-local module names (stdlib, third-party) are collected in the set.
            inlined_dag_content = inline_imports(
                py_file,
                dags_repo_folder,
                non_local_modules_import=non_local_modules_import_detected,
            )

            # Compile the inlined source. The compiler performs constant folding,
            # resolving simple string concatenations into their final values,
            # which improves the accuracy of string-based pattern detection.
            inlined_dag_compiled = compile(inlined_dag_content, f"inlined_{py_file.name}", "exec")

            # Extract all string constants from the compiled code object.
            strings = extract_strings(inlined_dag_compiled)
            strings_lowered = [s.lower().strip() for s in strings]

            # Check non-local imported modules against the suspicious module list.
            for module_name in non_local_modules_import_detected:
                if module_name in SUSPICIOUS_MODULE_NAMES:
                    reason = SUSPICIOUS_MODULE_NAMES[module_name]
                    print(f"[!] Suspicious module '{module_name}' detected: {reason}.")
                    warnings_count += 1

            # Check extracted string constants against the suspicious string list.
            for string in strings_lowered:
                if string in SUSPICIOUS_STRINGS:
                    reason = SUSPICIOUS_STRINGS[string]
                    print(f"[!] Suspicious string '{string}' detected: {reason}.")
                    warnings_count += 1

    print(f"{warnings_count} suspicious element(s) detected.")