Das-OS/srcpkgs/das-shell/files/python/das_shell/banner.py

"""Login banner model of dbanner: a grid of styled characters, fields and widgets.

A design is JSON:
    {"name": "...", "title": {"uk": ..., "en": ...}, "local": false,
     "vars": {"url": "https://..."},
     "styles": {"a": {"fg": "blue", "bg": null, "bold": false}, ...},
     "lines": ["####", ...], "style_map": ["aaaa", ...],
     "fields": [{"x": 4, "y": 1, "width": 20, "align": "left", "text": "{ip}", "style": "b"}],
     "widgets_position": "below", "widgets": [{"type": "builtin", "name": "disks"}, ...]}

style_map has one style letter per character of lines; a space is the
terminal's own style. This module only needs the standard library: the
banner is printed at every SSH login.
"""

import copy
import json
import os
import re
import subprocess
import unicodedata

from . import DATA_DIR, ETC_DIR, colors, config_dir
from .i18n import UK, title

STYLE_KEYS = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
ALIGNS = ("left", "center", "right")
SYSTEM_FILE = os.path.join(ETC_DIR, "banner.json")
OFF_FILE = os.path.join(ETC_DIR, "banner.off")


def char_width(ch):
    return 2 if unicodedata.east_asian_width(ch) in ("W", "F") else 1


class Grid:
    """Characters with a style key each; the canvas of the editor."""

    def __init__(self, width=40, height=8, styles=None):
        self.width, self.height = width, height
        self.cells = [[[" ", " "] for _ in range(width)] for _ in range(height)]
        self.styles = styles if styles is not None else {}

    # --- conversion -------------------------------------------------------
    @classmethod
    def from_design(cls, design):
        lines = design.get("lines") or []
        smap = design.get("style_map") or []
        width = max([len(line) for line in lines] + [1])
        grid = cls(width, max(len(lines), 1), copy.deepcopy(design.get("styles") or {}))
        for y, line in enumerate(lines):
            keys = smap[y] if y < len(smap) else ""
            for x, ch in enumerate(line):
                key = keys[x] if x < len(keys) else " "
                grid.cells[y][x] = [ch, key if key in grid.styles else " "]
        return grid

    def to_design(self, design):
        """Stores the grid into a copy of the design, dropping unused styles."""
        d = copy.deepcopy(design)
        lines, smap = [], []
        for row in self.cells:
            line = "".join(c[0] for c in row)
            keys = "".join(c[1] for c in row)
            # Trailing blank cells carry nothing
            n = len(line)
            while n and line[n - 1] == " " and keys[n - 1] == " ":
                n -= 1
            lines.append(line[:n])
            smap.append(keys[:n])
        while lines and not lines[-1]:
            lines.pop()
            smap.pop()
        used = {k for keys in smap for k in keys} | {f.get("style") for f in d.get("fields", [])}
        d["styles"] = {k: v for k, v in self.styles.items() if k in used}
        d["lines"], d["style_map"] = lines, smap
        d["canvas"] = {"width": self.width, "height": self.height}
        return d

    def copy(self):
        g = Grid(self.width, self.height, copy.deepcopy(self.styles))
        g.cells = [[list(c) for c in row] for row in self.cells]
        return g

    # --- styles -------------------------------------------------------------
    def style_key(self, fg=None, bg=None, bold=False):
        """The key of a style, added when it is new."""
        if not fg and not bg and not bold:
            return " "
        want = {"fg": fg, "bg": bg, "bold": bool(bold)}
        for key, st in self.styles.items():
            if {"fg": st.get("fg"), "bg": st.get("bg"), "bold": bool(st.get("bold"))} == want:
                return key
        for key in STYLE_KEYS:
            if key not in self.styles:
                self.styles[key] = want
                return key
        raise ValueError("too many styles")

    def style(self, key):
        st = self.styles.get(key) or {}
        return st.get("fg"), st.get("bg"), bool(st.get("bold"))

    # --- drawing --------------------------------------------------------------
    def inside(self, x, y):
        return 0 <= x < self.width and 0 <= y < self.height

    def get(self, x, y):
        return self.cells[y][x] if self.inside(x, y) else [" ", " "]

    def put(self, x, y, ch, key):
        if self.inside(x, y):
            self.cells[y][x] = [ch, key]

    def resize(self, width, height):
        width, height = max(1, width), max(1, height)
        cells = [[[" ", " "] for _ in range(width)] for _ in range(height)]
        for y in range(min(height, self.height)):
            for x in range(min(width, self.width)):
                cells[y][x] = self.cells[y][x]
        self.width, self.height, self.cells = width, height, cells

    def text(self, x, y, text, key):
        for i, ch in enumerate(text):
            self.put(x + i, y, ch, key)

    def line(self, x0, y0, x1, y1, ch, key):
        """A straight line; horizontal and vertical ones use ch as given."""
        dx, dy = abs(x1 - x0), -abs(y1 - y0)
        sx, sy = (1 if x0 < x1 else -1), (1 if y0 < y1 else -1)
        err = dx + dy
        while True:
            self.put(x0, y0, ch, key)
            if x0 == x1 and y0 == y1:
                return
            e2 = 2 * err
            if e2 >= dy:
                err += dy
                x0 += sx
            if e2 <= dx:
                err += dx
                y0 += sy

    def box(self, x0, y0, x1, y1, kind, key):
        x0, x1 = sorted((x0, x1))
        y0, y1 = sorted((y0, y1))
        tl, tr, bl, br, h, v = BOXES[kind]
        if x0 == x1 or y0 == y1:
            self.line(x0, y0, x1, y1, v if x0 == x1 else h, key)
            return
        for x in range(x0 + 1, x1):
            self.put(x, y0, h, key)
            self.put(x, y1, h, key)
        for y in range(y0 + 1, y1):
            self.put(x0, y, v, key)
            self.put(x1, y, v, key)
        self.put(x0, y0, tl, key)
        self.put(x1, y0, tr, key)
        self.put(x0, y1, bl, key)
        self.put(x1, y1, br, key)

    def fill(self, x0, y0, x1, y1, ch, key):
        x0, x1 = sorted((x0, x1))
        y0, y1 = sorted((y0, y1))
        for y in range(y0, y1 + 1):
            for x in range(x0, x1 + 1):
                self.put(x, y, ch, key)

    def region(self, x0, y0, x1, y1):
        """A copy of a rectangle as a new grid with the same styles."""
        x0, x1 = sorted((x0, x1))
        y0, y1 = sorted((y0, y1))
        g = Grid(x1 - x0 + 1, y1 - y0 + 1, copy.deepcopy(self.styles))
        for y in range(g.height):
            for x in range(g.width):
                g.cells[y][x] = list(self.get(x0 + x, y0 + y))
        return g

    def stamp(self, x, y, other, transparent=True):
        """Draws another grid here; its styles are merged into ours."""
        keymap = {" ": " "}
        for key in {c[1] for row in other.cells for c in row}:
            if key != " ":
                fg, bg, bold = other.style(key)
                keymap[key] = self.style_key(fg, bg, bold)
        for dy, row in enumerate(other.cells):
            for dx, (ch, key) in enumerate(row):
                if transparent and ch == " " and key == " ":
                    continue
                self.put(x + dx, y + dy, ch, keymap.get(key, " "))


