stack_traces.py

import inspect
from itertools import chain
import os
from typing import Any
from typing import Iterable
from typing import Optional

from ddtrace._trace.span import Span
from ddtrace.appsec._constants import STACK_TRACE
from ddtrace.internal import span_bus
from ddtrace.internal.settings.asm import config as asm_config


_INTERNAL_FRAMES = tuple(
    os.path.sep + "ddtrace" + os.path.sep + sub + os.path.sep for sub in ("contrib", "internal", "appsec")
) + (os.path.sep + "wrapt" + os.path.sep,)


def _get_stack_span() -> Optional[Span]:
    span = span_bus.get_span()
    if span is not None:
        return span._service_entry_span
    return span_bus.get_root_span()


def report_stack(
    message: Optional[str] = None,
    span: Optional[Span] = None,
    crop_stack: Optional[str] = None,
    stack_id: Optional[str] = None,
    namespace: str = STACK_TRACE.RASP,
) -> bool:
    """Report an AppSec stack trace to a span."""
    if not asm_config._ep_stack_trace_enabled:
        return False
    if namespace == STACK_TRACE.RASP and not (asm_config._asm_enabled and asm_config._ep_enabled):
        return False
    if namespace == STACK_TRACE.IAST and not asm_config._iast_enabled:
        return False

    if namespace == STACK_TRACE.IAST and asm_config._iast_use_root_span:
        span = span_bus.get_root_span()
    elif span is None:
        span = _get_stack_span()

    if span is None or stack_id is None:
        return False
    appsec_traces = span._get_struct_tag(STACK_TRACE.TAG) or {}
    current_list = appsec_traces.get(namespace, [])
    total_length = len(current_list)

    if asm_config._ep_max_stack_traces and total_length >= asm_config._ep_max_stack_traces:
        return False

    stack = inspect.stack()
    crop_index = 0
    if crop_stack is not None:
        for i, frame in enumerate(stack):
            if frame.frame.f_code.co_name == crop_stack:
                crop_index = i + 1
                break
        while crop_index < len(stack) and any(d in stack[crop_index].filename for d in _INTERNAL_FRAMES):
            crop_index += 1
    if crop_index:
        stack = stack[crop_index:]
    res: dict[str, Any] = {
        "language": "python",
        "id": stack_id,
    }
    if message is not None:
        res["message"] = message
    if len(stack) > asm_config._ep_max_stack_trace_depth > 0:
        top_stack = int(asm_config._ep_max_stack_trace_depth * asm_config._ep_stack_top_percent / 100)
        bottom_stack = asm_config._ep_max_stack_trace_depth - top_stack
        iterator: Iterable[int] = chain(range(top_stack), range(len(stack) - bottom_stack, len(stack)))
    else:
        iterator = range(len(stack))
    frames = [
        {
            "id": i,
            "function": getattr(stack[i].frame.f_code, "co_qualname", stack[i].frame.f_code.co_name),
            "file": stack[i].filename,
            "line": stack[i].lineno,
        }
        for i in iterator
    ]
    res["frames"] = frames
    current_list.append(res)
    appsec_traces[namespace] = current_list
    span._set_struct_tag(STACK_TRACE.TAG, appsec_traces)
    return True