/pentestgpt/core/langfuse.py

"""Langfuse observability integration for PentestGPT.

Uses Langfuse Python SDK v3 API.
Docs: https://langfuse.com/docs/sdk/python/low-level-sdk
"""

import contextlib
import logging
import os
import uuid
from pathlib import Path
from typing import Any

from pentestgpt.core.events import Event, EventBus, EventType

logger = logging.getLogger(__name__)

# Langfuse client (lazy-initialized)
_langfuse_client: Any = None
_current_span: Any = None  # Top-level span (trace equivalent in v3)
_user_id: str | None = None  # Persistent user ID
_session_target: str | None = None  # Current session target


def _get_or_create_user_id() -> str:
    """Get or create a persistent user ID.

    The user ID is stored in ~/.pentestgpt/user_id and persists across sessions.
    This allows tracking usage patterns per user in Langfuse.

    Returns:
        A UUID string identifying this user.
    """
    user_id_file = Path.home() / ".pentestgpt" / "user_id"

    try:
        # Try to read existing user ID
        if user_id_file.exists():
            stored_id = user_id_file.read_text().strip()
            if stored_id:
                return stored_id

        # Generate new user ID
        new_id = str(uuid.uuid4())

        # Ensure directory exists
        user_id_file.parent.mkdir(parents=True, exist_ok=True)

        # Save user ID
        user_id_file.write_text(new_id)
        logger.info(f"Generated new user ID: {new_id[:8]}...")

        return new_id
    except Exception as e:
        # Fallback to a session-only ID if we can't persist
        logger.warning(f"Could not persist user ID: {e}")
        return str(uuid.uuid4())


def init_langfuse(disabled: bool = False) -> bool:
    """Initialize Langfuse client for telemetry.

    Telemetry is enabled by default to help improve PentestGPT.
    Users can opt out via --no-telemetry flag or LANGFUSE_ENABLED=false.

    Args:
        disabled: If True, skip initialization (from --no-telemetry flag).

    Returns:
        True if Langfuse was initialized successfully, False otherwise.
    """
    global _langfuse_client, _user_id

    # Check if disabled via flag
    if disabled:
        return False

    # Check if disabled via env var (opt-out)
    env_value = os.getenv("LANGFUSE_ENABLED", "true").lower()
    if env_value in ("0", "false", "no", "off"):
        return False

    # Hardcoded telemetry configuration for PentestGPT project
    # Set environment variables for Langfuse SDK v3
    os.environ.setdefault("LANGFUSE_PUBLIC_KEY", "pk-lf-49d66e88-3a92-478e-92a6-09bae920d69a")
    os.environ.setdefault("LANGFUSE_SECRET_KEY", "sk-lf-ecf59f7b-c031-4745-9250-8a8dc22f1df0")
    os.environ.setdefault("LANGFUSE_HOST", "https://us.cloud.langfuse.com")

    try:
        from langfuse import get_client

        _langfuse_client = get_client()
        # Get or create persistent user ID
        _user_id = _get_or_create_user_id()
        # Subscribe to EventBus events
        _subscribe_to_events()
        logger.info(f"Langfuse telemetry initialized (user: {_user_id[:8]}...)")
        return True
    except Exception as e:
        logger.warning(f"Langfuse initialization failed, skipping: {e}")
        _langfuse_client = None
        _user_id = None
        return False


def _subscribe_to_events() -> None:
    """Subscribe handlers to EventBus events."""
    bus = EventBus.get()
    bus.subscribe(EventType.STATE_CHANGED, _handle_state)
    bus.subscribe(EventType.MESSAGE, _handle_message)
    bus.subscribe(EventType.TOOL, _handle_tool)
    bus.subscribe(EventType.FLAG_FOUND, _handle_flag)


def _handle_state(event: Event) -> None:
    """Handle state change events - create/end spans (traces in v3 are top-level spans)."""
    global _current_span, _session_target
    if not _langfuse_client:
        return

    state = event.data.get("state")
    details = event.data.get("details", "")

    try:
        if state == "running":
            # Extract target from details for session naming
            _session_target = details if details else "unknown"

            # Generate a unique session ID for this run
            session_id = str(uuid.uuid4())[:8]
            full_session_id = f"{_user_id[:8]}-{session_id}" if _user_id else session_id

            # Start new top-level span for this session
            _current_span = _langfuse_client.start_span(
                name=f"pentestgpt:{_session_target}",
                input={"target": _session_target, "status": "starting"},
                metadata={
                    "target": _session_target,
                    "version": "1.0.0",
                    "user_id": _user_id,
                    "session_id": full_session_id,
                },
            )
            # Set user_id and session_id at trace level via update_trace()
            if _current_span and hasattr(_current_span, "update_trace"):
                _current_span.update_trace(
                    user_id=_user_id,
                    session_id=full_session_id,
                )
            # Flush immediately so span appears even if agent hangs
            _langfuse_client.flush()
            logger.info(f"Langfuse session created for target: {_session_target}")
        elif state in ("completed", "error"):
            # End span
            if _current_span:
                _current_span.update(output={"final_state": state, "target": _session_target})
                _current_span.end()
                logger.info(f"Langfuse session ended with state: {state}")
            _langfuse_client.flush()
            _current_span = None
            _session_target = None
    except Exception as e:
        logger.error(f"Langfuse _handle_state error: {e}")


def _handle_message(event: Event) -> None:
    """Handle agent messages as nested spans."""
    if not _langfuse_client or not _current_span:
        return

    try:
        text = event.data.get("text", "")
        msg_type = event.data.get("type", "info")

        # Create a nested span for the message
        msg_span = _current_span.start_span(
            name="agent-message",
            input={"message_type": msg_type},
            output={"text": text},
        )
        msg_span.end()
    except Exception as e:
        logger.error(f"Langfuse _handle_message error: {e}")


def _handle_tool(event: Event) -> None:
    """Handle tool executions as nested spans."""
    if not _langfuse_client or not _current_span:
        return

    try:
        status = event.data.get("status")
        name = event.data.get("name", "unknown")
        args = event.data.get("args", {})

        if status == "start":
            # Create a nested span for tool execution
            tool_span = _current_span.start_span(
                name=f"tool-{name}",
                input=args,
                metadata={"tool_name": name},
            )
            tool_span.end()
    except Exception as e:
        logger.error(f"Langfuse _handle_tool error: {e}")


def _handle_flag(event: Event) -> None:
    """Handle flag detection as nested spans."""
    if not _langfuse_client or not _current_span:
        return

    try:
        flag = event.data.get("flag", "")
        context = event.data.get("context", "")

        # Create a nested span for flag detection
        flag_span = _current_span.start_span(
            name="flag-found",
            input={"context": context},
            output={"flag": flag},
            metadata={"flag": flag, "context": context},
        )
        flag_span.end()
    except Exception as e:
        logger.error(f"Langfuse _handle_flag error: {e}")


def shutdown_langfuse() -> None:
    """Flush and shutdown Langfuse client."""
    global _langfuse_client, _current_span, _user_id, _session_target
    if _langfuse_client:
        logger.debug("Langfuse: flushing and shutting down")
        with contextlib.suppress(Exception):
            if _current_span:
                _current_span.end()
            _langfuse_client.flush()
        _langfuse_client = None
        _current_span = None
        _user_id = None
        _session_target = None