phonerail/rack/backend/prrack/main.py

"""入口。

生产（服务器上，systemd 起）：
    python -m prrack.main --db /root/phonerail-rack/data/rack.db \\
        --mqtt-env /root/phonerail/mqtt.env --static /root/phonerail-rack/frontend/dist \\
        --fw-dir /root/phonerail/fw
本机开发（不连 broker，起 N 块模拟板子，每秒一条遥测）：
    python -m prrack.main --sim 2 --db dev.db --static ../frontend/dist

⛔ 只许绑回环地址：api.py 按 L2 信任门户注入的 X-Auth-User，这个信任只在「外面够不着本进程」时成立。
"""
import argparse
import asyncio
import ipaddress
import logging
import os
import sys

import uvicorn

from .api import create_app
from .bus import MemBroker, MemBus, PahoBus
from .db import Db, now_ms
from .firmware import FwStore
from .fleet import FleetPoller
from .hub import Hub
from .recovery import RecoveryManager
from .simboard import HUNG, SimBoard

log = logging.getLogger("prrack")

FW_GC_S = 3600  # 固件目录多久清一次（firmware.FwStore.prune）


def read_env(path: str) -> dict:
    out = {}
    with open(path, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line and not line.startswith("#") and "=" in line:
                k, v = line.split("=", 1)
                out[k.strip()] = v.strip().strip('"')
    return out


def build(args):
    db = Db(args.db)
    sims = []
    if args.sim:
        broker = MemBroker()
        bus = MemBus(broker)
        for i in range(args.sim):
            present = (True, True, True, False, True) if i == 0 else (True,) * 5
            sb = SimBoard(MemBus(broker), "5157ab%06x" % i, present=present, seed=i + 1)
            if i == 0:
                sb.set_phone(2, HUNG)  # 留一台「卡死」的，给自动恢复看
            sims.append(sb)
    else:
        env = read_env(args.mqtt_env)
        bus = PahoBus(args.mqtt_host, args.mqtt_port, env.get("BACKEND_USER", "phonerail-backend"),
                      env["BACKEND_PASS"], "phonerail-rack-%d" % os.getpid())
    hub = Hub(db, bus)
    rec = RecoveryManager(hub, now_ms)

    async def up():
        await hub.start()
        for sb in sims:
            sb.start()
        if sims:
            async def ticker():
                while True:
                    for sb in sims:
                        sb.tick(now_ms())
                    await asyncio.sleep(1)
            asyncio.ensure_future(ticker())

    async def down():
        await hub.stop()
        db.close()

    # 手机自己报的那一半（在线 / 温度）从 AppTask 后端读（fleet.py）。没配令牌文件就不读，界面上那一行不出现。
    token = ""
    if args.apptask_env and os.path.isfile(args.apptask_env):
        token = read_env(args.apptask_env).get("APPTASK_TOKEN", "")
    if token:
        async def fleet_up():
            asyncio.ensure_future(FleetPoller(hub, args.apptask_url, token).run())
    else:
        fleet_up = None
        log.info("没有 AppTask 令牌（--apptask-env）：不读手机温度 / 在线状态")

    static = args.static if args.static and os.path.isdir(args.static) else None
    fw = FwStore(db, args.fw_dir, args.fw_base) if args.fw_dir else None
    startup = [up] + ([fleet_up] if fleet_up else [])
    if fw is not None:
        async def fw_gc():
            async def loop():
                while True:
                    try:
                        gone = fw.prune()
                        if gone:
                            log.info("固件目录清理：删了 %d 个（%s）", len(gone), " ".join(gone))
                    except Exception:  # 清理出错不能拖垮后端
                        log.exception("固件目录清理失败")
                    await asyncio.sleep(FW_GC_S)
            asyncio.ensure_future(loop())
        startup.append(fw_gc)
    app = create_app(hub, rec, static, startup=startup, shutdown=[down], fw=fw)
    return app


def ws_library_ok() -> bool:
    for mod in ("websockets", "wsproto"):
        try:
            __import__(mod)
            return True
        except ImportError:
            pass
    return False


def main(argv=None):
    ap = argparse.ArgumentParser(description="PhoneRail 机架管理后端")
    ap.add_argument("--db", default="rack.db")
    ap.add_argument("--static", default=None)
    ap.add_argument("--bind", default="127.0.0.1")
    ap.add_argument("--port", type=int, default=7764)
    ap.add_argument("--mqtt-host", default="127.0.0.1")
    ap.add_argument("--mqtt-port", type=int, default=1884)
    ap.add_argument("--mqtt-env", default="/root/phonerail/mqtt.env")
    ap.add_argument("--fw-dir", default=None,
                    help="固件下载目录（nginx 7765 的 /fw/，生产是 /root/phonerail/fw）；不给就没有 OTA")
    ap.add_argument("--fw-base", default="http://home.example:7765/fw",
                    help="板子下载固件用的 URL 前缀（板子只能走公网）")
    ap.add_argument("--apptask-url", default="http://127.0.0.1:7736", help="AppTask 后端（读机架手机的在线状态 / 温度）")
    ap.add_argument("--apptask-env", default=None, help="含 APPTASK_TOKEN= 的 env 文件（生产是 /root/apptask/apptask-py.env）")
    ap.add_argument("--sim", type=int, default=0, help="不连 broker，起 N 块模拟板子")
    args = ap.parse_args(argv)
    if not ipaddress.ip_address(args.bind).is_loopback:
        sys.exit("拒绝绑 %s：只许回环地址（见文件头）" % args.bind)
    if not ws_library_ok():
        sys.exit("没装 websockets / wsproto：uvicorn 会把 /api/ws 当普通 GET，界面能开但永远不刷新")
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s")
    uvicorn.run(build(args), host=args.bind, port=args.port, log_level="info")


if __name__ == "__main__":
    main()