DataOps-platform/app/runner/nodes.py

"""Fail-closed Runner node registry and governed executor implementations."""

from __future__ import annotations

import ipaddress
import json
import math
import multiprocessing
import re
from typing import Any, Mapping
from urllib.parse import urlsplit

import requests
from sqlalchemy import text


class NodeExecutionError(ValueError):
    def __init__(self, message, *, commit_outcome="not_applicable"):
        super().__init__(message)
        self.commit_outcome = commit_outcome


def _statement(config):
    statement = config.get("statement")
    if not isinstance(statement, str) or not statement.strip():
        raise NodeExecutionError("SQL statement is required")
    if "\x00" in statement:
        raise NodeExecutionError("SQL statement contains invalid characters")
    return statement.strip()


def _parameters(config, runtime_parameters=None):
    parameters = config.get("parameters", {})
    if not isinstance(parameters, dict):
        raise NodeExecutionError("SQL parameters must be an object")
    runtime_parameters = runtime_parameters or {}
    resolved = {}
    pattern = re.compile(r"^\$\{parameters\.([A-Za-z][A-Za-z0-9_]*)\}$")
    for key, value in parameters.items():
        if isinstance(value, str):
            match = pattern.fullmatch(value)
            if match:
                parameter_name = match.group(1)
                if parameter_name not in runtime_parameters:
                    raise NodeExecutionError(
                        f"runtime parameter {parameter_name} is required"
                    )
                value = runtime_parameters[parameter_name]
            elif "${parameters." in value:
                raise NodeExecutionError(
                    "runtime parameter placeholders must occupy the whole value"
                )
        resolved[key] = value
    return resolved


def _sql_tokens(statement):
    without_strings = re.sub(r"'(?:''|[^'])*'", "''", statement)
    without_comments = re.sub(
        r"/\*.*?\*/|--[^\r\n]*",
        " ",
        without_strings,
        flags=re.DOTALL,
    )
    return re.findall(r"[A-Za-z_]+|;", without_comments.upper())


def _assert_read_only(statement):
    tokens = _sql_tokens(statement)
    if not tokens or tokens[0] not in {"SELECT", "WITH"}:
        raise NodeExecutionError("sql.query accepts read-only statements")
    forbidden = {
        "ALTER",
        "CALL",
        "COPY",
        "CREATE",
        "DELETE",
        "DROP",
        "EXEC",
        "EXECUTE",
        "GRANT",
        "INSERT",
        "INTO",
        "LOCK",
        "MERGE",
        "PROGRAM",
        "REVOKE",
        "TRUNCATE",
        "UPDATE",
    }
    if ";" in tokens or forbidden.intersection(tokens):
        raise NodeExecutionError("sql.query accepts read-only statements")
    if "FOR" in tokens and "UPDATE" in tokens:
        raise NodeExecutionError("sql.query accepts read-only statements")


def _assert_dml(statement):
    tokens = _sql_tokens(statement)
    if (
        not tokens
        or tokens[0] not in {"INSERT", "UPDATE", "DELETE", "MERGE"}
        or ";" in tokens
    ):
        raise NodeExecutionError(
            "sql.execute accepts one parameterized DML statement"
        )


class SqlQueryExecutor:
    def __init__(self, manager, *, max_rows=1000):
        self.manager = manager
        self.max_rows = int(max_rows)

    def execute(self, node, runtime_parameters, **_kwargs):
        config = node.get("config") or {}
        statement = _statement(config)
        _assert_read_only(statement)
        parameters = _parameters(config, runtime_parameters)
        with self.manager.connect(
            node.get("data_source_uid"),
            purpose="dataflow_read",
        ) as connection:
            rows = connection.execute(
                text(statement),
                parameters,
            ).mappings().all()
        bounded = [dict(row) for row in rows[: self.max_rows]]
        return {
            "rows": bounded,
            "row_count": len(bounded),
            "truncated": len(rows) > self.max_rows,
        }


class SqlExecuteExecutor:
    def __init__(self, manager):
        self.manager = manager

    def execute(
        self,
        node,
        runtime_parameters,
        *,
        write_authorized=False,
        **_kwargs,
    ):
        if node.get("purpose") != "write" or not write_authorized:
            raise NodeExecutionError("governed write authorization is required")
        idempotency = node.get("idempotency")
        if (
            not isinstance(idempotency, dict)
            or idempotency.get("strategy")
            not in {"partition_replace", "upsert", "deduplication_key"}
            or not str(idempotency.get("key") or "").strip()
        ):
            raise NodeExecutionError("write idempotency is required")
        config = node.get("config") or {}
        statement = _statement(config)
        _assert_dml(statement)
        try:
            with self.manager.connect(
                node.get("data_source_uid"),
                purpose="dataflow_write",
            ) as connection:
                result = connection.execute(
                    text(statement),
                    _parameters(config, runtime_parameters),
                )
        except Exception as exc:
            commit_outcome = getattr(exc, "commit_outcome", "not_committed")
            raise NodeExecutionError(
                "governed SQL write failed",
                commit_outcome=commit_outcome,
            ) from exc
        return {
            "affected_rows": max(0, int(result.rowcount or 0)),
            "commit_outcome": "committed",
        }


