docker-migrate/core/main.py

# -*- coding: utf-8 -*-
"""
main.py — Точка входа. Меню выбора режима.
"""

import sys
import os
import argparse
import signal
import atexit

_PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if _PROJECT_ROOT not in sys.path:
    sys.path.insert(0, _PROJECT_ROOT)

from core.color import banner, menu, prompt, confirm, info, error as cerror, warn
from core import state


def _save_on_exit():
    """Сохраняет состояние при любом завершении (Ctrl+C, SIGTERM, os._exit)."""
    state.flush_state()


def _handle_sigterm(signum, frame):
    """Обработчик SIGTERM (kill, systemd stop)."""
    state.set_interrupted()
    state.flush_state()
    sys.exit(128 + signum)


def _handle_keyboard_interrupt():
    """Обработчик Ctrl+C — сохраняет состояние до выхода."""
    state.set_interrupted()
    state.flush_state()
    print()
    warn("Прервано (Ctrl+C). Состояние сохранено. Запустите docker-migrate --resume для продолжения.")
    sys.exit(130)

# Регистрируем сохранение state при любом завершении
atexit.register(_save_on_exit)
signal.signal(signal.SIGTERM, _handle_sigterm)


def main():
    # Инициализируем state сразу, чтобы _LAST_STATE был готов при любом выходe
    state.load_state()
    
    parser = argparse.ArgumentParser(description="Docker Service Migration Tool")
    parser.add_argument("--mode", choices=["source", "target"], help="Режим работы")
    parser.add_argument("--resume", action="store_true", help="Продолжить после ошибки")
    parser.add_argument("--status", action="store_true", help="Показать статус")
    parser.add_argument("--logs", action="store_true", help="Показать логи ошибки")
    parser.add_argument("--remote-dir", help="Директория архива на target")
    parser.add_argument("--no-color", action="store_true", help="Отключить цвета")
    args = parser.parse_args()

    if args.no_color:
        os.environ["NO_COLOR"] = "1"

    if args.status:
        state.show_status()
        sys.exit(0)

    if args.logs:
        state.show_logs()
        sys.exit(0)

    if args.resume:
        _do_resume()
        sys.exit(0)

    if args.mode == "source":
        from source.source import run_source_mode
        try:
            run_source_mode()
        except KeyboardInterrupt:
            _handle_keyboard_interrupt()
        sys.exit(0)
    elif args.mode == "target":
        if args.remote_dir:
            state.set_stage("INIT", mode="target", target_remote_dir=args.remote_dir)
        from target.target import run_target_mode
        try:
            run_target_mode()
        except KeyboardInterrupt:
            _handle_keyboard_interrupt()
        sys.exit(0)

    # Интерактивный режим
    # Проверим, не было ли прерывания в прошлом запуске
    recovered = state.recover_after_interrupt()
    if recovered:
        banner()
        warn(f"⚠ Обнаружено прерывание прошлого запуска!")
        info(f"  Шаг: {recovered['stage']}")
        info(f"  Режим: {recovered['mode']}")
        print(f"  {recovered['hint']}")
        print()
        if confirm("Продолжить с места прерывания", default="y"):
            try:
                _do_resume()
            except KeyboardInterrupt:
                _handle_keyboard_interrupt()
            sys.exit(0)
        else:
            info("Сбрасываем прерванное состояние. Можете начать заново.")
            state.set_stage("INIT", mode=None, interrupted_at=None, resumable_hint=None)

    banner()
    while True:
        st = state.load_state()
        # Resume показываем только если есть незавершённая работа (stage != INIT, != DONE)
        stage = st.get("stage", "INIT")
        has_resume = bool(st.get("mode") and stage not in ("INIT", "DONE"))
        menu(has_resume=has_resume)
        try:
            choice = prompt("Ваш выбор", default="0").strip()
        except EOFError:
            break
        except KeyboardInterrupt:
            _handle_keyboard_interrupt()
        if choice == "1":
            from source.source import run_source_mode
            try:
                run_source_mode()
            except KeyboardInterrupt:
                _handle_keyboard_interrupt()
            break
        elif choice == "2":
            from target.target import run_target_mode
            try:
                run_target_mode()
            except KeyboardInterrupt:
                _handle_keyboard_interrupt()
            break
        elif choice == "3":
            if not has_resume:
                warn("Нет сохранённого состояния для продолжения")
                continue
            try:
                _do_resume()
            except KeyboardInterrupt:
                _handle_keyboard_interrupt()
            break
        elif choice == "4":
            _do_status_logs()
        elif choice == "0":
            info("Выход")
            sys.exit(0)
        else:
            warn("Неверный выбор")


def _do_resume():
    from core.fsm import FSM
    st = state.load_state()
    stage = st.get("stage")
    mode = st.get("mode")
    if not mode or stage == "INIT":
        info("Нет сохранённого состояния для продолжения")
        return
    # Сбрасываем флаг прерывания, если был
    state.set_stage(stage, interrupted_at=None, resumable_hint=None)
    fsm = FSM(mode=mode)
    try:
        fsm.resume_from(stage)
    except KeyboardInterrupt:
        _handle_keyboard_interrupt()
    except Exception as e:
        cerror(f"Ошибка при resume: {e}")


def _do_status_logs():
    from core import state as stmod
    while True:
        print("\n")
        print("  1  Показать статус")
        print("  2  Показать лог последней ошибки")
        print("  0  Назад")
        try:
            c = prompt("Выбор", default="0").strip()
        except EOFError:
            break
        except KeyboardInterrupt:
            _handle_keyboard_interrupt()
        if c == "1":
            stmod.show_status()
        elif c == "2":
            stmod.show_logs()
        elif c == "0":
            break


if __name__ == "__main__":
    main()