violin/plugins/violin_guard/bash_ast.py

"""Bash AST parsing powered by bashlex."""

from __future__ import annotations

import shlex
from dataclasses import dataclass, field
from typing import Any

import bashlex


@dataclass
class CommandSegment:
    raw_text: str
    words: list[str] = field(default_factory=list)
    executable: str = ""
    redirects: list[str] = field(default_factory=list)


class _CommandVisitor:
    """Traverse bashlex AST nodes to extract command segments and word tokens."""

    def __init__(self, command: str):
        self.command = command
        self.segments: list[CommandSegment] = []
        self.words: list[str] = []

    def visit(self, node: Any) -> None:
        kind = getattr(node, "kind", None)
        if kind == "command":
            start, end = node.pos
            segment_text = self.command[start:end]
            words = self._collect_words(node)
            executable = self._extract_executable(words)
            redirects = [
                str(getattr(getattr(part, "output", None), "word", ""))
                for part in getattr(node, "parts", [])
                if getattr(part, "kind", None) == "redirect"
                and getattr(getattr(part, "output", None), "word", None)
            ]
            self.segments.append(
                CommandSegment(
                    raw_text=segment_text,
                    words=words,
                    executable=executable,
                    redirects=redirects,
                )
            )

        if hasattr(node, "parts"):
            for child in node.parts:
                self.visit(child)
        if hasattr(node, "command") and getattr(node, "command", None):
            self.visit(node.command)
        if hasattr(node, "list") and getattr(node, "list", None):
            for item in getattr(node, "list", []):
                self.visit(item)

    def _collect_words(self, node: Any) -> list[str]:
        words: list[str] = []

        def collect(n):
            kind = getattr(n, "kind", None)
            if kind == "word" and hasattr(n, "word"):
                words.append(n.word)
                self.words.append(n.word)
            elif kind == "redirect":
                output = getattr(n, "output", None)
                if output is not None and hasattr(output, "word"):
                    words.append(output.word)
                    self.words.append(output.word)
            if hasattr(n, "parts"):
                for child in n.parts:
                    collect(child)
            if hasattr(n, "command") and getattr(n, "command", None):
                collect(n.command)
            if hasattr(n, "list") and getattr(n, "list", None):
                for item in getattr(n, "list", []):
                    collect(item)

        collect(node)
        return words

    @staticmethod
    def _extract_executable(words: list[str]) -> str:
        for word in words:
            if "=" in word and not word.startswith("-") and not word.startswith("/"):
                continue
            if word.lower() in {"command", "env", "exec", "nice", "sudo", "timeout"}:
                continue
            base = word.rsplit("/", 1)[-1].rsplit("\\", 1)[-1].lower()
            return base
        return ""


def parse_bash_segments(command: str) -> list[CommandSegment]:
    """Parse shell command into AST segments using bashlex."""
    if not command or not command.strip():
        return []
    try:
        nodes = bashlex.parse(command)
        visitor = _CommandVisitor(command)
        for node in nodes:
            visitor.visit(node)
        if visitor.segments:
            return visitor.segments
    except Exception:
        pass

    words = command.split()
    exec_name = _CommandVisitor._extract_executable(words)
    return [CommandSegment(raw_text=command, words=words, executable=exec_name)]


def extract_all_command_words(command: str) -> list[str]:
    """Extract all word tokens across subcommands, pipelines, and subshells via bashlex AST."""
    if not command or not command.strip():
        return []
    try:
        nodes = bashlex.parse(command)
        visitor = _CommandVisitor(command)
        for node in nodes:
            visitor.visit(node)
        if visitor.words:
            all_words: list[str] = []
            for w in visitor.words:
                cleaned = w.strip("'\"`")
                if cleaned:
                    all_words.append(cleaned)
                    if any(char in cleaned for char in (" ", ";", "|", "&", ">", "<")):
                        for sub in cleaned.replace(";", " ").replace("|", " ").split():
                            sub_clean = sub.strip("'\"`")
                            if sub_clean:
                                all_words.append(sub_clean)
            return list(dict.fromkeys(all_words))
    except Exception:
        pass

    try:
        return list(dict.fromkeys(shlex.split(command, posix=True)))
    except ValueError:
        return list(dict.fromkeys(command.split()))