Bridge-and-Join-s/deploy/iso/bj-install-progress

#!/usr/bin/python3
# Экран установки Bridge&Join на первой консоли — вместо технического журнала установщика Ubuntu
# («start: subiquity/Install/install/curtin_install/…»). Запускается из early-commands автоустановки
# (deploy/iso/ubuntu/user-data.tmpl: bootcmd — как можно раньше, early-commands — запасной) через
# bj-install-progress-start отдельной службой с выводом на tty1 и держит консоль за собой: экран установщика
# Ubuntu (snap.subiquity.subiquity-service) останавливает, как только тот появится. Сама установка идёт как шла: этот экран только
# читает её события из журнала (SUBIQUITY_EVENT_TYPE / SUBIQUITY_CONTEXT_NAME) и состояние сервера
# установщика (/run/subiquity/server-state).
#
# Ошибка установки — экран говорит об этом и возвращает стандартный экран установщика Ubuntu (он
# показывает подробности и открывает командную строку); то же, если сам этот экран упадёт
# (обёртка в early-commands).
#
# Шрифт консоли установщика (subiquity.psf, 512 знаков): есть кириллица, рамки, ✓ ▶ ✗ █ ░ · •;
# нет ○ ◐ ⚠ ▓ — их не использовать.
#
# Проверка без установки: bj-install-progress --demo <файл событий JSON> [--html кадр.html]

import json
import math
import os
import select
import shutil
import subprocess
import sys
import time

STAGES = [
    # заголовок, вес в общей полосе, типичная длительность этапа (с) — для плавной полосы внутри этапа
    ("Подготовка", 4, 25),
    ("Разметка диска", 4, 20),
    ("Копирование системы на диск", 34, 180),
    ("Настройка системы и загрузчика", 32, 180),
    ("Пакеты и обновления безопасности", 18, 150),
    ("Компоненты Bridge&Join", 8, 20),
]
DONE_STAGE = len(STAGES)
STATE_FILE = os.environ.get("BJ_PROGRESS_STATE", "/run/subiquity/server-state")  # переменная — только для проверки
CLIENT_UNIT = "snap.subiquity.subiquity-service.service"

# цвета консоли Linux
RESET, BOLD, DIM = "\033[0m", "\033[1m", "\033[2m"
GREEN, CYAN, RED, WHITE, GREY = "\033[32m", "\033[36m", "\033[31m", "\033[97m", "\033[90m"


def stage_of(name, message, state):
    """Этап по событию установщика (имя контекста и сообщение) и состоянию сервера; -1 — не меняет этап."""
    n, m = name or "", message or ""
    if state == "DONE":
        return DONE_STAGE
    if state == "LATE_COMMANDS" or n.startswith("subiquity/Late/run"):
        return 5
    if state == "UU_RUNNING" or "/postinstall" in n:
        return 4
    if "stage-curthooks" in n or "curthooks step" in m or "swap step" in m:
        return 3
    if "stage-extract" in n or "extract step" in m:
        return 2
    if "stage-partitioning" in n or "partitioning step" in m or "formatting step" in m:
        return 1
    return -1


class Progress:
    def __init__(self):
        self.stage = 0
        self.stage_started = time.time()
        self.started = time.time()
        self.state = ""
        self.error = False
        self.detail = ""

    def feed(self, ev):
        """Событие журнала (JSON journalctl)."""
        name = ev.get("SUBIQUITY_CONTEXT_NAME", "")
        msg = (ev.get("MESSAGE") or "").strip()
        etype = ev.get("SUBIQUITY_EVENT_TYPE", "")
        if etype == "error":
            self.detail = msg[:200]
        st = stage_of(name, msg, "")
        self.advance(st)
        if self.stage == 4 and etype == "start":
            if "/install_" in name and ": installing " in msg:
                self.detail = "ставится " + msg.split(": installing ", 1)[1]
            elif "updates" in msg:
                self.detail = "загрузка и установка обновлений безопасности"

    def advance(self, st):
        if st > self.stage:
            self.stage, self.stage_started, self.detail = st, time.time(), ""

    def set_state(self, state):
        self.state = state
        if state == "ERROR":
            self.error = True
        self.advance(stage_of("", "", state))
        if self.stage == 0 and not self.detail.startswith("ставится"):
            self.detail = "запуск установщика…" if state in ("", "STARTING_UP", "CLOUD_INIT_WAIT", "EARLY_COMMANDS") else ""

    def percent(self):
        if self.stage >= DONE_STAGE:
            return 100
        done = sum(w for _, w, _ in STAGES[: self.stage])
        _, w, tau = STAGES[self.stage]
        f = min(0.95, 1 - math.exp(-(time.time() - self.stage_started) / tau))
        return min(99, int(done + w * f))