BOXES = {
    "single": "┌┐└┘─│",
    "double": "╔╗╚╝═║",
    "rounded": "╭╮╰╯─│",
    "heavy": "┏┓┗┛━┃",
    "ascii": "++++-|",
    "hash": "######",
    "block": "██████",
}


# --- designs ------------------------------------------------------------------

def normalize(design):
    d = copy.deepcopy(design)
    d.setdefault("name", "custom")
    d.setdefault("local", False)
    d.setdefault("vars", {})
    d.setdefault("styles", {})
    d.setdefault("lines", [])
    d.setdefault("style_map", [])
    d.setdefault("fields", [])
    d.setdefault("widgets", [])
    if d.get("widgets_position") not in ("below", "right"):
        d["widgets_position"] = "below"
    for f in d["fields"]:
        f.setdefault("align", "left")
        f.setdefault("style", " ")
        f.setdefault("width", len(f.get("text", "")))
    return d


def preset_dir():
    return os.path.join(DATA_DIR, "banner", "presets")


def art_dir():
    return os.path.join(DATA_DIR, "banner", "art")


def drafts_dir():
    return os.path.join(config_dir(), "banner")


def _json_names(directory):
    try:
        return sorted(f[:-5] for f in os.listdir(directory) if f.endswith(".json"))
    except OSError:
        return []


def list_presets():
    out = []
    for name in sorted(_json_names(preset_dir()), key=lambda n: (n != "das-classic", n)):
        try:
            out.append((name, title(load_file(os.path.join(preset_dir(), name + ".json")).get("title"), name)))
        except (OSError, ValueError):
            continue
    return out


def list_drafts():
    return [(n, os.path.join(drafts_dir(), n + ".json")) for n in _json_names(drafts_dir())]


def list_art():
    """(name, title, design) of the ready-made pictures and the user's own."""
    out = []
    for directory in (art_dir(), os.path.join(config_dir(), "banner-art")):
        for name in _json_names(directory):
            try:
                d = load_file(os.path.join(directory, name + ".json"))
            except (OSError, ValueError):
                continue
            out.append((name, title(d.get("title"), name), d))
    return out


def load_file(path):
    with open(path, encoding="utf-8") as f:
        return normalize(json.load(f))


def load(spec):
    """A preset name, a draft name or a path to a design file."""
    if os.sep in spec or spec.endswith(".json"):
        return load_file(spec)
    for directory in (preset_dir(), drafts_dir()):
        path = os.path.join(directory, spec + ".json")
        if os.path.exists(path):
            return load_file(path)
    raise FileNotFoundError(spec)


