KubeDiagrams/webapp/backend/services/utils.py

"""Utilities for diagram generation services."""
import base64
import json
import os
import re
import shlex
import subprocess

from constants import TEXT_FORMATS, MIME_TYPES
from utils import InputValidator

def redact_temp_paths(text: str, *paths: str) -> str:  
    for path in paths:
        if path:
            text = text.replace(path, os.path.basename(path))
    return text


def has_fatal_error(stdout_txt: str, stderr_txt: str) -> bool:
    """
    Check whether subprocess output contains a fatal error marker.

    Args:
        stdout_txt: Standard output text
        stderr_txt: Standard error text

    Returns:
        bool: True if a fatal error was detected
    """
    return ("error:" in (stdout_txt or "").lower()) or ("error:" in (stderr_txt or "").lower())


_SAFE_FORMAT_EXTENSIONS = {fmt: fmt for fmt in MIME_TYPES}


def get_safe_format_extension(output_format: str) -> str:
    """
    Look up output_format in a hardcoded allowlist and return the matching
    literal extension string.

    Args:
        output_format: Output format to check

    Returns:
        str: The same format string, sourced from a hardcoded mapping

    Raises:
        ValueError: If output_format is not in MIME_TYPES
    """
    try:
        return _SAFE_FORMAT_EXTENSIONS[output_format]
    except KeyError:
        raise ValueError(f"Unsupported output format: {output_format!r}")


def parse_extra_args(extra_args: str, tool: str) -> list[str]:
    """
    Parse a string of extra CLI arguments using shell-like tokenization,
    rejecting any flag not in that tool's allowlist (EXTRA_ARGS_ALLOWED_FLAGS).

    Args:
        extra_args: Space-separated argument string (may include quoted tokens)
        tool: Key into EXTRA_ARGS_ALLOWED_FLAGS identifying the target CLI tool

    Returns:
        list[str]: Parsed argument tokens

    Raises:
        ValueError: If the argument string has invalid shell syntax, or
            contains a flag not allowed for this tool
    """
    if not extra_args or not extra_args.strip():
        return []
    try:
        tokens = shlex.split(extra_args.strip())
    except Exception:
        raise ValueError("Invalid extraArgs: could not parse the value (check for unmatched quotes).")

    bad_flag = InputValidator.find_disallowed_flag(tokens, tool)
    if bad_flag:
        raise ValueError(f"Extra arg flag '{bad_flag}' is not allowed.")

    return tokens


def encode_content(content: bytes, output_format: str) -> str:
    """
    Encode diagram content for JSON transport.

    Text-based formats (SVG, DOT, DOT_JSON, DRAWIO) are decoded as UTF-8.
    Binary formats (PNG, JPG, PDF) are base64-encoded.

    Args:
        content: Raw file content
        output_format: Output format string (e.g. 'png', 'svg')

    Returns:
        str: Encoded content ready for the API response
    """
    if output_format in TEXT_FORMATS:
        return content.decode("utf-8")
    return base64.b64encode(content).decode("utf-8")


_GITHUB_ICONS_BASE = (
    "https://raw.githubusercontent.com/mingrammer/diagrams/refs/heads/master/"
)


def _local_path_to_github_url(path: str) -> str:
    """
    Convert a local kube-diagrams icon path to its GitHub CDN URL.

    The .dot file generated by kube-diagrams references icons via local
    paths (site-packages/resources/...). Running `dot -Tjson` on it preserves
    those local paths as-is, so this restores the GitHub URLs the browser
    needs to load the icons.
    """
    m = re.search(r'resources/.+', path)
    return _GITHUB_ICONS_BASE + m.group() if m else path


def dot_to_dot_json(dot_path: str, dot_json_path: str) -> bool:
    """
    Convert a .dot file into a .dot_json file with layout coordinates via
    `dot -Tjson`, then fix local icon paths into GitHub CDN URLs.

    Going through a single kube-diagrams execution producing a .dot file,
    then converting it locally via `dot -Tjson`, avoids the mismatched
    gvids that occur when kube-diagrams is invoked twice for the same
    diagram (once per output format).

    Args:
        dot_path: Path to the source .dot file.
        dot_json_path: Output path for the enriched .dot_json file.

    Returns:
        bool: True if the conversion succeeded, False otherwise.
    """
    try:
        result = subprocess.run(
            ["dot", "-Tjson", dot_path],
            capture_output=True,
            check=False
        )
        if result.returncode != 0 or not result.stdout:
            return False

        data = json.loads(result.stdout.decode("utf-8"))

        # Fix local paths into GitHub URLs in both image attributes and HTML labels
        for obj in data.get("objects", []):
            if obj.get("image"):
                obj["image"] = _local_path_to_github_url(obj["image"])
            if obj.get("label") and "<img" in str(obj["label"]):
                obj["label"] = re.sub(
                    r'src="([^"]*resources/[^"]+)"',
                    lambda m: f'src="{_local_path_to_github_url(m.group(1))}"',
                    obj["label"],
                )

        with open(dot_json_path, "w", encoding="utf-8") as f:
            json.dump(data, f)
        return True
    except Exception:
        return False


