bmc_bash.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import json
import time
import shlex
import argparse
import logging

import paramiko  # keep for your environment usage

import lib.ssh_util as ssh_util
from lib.init_logger import init_logger
import lib.bmc_util as bu

# ---- Paths ----

fw_map = {
    "bic": "/home/billy/Desktop/meta/sitv3/build/zephyr/SB_SI.pldm"
}

CONSOLE_SERVICES = [
    "obmc-console@ttyUSBdevice1.service",
    "obmc-console@ttyUSBdevice15.service",
    "obmc-console@ttyUSBdevice29.service",
    "obmc-console@ttyUSBdevice43.service",
]

SESSION_DEFAULT = "/tmp/mmc_log_session.json"

logger = init_logger("bmc_reboot.log", verbose=True)

# Prevent duplicate prints (root handler + custom handler)
logger.propagate = False
root = logging.getLogger()
if root.handlers:
    root.handlers.clear()


# -------------------------
# Helpers
# -------------------------
def _quote(s: str) -> str:
    return shlex.quote(str(s))


def _write_session(path: str, data: dict):
    with open(path, "w", encoding="utf-8") as f:
        json.dump(data, f)


def _read_session(path: str) -> dict:
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def _systemctl_stop_services(ssh, services):
    for svc in services:
        cmd = f"systemctl stop {svc}"
        logger.info(f"[BMC] {cmd}")
        stdin, stdout, stderr = ssh.exec_command(cmd)
        rc = stdout.channel.recv_exit_status()
        if rc != 0:
            err = (stderr.read() or b"").decode(errors="ignore").strip()
            logger.warning(f"stop failed: {svc} (rc={rc}) {err}")


def _systemctl_start_services(ssh, services):
    for svc in services:
        cmd = f"systemctl start {svc}"
        logger.info(f"[BMC] {cmd}")
        stdin, stdout, stderr = ssh.exec_command(cmd)
        rc = stdout.channel.recv_exit_status()
        if rc != 0:
            err = (stderr.read() or b"").decode(errors="ignore").strip()
            logger.warning(f"start failed: {svc} (rc={rc}) {err}")


def _parse_ports(ports_str: str):
    ports = []
    for x in ports_str.split(","):
        x = x.strip()
        if not x:
            continue
        ports.append(int(x))
    return tuple(ports)


def _get_selected_ttys(ports):
    mmc_ports = ["tty_SITV_0", "tty_SITV_1", "tty_SITV_2", "tty_SITV_3"]
    for p in ports:
        if p < 0 or p >= len(mmc_ports):
            raise SystemExit(f"Invalid port: {p}")
    return [mmc_ports[p] for p in ports]


# -------------------------
# Core Actions
# -------------------------
def fw_update(ssh, ip, component="bic", remote_dir="/tmp"):
    """
    Upload PLDM image to remote_dir then run:
      cd remote_dir && pldm-update ag <filename>
    Workaround: absolute path causes "software ID is invalid" in your environment.
    """
    target = {"ip": ip, "port": 22, "username": "root", "password": "0penBmc"}

    if component not in fw_map:
        logger.error(f"Unknown component: {component}")
        return False

    local_path = fw_map[component]
    file_name = os.path.basename(local_path)
    remote_path = f"{remote_dir}/{file_name}"

    if not ssh_util.push_file(target, local_path, remote_dir=remote_dir):
        logger.error(f"push file failed: {local_path} -> {remote_path}")
        return False

    ssh_util.run_command(ssh, f"ls -l {remote_path}", logger)

    # IMPORTANT: cd then run with relative filename
    cmd = f"sh -lc 'cd {_quote(remote_dir)} && pldm-update ag {_quote(file_name)}'"
    ok = ssh_util.run_command(ssh, cmd, logger)
    if not ok:
        logger.error(f"fw update fail (cmd={cmd})")
        return False

    logger.info("fw update success")
    return True


# -------------------------
# CLI Commands (bash-callable)
# -------------------------
def cmd_mmc_start(args):
    """
    Stop obmc-console services, clear logs once, write session file.
    NOTE: This does NOT keep serial open, so it cannot collect logs during fw-update.
    Use `run-update` if you need logs during fw update.
    """
    ssh = ssh_util.connect_bmc(args.ip, logger)

    ports = _parse_ports(args.ports)
    selected_ttys = _get_selected_ttys(ports)

    _systemctl_stop_services(ssh, CONSOLE_SERVICES)
    time.sleep(0.3)

    if args.clear_on_start:
        mmcs = [bu.SerialSSHClient(ssh, tty, logger) for tty in selected_ttys]
        for mmc in mmcs:
            try:
                mmc.clear_log()
            except Exception as e:
                logger.warning(f"clear_log failed on {getattr(mmc, 'port', 'unknown')}: {e}")
        for mmc in mmcs:
            mmc.close()

    sess = {
        "ip": args.ip,
        "ports": list(ports),
        "ttys": selected_ttys,
        "services": CONSOLE_SERVICES,
    }
    _write_session(args.session, sess)
    logger.info(f"MMC session created: {args.session} (ports={ports})")


def cmd_mmc_dump(args):
    """
    Dump MMC logs once (does NOT start/stop services).
    """
    sess = _read_session(args.session)
    ssh = ssh_util.connect_bmc(sess["ip"], logger)

    mmcs = [bu.SerialSSHClient(ssh, tty, logger) for tty in sess["ttys"]]
    logger.info("========== MMC LOG DUMP ==========")
    for mmc in mmcs:
        try:
            logger.info(f"-------sitv {mmc.port[-1]}--------")
            mmc.log_inf()
            if args.clear_after_dump:
                mmc.clear_log()
        except Exception as e:
            logger.warning(f"log dump failed on {getattr(mmc, 'port', 'unknown')}: {e}")
    for mmc in mmcs:
        mmc.close()


