cleveragents-core/features/steps/security_pyyaml_dependency_steps.py

"""Step definitions for PyYAML security dependency verification (#9055)."""

from __future__ import annotations

import importlib
import re
from pathlib import Path
from typing import Any

from behave import given, then, when
from packaging.version import parse as parse_version
from behave.runner import Context


@given('the pyproject.toml file exists at "{path}"')
def step_pyproject_exists(context: Context, path: str) -> None:
    """Verify pyproject.toml exists at the given path."""
    pyproject_path = Path(path)
    assert pyproject_path.exists(), f"pyproject.toml not found at {path}"
    context.pyproject_path = pyproject_path


@when("I read the project dependencies from pyproject.toml")
def step_read_project_dependencies(context: Context) -> None:
    """Read the [project.dependencies] list from pyproject.toml."""
    content = context.pyproject_path.read_bytes()
    data: dict[str, Any] = _load_toml(content)
    context.project_dependencies = data.get("project", {}).get("dependencies", [])


@then('the dependency list should include a package matching "{pattern}"')
def step_dependency_matches_pattern(context: Context, pattern: str) -> None:
    """Assert that a dependency in the project dependencies matches the given regex pattern."""
    deps: list[str] = context.project_dependencies
    found = any(re.search(pattern, dep) for dep in deps)
    assert found, (
        f"No dependency matching '{pattern}' found in "
        f"[project.dependencies]. Current deps: {deps}"
    )


@when("I read the PyYAML dependency specification from pyproject.toml")
def step_read_pyyaml_dependency(context: Context) -> None:
    """Read and parse the PyYAML dependency from pyproject.toml."""
    content = context.pyproject_path.read_bytes()
    data: dict[str, Any] = _load_toml(content)
    deps: list[str] = data.get("project", {}).get("dependencies", [])

    # Find the pyyaml dependency entry
    pyyaml_dep = None
    for dep in deps:
        if dep.strip().startswith("pyyaml"):
            pyyaml_dep = dep.strip()
            break

    assert pyyaml_dep is not None, (
        f"PyYAML dependency not found in [project.dependencies]. Current deps: {deps}"
    )

    # Parse the minimum version from the constraint
    # Examples: "pyyaml>=6.0.3", "pyyaml>=6.0", "pyyaml==6.0.2"
    match = re.search(r"(?:>=|==|~>)([\d]+(?:\.[\d]+)*(?:\.\w+)*)", pyyaml_dep)
    assert match is not None, f"Could not parse version from PyYAML dependency: {pyyaml_dep}"

    context.pyyaml_version_spec = pyyaml_dep
    context.pyyaml_min_version = match.group(1)


@then("the minimum version should be at least \"{expected}\"")
def step_pyyaml_minimum_version(context: Context, expected: str) -> None:
    """Assert that the PyYAML minimum version is >= expected version."""
    min_version = context.pyyaml_min_version
    assert parse_version(min_version) >= parse_version(expected), (
        f"PyYAML minimum version {min_version} is less than required {expected}. "
        f"Dependency spec: {context.pyyaml_version_spec}"
    )


@when('I attempt to import the "{module_name}" module')
def step_attempt_import(context: Context, module_name: str) -> None:
    """Attempt to import the given module and record the result."""
    try:
        context.imported_module = importlib.import_module(module_name)
        context.import_error = None
    except ImportError as exc:
        context.imported_module = None
        context.import_error = exc


@then("the import should succeed without errors")
def step_import_succeeded(context: Context) -> None:
    """Assert that the previous import attempt succeeded."""
    assert context.import_error is None, f"Import failed with: {context.import_error}"
    assert context.imported_module is not None, "Module was not imported"


@then('the version constraint in uv.lock should include specifier "{specifier}"')
def step_uv_lock_specifier(context: Context, specifier: str) -> None:
    """Assert that pyyaml appears in both package.dependencies and requires-dist in uv.lock."""
    lock_path = context.pyproject_path.parent / "uv.lock"
    content = lock_path.read_bytes().decode("utf-8")

    # Check that pyyaml is in the cleveragents package dependencies section
    assert (
        '[[package]]\nname = "cleveragents"' in content
    ), "cleveragents package section not found in uv.lock"

    # Verify pyyaml entry exists in file after cleveragents section
    cleveragents_match = re.search(
        r'\[\[package\]\]\s*name = "cleveragents"(.*?)^\[\[package\]\]',
        content,
        re.DOTALL | re.MULTILINE,
    )
    assert cleveragents_match, "Could not find cleveragents package section"

    cleveragents_block = cleveragents_match.group(1)

    assert '{ name = "pyyaml" }' in cleveragents_block, (
        f"pyyaml not found in cleveragents package dependencies in uv.lock"
    )


def _load_toml(content: bytes) -> dict[str, Any]:
    """Load a TOML file from bytes. Uses tomllib if available, else fallback."""
    try:
        import tomllib

        return tomllib.loads(content.decode("utf-8"))
    except ImportError:  # pragma: no cover — Python <3.11
        # Fallback for older Python versions without tomllib
        import tomlkit

        parsed_doc = tomlkit.parse(content.decode("utf-8"))
        return _toml_to_python(parsed_doc)


def _toml_to_python(data: object) -> dict[str, Any]:
    """Convert a TOML document from tomlkit (or tomllib) to plain Python dicts."""
    if isinstance(data, dict):
        def _flatten(d: dict[Any, Any]) -> dict[str, Any]:
            result: dict[str, Any] = {}
            for key, value in d.items():
                if isinstance(value, dict):
                    result[key] = _flatten(value)
                else:
                    result[key] = value
            return result
        return _flatten(data)
    raise TypeError(f"Expected dict-like TOML table, got {type(data).__name__}")