def mmss(sec):
    sec = int(max(0, sec))
    return "%d:%02d" % (sec // 60, sec % 60)


def frame(p, cols, rows):
    """Кадр: список строк, каждая — список кусков (текст, цвет)."""
    W = 66
    inner = W - 2
    lines = []

    def row(parts=()):
        parts, text = list(parts), ""
        for i, (t, c) in enumerate(parts):  # строка длиннее рамки — обрезать, рамка не должна съезжать
            if len(text) + len(t) > inner:
                parts = parts[:i] + [(t[: inner - len(text)], c)]
                text += parts[-1][0]
                break
            text += t
        pad = inner - len(text)
        lines.append([("║", GREY)] + list(parts) + [(" " * max(0, pad), ""), ("║", GREY)])

    lines.append([("╔" + "═" * inner + "╗", GREY)])
    row()
    row([("   Bridge&Join", BOLD + WHITE)])
    row([("   шлюз, объединяющий системы финансовой организации", DIM)])
    row()
    row([("   Установка сервера", BOLD)])
    row()
    for i, (title, _, _) in enumerate(STAGES):
        if i < p.stage:
            row([("   ✓  ", GREEN), (title, "")])
        elif i == p.stage and p.error:
            row([("   ✗  ", RED), (title, BOLD + RED)])
        elif i == p.stage:
            t = mmss(time.time() - p.stage_started)
            gap = inner - 6 - len(title) - len(t) - 3
            row([("   ▶  ", CYAN), (title, BOLD + CYAN), (" " * max(1, gap), ""), (t, DIM)])
        else:
            row([("   ·  ", GREY), (title, GREY)])
    row()
    if p.error:
        row([("   ✗ Установка не удалась.", BOLD + RED)])
        row([("     Сейчас откроется экран установщика Ubuntu", "")])
        row([("     с подробностями и командной строкой.", "")])
    else:
        pct = p.percent()
        bar_w = 46
        full = int(bar_w * pct / 100)
        row([("   ", ""), ("█" * full, CYAN), ("░" * (bar_w - full), GREY), ("  %3d %%" % pct, BOLD)])
        row([("   " + p.detail, DIM)] if p.detail and p.stage < DONE_STAGE else [])  # подпись под полосой
        row()
        if p.stage >= DONE_STAGE:
            row([("   ✓ Установка завершена — сервер перезагружается.", BOLD + GREEN)])
            row([("     После загрузки на этом экране будет адрес для входа.", "")])
        else:
            row([("   Не выключайте сервер: он перезагрузится сам и покажет", "")])
            row([("   на этом экране адрес для входа в браузере.", "")])
    row()
    lines.append([("╚" + "═" * inner + "╝", GREY)])
    foot = "прошло " + mmss(time.time() - p.started)
    lines.append([(" " * max(0, (W - len(foot)) // 2) + foot, GREY)])
    # по центру экрана
    left = max(0, (cols - W) // 2)
    top = max(0, (rows - len(lines)) // 2)
    return [[("", "")]] * top + [[(" " * left, "")] + ln for ln in lines]


def to_ansi(fr):
    out = ["\033[H\033[2J"]
    for ln in fr:
        out.append("".join((c + t + RESET) if c else t for t, c in ln) + "\n")
    return "".join(out)


def to_html(fr):
    css = {GREEN: "#4e9a06", CYAN: "#06989a", RED: "#cc0000", GREY: "#888", DIM: "#9a9a9a", BOLD: "", WHITE: "#fff"}
    html = []
    for ln in fr:
        h = ""
        for t, c in ln:
            t = t.replace("&", "&amp;").replace("<", "&lt;")
            if not c:
                h += t
                continue
            color = next((css[k] for k in (GREEN, CYAN, RED, GREY, WHITE, DIM) if k in c), "#ddd")
            weight = "700" if BOLD in c else "400"
            h += '<span style="color:%s;font-weight:%s">%s</span>' % (color, weight, t)
        html.append(h)
    return ('<!doctype html><meta charset="utf-8"><body style="margin:0;background:#000">'
            '<pre style="margin:0;padding:8px;font:16px/1 \'DejaVu Sans Mono\',monospace;color:#ddd;background:#000">'
            + "\n".join(html) + "</pre>")


def read_state():
    try:
        with open(STATE_FILE) as f:
            return f.read().strip()
    except OSError:
        return ""


def demo(path, html_out):
    """Прогон по записанным событиям: кадр в HTML (для проверки вида без установки)."""
    p = Progress()
    for line in open(path, encoding="utf-8"):
        line = line.strip()
        if not line:
            continue
        ev = json.loads(line)
        if "STATE" in ev:
            p.set_state(ev["STATE"])
        else:
            p.feed(ev)
        if "ELAPSED" in ev:  # сдвиг времени для наглядности кадра
            p.stage_started -= ev["ELAPSED"]
            p.started -= ev["ELAPSED"] * 3
    fr = frame(p, 100, 37)
    if html_out:
        open(html_out, "w", encoding="utf-8").write(to_html(fr))
    else:
        sys.stdout.write(to_ansi(fr))


LIVE = "BJ_PROGRESS_STATE" not in os.environ  # настоящая установка (не проверка)
TTY = "/dev/tty1"


class Console:
    """Вывод на первую консоль. Экран установщика Ubuntu при остановке «вешает» консоль (vhangup): открытый
    до этого вывод отвечает EIO (так упал первый вариант экрана) — тогда консоль открывается заново."""

    def __init__(self):
        self.f = None

    def _open(self):
        if not LIVE:
            return sys.stdout
        return open(TTY, "w", encoding="utf-8", errors="replace")

    def write(self, s):
        for _ in range(5):
            try:
                if self.f is None:
                    self.f = self._open()
                self.f.write(s)
                self.f.flush()
                return
            except OSError:
                if self.f is not None and self.f is not sys.stdout:
                    try:
                        self.f.close()
                    except OSError:
                        pass
                self.f = None
                time.sleep(0.3)

    def size(self):
        try:
            return os.get_terminal_size(self.f.fileno()) if self.f is not None else os.get_terminal_size(sys.stdout.fileno())
        except OSError:
            return os.terminal_size((100, 37))


OUT = Console()
PSF = "/snap/subiquity/current/subiquity.psf"  # шрифт консоли установщика: кириллица и рамки


def hold_console(state, font_set):
    """Консоль — за этим экраном: экран установщика Ubuntu, если запустился, остановить (кроме ошибки и
    запроса подтверждения — тогда он нужен). Шрифт с кириллицей — как только смонтирован snap установщика."""
    if not LIVE:
        return font_set
    if state not in ("ERROR", "NEEDS_CONFIRMATION"):
        if subprocess.call(["systemctl", "is-active", "--quiet", CLIENT_UNIT]) == 0:
            subprocess.call(["systemctl", "stop", "--no-block", CLIENT_UNIT], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
            sys.stderr.write("bj-install-progress: экран установщика Ubuntu остановлен\n")
    if not font_set and os.path.exists(PSF):
        font_set = subprocess.call(["setfont", "-C", TTY, PSF], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) == 0
    return font_set


def main():
    if len(sys.argv) >= 3 and sys.argv[1] == "--demo":
        demo(sys.argv[2], sys.argv[4] if len(sys.argv) >= 5 and sys.argv[3] == "--html" else "")
        return 0
    try:
        return run()
    except Exception:  # упал — в журнал; обёртка вернёт экран установщика Ubuntu
        import traceback
        traceback.print_exc()
        return 1


def run():
    sys.stderr.write("bj-install-progress: экран установки запущен\n")
    font_set = hold_console(read_state(), False)
    OUT.write("\033[?25l")  # без мигающего курсора
    p = Progress()
    j = subprocess.Popen(["journalctl", "-b", "-o", "json", "-f", "-n", "all", "-q",
                          "SUBIQUITY_EVENT_TYPE=start", "SUBIQUITY_EVENT_TYPE=finish", "SUBIQUITY_EVENT_TYPE=error"],
                         stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, bufsize=0)
    buf = b""
    error_at = 0.0
    while True:
        r, _, _ = select.select([j.stdout], [], [], 1.0)
        if r:
            chunk = os.read(j.stdout.fileno(), 65536)
            if not chunk and j.poll() is not None:
                return 1  # журнал недоступен — обёртка вернёт экран установщика
            buf += chunk
            while b"\n" in buf:
                line, buf = buf.split(b"\n", 1)
                try:
                    p.feed(json.loads(line.decode("utf-8", "replace")))
                except ValueError:
                    pass
        p.set_state(read_state())
        font_set = hold_console(p.state, font_set)
        if p.state == "NEEDS_CONFIRMATION":  # установщик ждёт подтверждения — это делает его экран
            j.terminate()
            OUT.write("\033[?25h\033[H\033[2J")
            return 3
        cols, rows = OUT.size()
        OUT.write(to_ansi(frame(p, cols, rows)))
        if p.error:
            if not error_at:
                error_at = time.time()
            elif time.time() - error_at > 6:
                j.terminate()
                OUT.write("\033[?25h\033[H\033[2J")
                return 2  # обёртка запустит экран установщика Ubuntu: подробности и командная строка


if __name__ == "__main__":
    sys.exit(main())