def _restricted_python_child(
    handler,
    arguments,
    result_queue,
    cpu_seconds,
    memory_bytes,
):
    import resource
    import socket

    def set_soft_limit(kind, requested):
        _soft, hard = resource.getrlimit(kind)
        limit = (
            requested
            if hard == resource.RLIM_INFINITY
            else min(requested, hard)
        )
        try:
            resource.setrlimit(kind, (limit, hard))
        except (OSError, ValueError):
            # macOS does not implement RLIMIT_AS; the Linux Runner image does.
            if kind != resource.RLIMIT_AS:
                raise

    set_soft_limit(resource.RLIMIT_CPU, cpu_seconds)
    set_soft_limit(resource.RLIMIT_AS, memory_bytes)
    set_soft_limit(resource.RLIMIT_NOFILE, 32)

    def deny_network(*_args, **_kwargs):
        raise PermissionError("network access is disabled")

    socket.socket = deny_network
    socket.create_connection = deny_network
    try:
        result = handler(arguments)
        json.dumps(result)
        result_queue.put(("success", result))
    except PermissionError:
        result_queue.put(("network_denied", None))
    except (TypeError, ValueError):
        result_queue.put(("invalid_result", None))
    except BaseException:
        result_queue.put(("failed", None))


class RestrictedPythonExecutor:
    """Run only pre-installed, reviewed handlers; arbitrary source is forbidden."""

    def __init__(
        self,
        handlers=None,
        *,
        timeout_seconds=10,
        memory_bytes=512 * 1024 * 1024,
    ):
        self.handlers = dict(handlers or {})
        self.timeout_seconds = float(timeout_seconds)
        self.memory_bytes = int(memory_bytes)
        if self.timeout_seconds <= 0 or self.memory_bytes < 64 * 1024 * 1024:
            raise ValueError("restricted python resource limit is invalid")

    def execute(self, node, _runtime_parameters, **_kwargs):
        config = node.get("config") or {}
        if set(config) - {"handler", "arguments"}:
            raise NodeExecutionError(
                "restricted python accepts only handler and arguments"
            )
        handler_name = config.get("handler")
        handler = self.handlers.get(handler_name)
        if handler is None:
            raise NodeExecutionError("python handler is not allowlisted")
        arguments = config.get("arguments", {})
        if not isinstance(arguments, dict):
            raise NodeExecutionError("python arguments must be an object")
        context = multiprocessing.get_context("fork")
        result_queue = context.Queue(maxsize=1)
        process = context.Process(
            target=_restricted_python_child,
            args=(
                handler,
                arguments,
                result_queue,
                max(1, int(math.ceil(self.timeout_seconds))),
                self.memory_bytes,
            ),
            daemon=True,
        )
        process.start()
        process.join(self.timeout_seconds)
        if process.is_alive():
            process.terminate()
            process.join(1)
            raise NodeExecutionError("python handler exceeded its time limit")
        if result_queue.empty():
            raise NodeExecutionError("python handler failed in the sandbox")
        state, result = result_queue.get_nowait()
        if state == "network_denied":
            raise NodeExecutionError("python handler network access is disabled")
        if state == "invalid_result":
            raise NodeExecutionError(
                "python handler result must be JSON serializable"
            )
        if state != "success":
            raise NodeExecutionError("python handler failed in the sandbox")
        return result


class GovernedHttpExecutor:
    def __init__(
        self,
        *,
        allowed_hosts,
        session=None,
        timeout_seconds=15,
        max_response_bytes=1_000_000,
    ):
        self.allowed_hosts = {
            str(host).strip().lower()
            for host in allowed_hosts
            if str(host).strip()
        }
        self.session = session or requests.Session()
        self.timeout_seconds = int(timeout_seconds)
        self.max_response_bytes = int(max_response_bytes)

    def _validate_url(self, raw_url):
        parts = urlsplit(str(raw_url or ""))
        host = (parts.hostname or "").lower()
        if (
            parts.scheme != "https"
            or not host
            or parts.username is not None
            or parts.password is not None
            or host not in self.allowed_hosts
        ):
            raise NodeExecutionError("HTTP target is not allowlisted")
        try:
            address = ipaddress.ip_address(host)
        except ValueError:
            address = None
        if address and (
            address.is_private
            or address.is_loopback
            or address.is_link_local
            or address.is_reserved
        ):
            raise NodeExecutionError("HTTP target is not allowlisted")
        return parts.geturl()

    def execute(self, node, _runtime_parameters, **_kwargs):
        config = node.get("config") or {}
        unknown = set(config) - {"method", "url", "headers", "json"}
        if unknown:
            raise NodeExecutionError("HTTP node contains unsupported fields")
        method = str(config.get("method", "GET")).upper()
        if method not in {"GET", "POST", "PUT", "PATCH"}:
            raise NodeExecutionError("HTTP method is not allowed")
        headers = config.get("headers", {})
        if not isinstance(headers, dict):
            raise NodeExecutionError("HTTP headers must be an object")
        sensitive = {
            "authorization",
            "cookie",
            "proxy-authorization",
            "x-api-key",
        }
        if sensitive.intersection(str(key).lower() for key in headers):
            raise NodeExecutionError("inline HTTP credentials are forbidden")
        response = self.session.request(
            method,
            self._validate_url(config.get("url")),
            headers=headers,
            json=config.get("json"),
            timeout=self.timeout_seconds,
            allow_redirects=False,
        )
        content = bytes(response.content or "")
        if len(content) > self.max_response_bytes:
            raise NodeExecutionError("HTTP response exceeds the configured limit")
        content_type = str(response.headers.get("content-type", "")).lower()
        body = response.json() if "application/json" in content_type else content.decode(
            "utf-8", errors="replace"
        )
        return {"status_code": int(response.status_code), "body": body}


class NodeRegistry:
    def __init__(self, executors):
        self.executors = dict(executors)

    def execute(
        self,
        node: Mapping[str, Any],
        parameters,
        *,
        write_authorized=False,
    ):
        executor = self.executors.get(node.get("type"))
        if executor is None or not callable(getattr(executor, "execute", None)):
            raise NodeExecutionError("node type is not registered")
        return executor.execute(
            node,
            parameters,
            write_authorized=write_authorized,
        )