KubeDiagrams/webapp/backend/services/helmfileService.py

"""Service for generating diagrams from Helmfiles."""
import subprocess
import os

from constants import MIME_TYPES
from utils import get_app_logger, log_unexpected_error
from .models import DiagramResult
from .file_manager import FileManager
from .utils import parse_extra_args, has_fatal_error, encode_content, dot_to_dot_json, get_safe_format_extension, redact_temp_paths

logger = get_app_logger(__name__)


def generate_from_helmfile(
    helmfile_content: str,
    output_format: str = "png",
    extra_args: str = "",
    without_namespace: bool = False
) -> DiagramResult:
    """
    Generate a diagram from a Helmfile.

    Args:
        helmfile_content: Contents of the Helmfile
        output_format: Output format
        extra_args: Additional arguments
        without_namespace: Hide namespaces

    Returns:
        DiagramResult: Result of the generation
    """
    safe_ext = get_safe_format_extension(output_format)

    with FileManager.create_temp_file(helmfile_content, suffix=".yaml", mode='wb') as temp_helmfile_path:
        output_path = temp_helmfile_path + f".{safe_ext}"
        dot_output_path = temp_helmfile_path + ".dot" if output_format == "dot_json" else None

        try:
            # Command helmfile template
            template_cmd = ["helmfile", "template", "-f", temp_helmfile_path]
            
            template_proc = subprocess.Popen(
                template_cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True
            )
            helm_output, helm_err = template_proc.communicate()

            if template_proc.returncode != 0 or has_fatal_error("", helm_err):
                FileManager.cleanup_files(output_path, dot_output_path)
                return DiagramResult(
                    success=False,
                    error="Helmfile template failed. See command output below.",
                    command=redact_temp_paths(" ".join(template_cmd), temp_helmfile_path),
                    stdout="",
                    stderr=helm_err or ""
                )

            # Command kube-diagrams
            cmd = ["kube-diagrams", "-", "-o", dot_output_path or output_path]
            if without_namespace:
                cmd.append("--without-namespace")
            if extra_args.strip():
                cmd.extend(parse_extra_args(extra_args, "kube-diagrams"))

            kube_proc = subprocess.run(
                cmd,
                input=helm_output,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True
            )

            stdout_output = kube_proc.stdout or ""
            stderr_output = kube_proc.stderr or ""

            if kube_proc.returncode != 0 or has_fatal_error(stdout_output, stderr_output):
                FileManager.cleanup_files(output_path, dot_output_path)
                return DiagramResult(
                    success=False,
                    error="kube-diagrams failed",
                    command=redact_temp_paths(f"{' '.join(template_cmd)} | {' '.join(cmd)}", temp_helmfile_path, output_path, dot_output_path),
                    stdout=stdout_output,
                    stderr=stderr_output
                )

            if output_format == "dot_json":
                if not os.path.exists(dot_output_path):
                    return DiagramResult(
                        success=False,
                        error=f"Output file not found: {os.path.basename(dot_output_path)}",
                        command=redact_temp_paths(f"{' '.join(template_cmd)} | {' '.join(cmd)}", temp_helmfile_path, output_path, dot_output_path),
                        stdout=stdout_output,
                        stderr=stderr_output
                    )
                if not dot_to_dot_json(dot_output_path, output_path):
                    FileManager.cleanup_files(dot_output_path)
                    return DiagramResult(
                        success=False,
                        error="dot -Tjson conversion failed (is graphviz installed?).",
                        command=redact_temp_paths(f"{' '.join(template_cmd)} | {' '.join(cmd)}", temp_helmfile_path, output_path, dot_output_path),
                        stdout=stdout_output,
                        stderr=stderr_output
                    )
            elif not os.path.exists(output_path):
                return DiagramResult(
                    success=False,
                    error=f"Output file not found: {os.path.basename(output_path)}",
                    command=redact_temp_paths(f"{' '.join(template_cmd)} | {' '.join(cmd)}", temp_helmfile_path, output_path, dot_output_path),
                    stdout=stdout_output,
                    stderr=stderr_output
                )

            content = FileManager.read_file_content(output_path, binary=True)
            encoded = encode_content(content, output_format)

            # Cleaning
            FileManager.cleanup_files(output_path, dot_output_path)

            return DiagramResult(
                success=True,
                diagram=encoded,
                mime_type=MIME_TYPES.get(output_format, "application/octet-stream"),
                filename=f"helmfile-diagram.{output_format}",
                message="Helmfile diagram successfully generated.",
                command=redact_temp_paths(f"{' '.join(template_cmd)} | {' '.join(cmd)}", temp_helmfile_path, output_path, dot_output_path),
                stdout=stdout_output,
                stderr=stderr_output
            )

        except Exception:
            FileManager.cleanup_files(output_path, dot_output_path)
            return DiagramResult(
                success=False,
                error=log_unexpected_error(logger, "generating diagram from Helmfile"),
                command=redact_temp_paths(" ".join(cmd), temp_helmfile_path, output_path, dot_output_path) if 'cmd' in locals() else None
            )