def cmd_mmc_exit(args):
    """
    Optionally dump logs once, start services back, and delete session file.
    """
    sess = _read_session(args.session)
    ssh = ssh_util.connect_bmc(sess["ip"], logger)

    if not args.no_dump:
        mmcs = [bu.SerialSSHClient(ssh, tty, logger) for tty in sess["ttys"]]
        logger.info("========== MMC LOG DUMP (ON EXIT) ==========")
        for mmc in mmcs:
            try:
                logger.info(f"-------sitv {mmc.port[-1]}--------")
                mmc.log_inf()
            except Exception as e:
                logger.warning(f"log dump failed on {getattr(mmc, 'port', 'unknown')}: {e}")
        for mmc in mmcs:
            mmc.close()

    _systemctl_start_services(ssh, sess.get("services", CONSOLE_SERVICES))

    try:
        os.remove(args.session)
    except OSError:
        pass

    logger.info("MMC session exited and cleaned up")


def cmd_fw_update(args):
    ssh = ssh_util.connect_bmc(args.ip, logger)
    ok = fw_update(ssh, args.ip, component=args.component, remote_dir=args.remote_dir)
    if not ok:
        raise SystemExit(1)


def cmd_run_update(args):
    """
    One-shot pipeline (RECOMMENDED):
      stop console services
      open serial and KEEP it open during fw update
      clear logs once at start (optional)
      fw update
      dump logs once at end (optional)
      close serial
      restore console services

    This is the only way to ensure mmc logs are collected during fw update
    without relying on BMC-side background log file mechanisms.
    """
    ssh = ssh_util.connect_bmc(args.ip, logger)

    ports = _parse_ports(args.ports)
    selected_ttys = _get_selected_ttys(ports)

    _systemctl_stop_services(ssh, CONSOLE_SERVICES)
    time.sleep(0.3)

    mmcs = []
    try:
        # Keep serial open for the whole operation
        mmcs = [bu.SerialSSHClient(ssh, tty, logger) for tty in selected_ttys]

        if args.clear_on_start:
            for mmc in mmcs:
                mmc.clear_log()

        ok = fw_update(ssh, args.ip, component=args.component, remote_dir=args.remote_dir)
        if not ok:
            raise SystemExit(1)

    finally:
        if mmcs and not args.no_dump:
            logger.info("========== MMC LOG DUMP (ON EXIT) ==========")
            for mmc in mmcs:
                try:
                    logger.info(f"-------sitv {mmc.port[-1]}--------")
                    mmc.log_inf()
                except Exception as e:
                    logger.warning(f"log dump failed on {getattr(mmc, 'port', 'unknown')}: {e}")

        for mmc in mmcs:
            try:
                mmc.close()
            except Exception:
                pass

        _systemctl_start_services(ssh, CONSOLE_SERVICES)
        logger.info("run-update done")


def main():
    parser = argparse.ArgumentParser(description="BMC helper CLI (bash-callable subcommands)")
    sub = parser.add_subparsers(dest="cmd", required=True)

    # mmc-start
    p = sub.add_parser("mmc-start", help="stop console services + clear mmc logs once + create session")
    p.add_argument("-ip", required=True, help="BMC IP")
    p.add_argument("--ports", default="0,1,2,3", help="MMC ports, e.g. 0,1,2,3")
    p.add_argument("--session", default=SESSION_DEFAULT, help="session file path")
    p.add_argument("--clear-on-start", dest="clear_on_start", action="store_true", default=True)
    p.add_argument("--no-clear-on-start", dest="clear_on_start", action="store_false")
    p.set_defaults(func=cmd_mmc_start)

    # mmc-dump
    p = sub.add_parser("mmc-dump", help="dump mmc logs once using existing session")
    p.add_argument("--session", default=SESSION_DEFAULT, help="session file path")
    p.add_argument("--clear-after-dump", action="store_true", help="clear logs after dump")
    p.set_defaults(func=cmd_mmc_dump)

    # mmc-exit
    p = sub.add_parser("mmc-exit", help="dump once (optional) + restore console services + remove session")
    p.add_argument("--session", default=SESSION_DEFAULT, help="session file path")
    p.add_argument("--no-dump", action="store_true", help="do not dump logs on exit")
    p.set_defaults(func=cmd_mmc_exit)

    # fw-update
    p = sub.add_parser("fw-update", help="upload PLDM image then run pldm-update ag in remote dir")
    p.add_argument("-ip", required=True, help="BMC IP")
    p.add_argument("--component", default="bic", choices=list(fw_map.keys()))
    p.add_argument("--remote-dir", default="/tmp", help="remote directory to store image")
    p.set_defaults(func=cmd_fw_update)

    # run-update (NEW)
    p = sub.add_parser("run-update", help="one-shot: keep serial open during fw update + dump once at end")
    p.add_argument("-ip", required=True, help="BMC IP")
    p.add_argument("--component", default="bic", choices=list(fw_map.keys()))
    p.add_argument("--remote-dir", default="/tmp", help="remote directory to store image")
    p.add_argument("--ports", default="0,1,2,3", help="MMC ports, e.g. 0,1,2,3")
    p.add_argument("--clear-on-start", dest="clear_on_start", action="store_true", default=True)
    p.add_argument("--no-clear-on-start", dest="clear_on_start", action="store_false")
    p.add_argument("--no-dump", action="store_true", help="do not dump logs on exit")
    p.set_defaults(func=cmd_run_update)

    args = parser.parse_args()
    args.func(args)


if __name__ == "__main__":
    main()