from frostfs_testlib.shell.interfaces import CommandInspector


class SudoInspector(CommandInspector):
    """Prepends command with sudo.

    If command is already prepended with sudo, then has no effect.
    """

    def inspect(self, original_command: str, command: str) -> str:
        if not command.startswith("sudo"):
            return f"sudo {command}"
        return command


class SuInspector(CommandInspector):
    """Allows to run command as another user via sudo su call

    If command is already prepended with sudo su, then has no effect.
    """

    def __init__(self, user: str) -> None:
        self.user = user

    def inspect(self, original_command: str, command: str) -> str:
        if not original_command.startswith("sudo su"):
            cmd = original_command.replace('"', '\\"').replace("\$", "\\\\\\$")
            return f'sudo su - {self.user} -c "{cmd}"'
        return original_command