def system_design():
    """The banner shown at login: /etc/das/banner.json or the das-classic preset."""
    try:
        return load_file(SYSTEM_FILE)
    except (OSError, ValueError):
        return load_file(os.path.join(preset_dir(), "das-classic.json"))


def dumps(design):
    return json.dumps(design, ensure_ascii=False, indent=2) + "\n"


def save(design, path):
    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
    tmp = path + ".tmp"
    with open(tmp, "w", encoding="utf-8") as f:
        f.write(dumps(design))
    os.replace(tmp, path)


# --- rendering ------------------------------------------------------------------

_VAR = re.compile(r"\{([A-Za-z_][A-Za-z0-9_]*)\}")
_LANG = re.compile(r"\[\[([^|\]]*)\|([^\]]*)\]\]")
FIGLET_FONTS = ("standard", "small", "big", "slant", "smslant", "mini", "banner", "block",
                "shadow", "smshadow", "lean", "script", "digital")


def expand(text, values):
    """[[Ukrainian|English]] in the user's language, {name} replaced by its
    value; unknown names stay as they are."""
    text = _LANG.sub(lambda m: m.group(1) if UK else m.group(2), text)

    def repl(m):
        value = values.get(m.group(1))
        return m.group(0) if value is None else str(value)
    return _VAR.sub(repl, text)


def figlet(text, font):
    """Big letters by figlet, or the text itself when figlet is missing."""
    try:
        out = subprocess.run(["figlet", "-f", font, "-w", "1000", text], capture_output=True,
                             text=True, timeout=2).stdout
    except (OSError, subprocess.SubprocessError):
        out = ""
    lines = out.rstrip("\n").split("\n") if out.strip() else [text]
    while lines and not lines[-1].strip():
        lines.pop()
    return lines


def fit(text, width, align):
    if len(text) > width:
        return text[:max(width - 1, 0)] + "…" if width else ""
    pad = width - len(text)
    if align == "right":
        return " " * pad + text
    if align == "center":
        return " " * (pad // 2) + text + " " * (pad - pad // 2)
    return text + " " * pad


def field_lines(f, values):
    """The lines a field draws: one, or several for big letters."""
    # An empty value leaves no dangling separator: "{host} · {url}" -> "host"
    text = expand(f.get("text", ""), values).strip(" ·|")
    width, align = int(f.get("width", 0)), f.get("align", "left")
    if f.get("font"):
        big = figlet(text, f["font"])
        if max(len(line) for line in big) <= width:
            # Align the block as a whole, so the letters stay together
            block = max(len(line) for line in big)
            return [fit(fit(line, block, "left"), width, align) for line in big]
        return [fit(text, width, align)]
    return [fit(text, width, align)]


def draw_fields(grid, design, values):
    for f in design.get("fields", []):
        key = f.get("style", " ")
        x, y = int(f.get("x", 0)), int(f.get("y", 0))
        for i, line in enumerate(field_lines(f, values)):
            # Spaces of a field do not cover the picture under it
            for j, ch in enumerate(line):
                if ch != " " or not f.get("transparent", True):
                    grid.put(x + j, y + i, ch, key if key in grid.styles else " ")


def grid_lines(grid, color=True):
    """The grid as terminal lines, trailing blanks removed."""
    out = []
    for row in grid.cells:
        n = len(row)
        while n and row[n - 1][0] == " " and (row[n - 1][1] == " " or not grid.style(row[n - 1][1])[1]):
            n -= 1
        line, current = [], " "
        for ch, key in row[:n]:
            if color and key != current:
                line.append(colors.sgr(*grid.style(key)) if key != " " else colors.RESET)
                current = key
            line.append(ch)
        if color and current != " ":
            line.append(colors.RESET)
        out.append("".join(line))
    while out and not out[-1].strip():
        out.pop()
    return out


def visible_len(text):
    return sum(char_width(c) for c in re.sub(r"\x1b\[[0-9;?]*[A-Za-z]", "", text))


def render(design, values, widget_blocks, color=True):
    """The whole banner as text: the picture with fields, then the widgets."""
    design = normalize(design)
    grid = Grid.from_design(design)
    if design.get("canvas"):
        grid.resize(max(grid.width, design["canvas"].get("width", 0)), grid.height)
    draw_fields(grid, design, values)
    art = grid_lines(grid, color)
    blocks = [b for b in widget_blocks if b]
    if not color:
        blocks = [[re.sub(r"\x1b\[[0-9;?]*[A-Za-z]", "", line) for line in b] for b in blocks]
    widget_lines = [line for b in blocks for line in b]
    if design["widgets_position"] == "right" and art:
        width = max(visible_len(line) for line in art) + 3
        out = []
        for i in range(max(len(art), len(widget_lines))):
            left = art[i] if i < len(art) else ""
            right = widget_lines[i] if i < len(widget_lines) else ""
            out.append((left + " " * (width - visible_len(left)) + right) if right else left)
        return "\n".join(out) + "\n"
    out = list(art)
    if widget_lines:
        if out:
            out.append("")
        out += widget_lines
    return "\n".join(out) + "\n"