doczyai-pipelines/archive/terminal.py

#!/usr/bin/env python3
import os
import re
import subprocess
import logging
import signal
from typing import Callable, Optional, Dict, Tuple, List

LOGGER = logging.getLogger(__name__)
LOGGER.setLevel(logging.INFO)
handler = logging.StreamHandler()
handler.setFormatter(
    logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
)
LOGGER.addHandler(handler)


class SensitiveFormatter(logging.Formatter):
    """Formatter that removes sensitive information from logs."""

    @staticmethod
    def mask_password(log: str) -> str:
        return re.sub(r"password=([^\s]+)", r"password=*****", log)

    @staticmethod
    def mask_api_key(log: str) -> str:
        return re.sub(r"api_key=([^\s]+)", r"api_key=*****", log)

    @staticmethod
    def _mask(s: str) -> str:
        filtered = SensitiveFormatter.mask_password(s)
        filtered = SensitiveFormatter.mask_api_key(filtered)
        return filtered

    def format(self, record) -> str:
        original = super().format(record)
        return self._mask(original)


def prepare_logger() -> logging.Logger:
    global LOGGER
    if LOGGER is not None:
        return LOGGER
    LOGGER = logging.getLogger(__name__)
    LOGGER.setLevel(logging.INFO)
    log_format = "%(asctime)s %(filename)s:%(lineno)-4s [%(levelname)s] %(message)s"
    handler = logging.StreamHandler()
    handler.setFormatter(SensitiveFormatter(log_format))
    LOGGER.addHandler(handler)
    return LOGGER


def get_blue_shade(log_prefix: str) -> str:
    """Returns an ANSI color code for a shade of blue based on the hash of the log_prefix."""
    blue_shades = [81, 87, 117, 153, 159, 195, 111, 45, 39]
    hash_value = hash(log_prefix)
    blue_index = hash_value % len(blue_shades)
    return f"\033[38;5;{blue_shades[blue_index]}m"


def signal_handler(sig, frame, process):
    print("Ctrl+C pressed! Sending SIGTERM to Terraform process...")
    process.terminate()


def run_command(
    command: str,
    log_output: bool = False,
    decorate_logs: bool = True,
    log_cmd: bool = False,
    log_prefix: str = "",
    envs: Optional[str] = None,
    secrets: Optional[Dict[str, str]] = None,
    failure_callback: Optional[Callable[[str], None]] = None,
    cwd: Optional[str] = None,
) -> Tuple[int, List[str]]:
    new_env = os.environ.copy()
    if secrets:
        for key, value in secrets.items():
            new_env[key] = os.path.expandvars(value)

    full_command = f"{envs} {command}" if envs else command
    process = subprocess.Popen(
        full_command,
        shell=True,
        cwd=cwd,
        env=new_env,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
    )

    # Setup signal handler
    original_sigint_handler = signal.getsignal(signal.SIGINT)
    signal.signal(signal.SIGINT, lambda sig, frame: signal_handler(sig, frame, process))

    log_prefix_env = os.getenv("LOG_PREFIX", "")
    log_prefix_full = (
        f"[{log_prefix_env}] {log_prefix}" if log_prefix_env else log_prefix
    )

    if log_cmd:
        LOGGER.info(
            f"{get_blue_shade(log_prefix_full)}{log_prefix_full}\033[0m: Running command: \033[1;34m{command}\033[0m"
        )
        if envs:
            LOGGER.info(f"Using additional envs: \033[1;34m{envs}\033[0m")
        if secrets:
            LOGGER.info(
                f"Using additional environment variables with secrets: \033[1;34m{list(secrets.keys())}\033[0m"
            )

    cmd_output = []
    while True:
        output = process.stdout.readline() if process.stdout else ""
        if output:
            output_line = output.strip()
            if log_output:
                log_statement = (
                    (
                        f"{get_blue_shade(log_prefix_full)}{log_prefix_full}\033[0m: "
                        f"{output_line}"
                    )
                    if log_prefix
                    else output_line
                )
                LOGGER.info(log_statement) if decorate_logs else print(log_statement)
            cmd_output.append(output_line)
        elif process.poll() is not None:
            break

    # Restore original signal handler
    signal.signal(signal.SIGINT, original_sigint_handler)

    rc = process.poll() or 0  # Ensure rc is an int, default to 0 if None
    if rc != 0 and failure_callback:
        failure_callback(f"Command returned code: {rc}")

    return rc, cmd_output


def decorate_successful(text: str) -> str:
    return f"\033[32;1m{text}\033[0m"


def decorate_white_bold(text: str) -> str:
    return f"\033[37;1m{text}\033[0m"


def decorate_warn(text: str) -> str:
    return f"\033[33;1m{text}\033[0m"