def dot_to_svg(dot_text: str) -> str | None:
    """
    Render DOT source to SVG via `dot -Tsvg`, then fix local icon paths into
    GitHub CDN URLs. dot embeds icons in the SVG as xlink:href pointing to
    the absolute local filesystem path used at generation time, which the
    browser can't resolve.

    Args:
        dot_text: DOT source text.

    Returns:
        str | None: The rendered SVG text, or None if the conversion failed.
    """
    try:
        result = subprocess.run(
            ["dot", "-Tsvg"],
            input=dot_text,
            capture_output=True,
            text=True,
            check=False
        )
        if result.returncode != 0 or not result.stdout:
            return None

        return re.sub(
            r'xlink:href="([^"]*resources/[^"]+)"',
            lambda m: f'xlink:href="{_local_path_to_github_url(m.group(1))}"',
            result.stdout,
        )
    except Exception:
        return None


def enrich_dot_json_with_positions(dot_json_path: str) -> None:
    """
    Enrich a .dot_json file with layout coordinates computed by `dot -Tjson`.

    Strategy:
      1. Read the .dot_json produced by kube-diagrams.
      2. Look for the matching .dot file (same path, .dot extension).
         If missing, skip (coordinates can't be recomputed without it).
      3. Run `dot -Tjson` on the .dot file to get the enriched JSON with pos/bb.
      4. Inject the `pos` (nodes) and `bb` (clusters/graph) fields into the .dot_json.
      5. Rewrite the enriched .dot_json.

    Args:
        dot_json_path: Path to the .dot_json file to enrich (modified in place).
    """
    dot_executable = "dot"
    dot_path = dot_json_path.replace(".dot_json", ".dot")

    if not os.path.exists(dot_path):
        return  # No source .dot file available

    try:
        # 1. Read the existing dot_json
        with open(dot_json_path, "r", encoding="utf-8") as f:
            original = json.load(f)

        # 2. Run dot -Tjson on the .dot file to get coordinates
        result = subprocess.run(
            [dot_executable, "-Tjson", dot_path],
            capture_output=True,
            check=False
        )
        if result.returncode != 0 or not result.stdout:
            return  # dot unavailable or failed

        enriched = json.loads(result.stdout.decode("utf-8"))

        # 3. Build a name -> coordinates mapping from the enriched JSON
        #    (dot -Tjson uses `name` to identify nodes/clusters)
        name_to_coords = {}
        for obj in enriched.get("objects", []):
            name = obj.get("name")
            if not name:
                continue
            coords = {}
            if "pos" in obj:
                coords["pos"] = obj["pos"]
            if "bb" in obj:
                coords["bb"] = obj["bb"]
            if "lp" in obj:
                coords["lp"] = obj["lp"]
            if "lwidth" in obj:
                coords["lwidth"] = obj["lwidth"]
            if "lheight" in obj:
                coords["lheight"] = obj["lheight"]
            if coords:
                name_to_coords[name] = coords

        # Global graph bounding box
        if "bb" in enriched:
            original["bb"] = enriched["bb"]

        # 4. Inject coordinates into the original dot_json (matched by name)
        for obj in original.get("objects", []):
            name = obj.get("name")
            if name and name in name_to_coords:
                obj.update(name_to_coords[name])

        # Edge coordinates: matched by (tail_name, head_name)
        # Build a gvid -> name index for the original nodes
        gvid_to_name = {o["_gvid"]: o.get("name", "") for o in original.get("objects", [])}
        # Index for the enriched JSON
        enriched_edge_map = {}
        for e in enriched.get("edges", []):
            tail_gvid = e.get("tail")
            head_gvid = e.get("head")
            enriched_objs = enriched.get("objects", [])
            tail_name = enriched_objs[tail_gvid]["name"] if tail_gvid is not None and tail_gvid < len(enriched_objs) else None
            head_name = enriched_objs[head_gvid]["name"] if head_gvid is not None and head_gvid < len(enriched_objs) else None
            if tail_name and head_name and "pos" in e:
                enriched_edge_map[(tail_name, head_name)] = e["pos"]

        for edge in original.get("edges", []):
            tail_name = gvid_to_name.get(edge.get("tail"))
            head_name = gvid_to_name.get(edge.get("head"))
            if tail_name and head_name:
                pos = enriched_edge_map.get((tail_name, head_name))
                if pos:
                    edge["pos"] = pos

        # 5. Rewrite the enriched file
        with open(dot_json_path, "w", encoding="utf-8") as f:
            json.dump(original, f)

    except Exception:
        pass  # On error, the file is left without coordinates