charging-cabinet/tools/mock_redis.py

#!/usr/bin/env python3
"""
Redis 模拟数据填充工具
- 从数据库读取前10个设备
- 生成随机状态写入Redis（设备级+板级+通道级）
- 300秒后自动过期（与设备上报TTL一致）

用法：
  python tools/mock_redis.py          # 填充模拟数据
  python tools/mock_redis.py --clear  # 清除所有模拟数据
"""
import os
import sys
import json
import random
import redis
import pymysql

# 默认指向 SSH 隧道（tools/dev-tunnel.ps1），不再使用已移除的 ZeroTier/nginx 转发
# Redis 配置
REDIS_HOST = os.environ.get("PMS_REDIS_HOST", "127.0.0.1")
REDIS_PORT = int(os.environ.get("PMS_REDIS_PORT", "16379"))
REDIS_PASSWORD = "Hbhyg731024@"
REDIS_DB = 5

# MySQL 配置
MYSQL_HOST = os.environ.get("PMS_MYSQL_HOST", "127.0.0.1")
MYSQL_PORT = int(os.environ.get("PMS_MYSQL_PORT", "13306"))
MYSQL_USER = "root"
MYSQL_PASSWORD = "Hbhyg731024@"
MYSQL_DB = "pms"

# 设备数量
DEVICE_COUNT = 10

# 通道状态: 0=空闲, 1=插入, 2=充电中, 3=充满, 4=故障
STATUS_IDLE = 0
STATUS_INSERTED = 1
STATUS_CHARGING = 2
STATUS_FULL = 3
STATUS_FAULT = 4


def get_devices():
    """从数据库获取前N个设备"""
    conn = pymysql.connect(
        host=MYSQL_HOST, port=MYSQL_PORT, user=MYSQL_USER,
        password=MYSQL_PASSWORD, database=MYSQL_DB, charset="utf8mb4"
    )
    try:
        with conn.cursor(pymysql.cursors.DictCursor) as cur:
            cur.execute("""
                SELECT c.imei, c.abstract_id, c.name,
                       (SELECT COUNT(*) FROM cabin_boards cb WHERE cb.cabinet_id = c.id) as board_count
                FROM cabinets c ORDER BY c.id LIMIT %s
            """, (DEVICE_COUNT,))
            return cur.fetchall()
    finally:
        conn.close()


def gen_channel_hex(status: int) -> str:
    """生成单通道hex（4字节=8个hex字符）"""
    # 字节1: bit0=在线 bit1=充电中 bit2=充满 bit3=故障
    byte1 = 0x01  # 在线
    if status == STATUS_CHARGING:
        byte1 |= 0x02
    elif status == STATUS_FULL:
        byte1 |= 0x04
    elif status == STATUS_FAULT:
        byte1 |= 0x08
    # 字节2-4: 随机（模拟电压电流等）
    byte2 = random.randint(0x10, 0x60)
    byte3 = random.randint(0x00, 0xFF)
    byte4 = random.randint(0x00, 0xFF)
    return f"{byte1:02x}{byte2:02x}{byte3:02x}{byte4:02x}"


def gen_board_status_hex(board_idx: int) -> str:
    """生成一块板6个通道的status_hex"""
    parts = []
    for ch in range(6):
        # 随机分配状态，充电和故障少一些
        r = random.random()
        if r < 0.02:
            status = STATUS_FAULT
        elif r < 0.25:
            status = STATUS_CHARGING
        elif r < 0.32:
            status = STATUS_FULL
        elif r < 0.42:
            status = STATUS_INSERTED
        else:
            status = STATUS_IDLE
        parts.append(gen_channel_hex(status))
    return "".join(parts)


def fill_mock_data():
    """填充模拟数据"""
    r = redis.Redis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, db=REDIS_DB, decode_responses=True)
    devices = get_devices()
    print(f"获取到 {len(devices)} 个设备")

    for dev in devices:
        imei = dev["imei"]
        board_count = dev.get("board_count") or 6
        device_key = f"device:{imei}"

        # 设备级
        r.hset(device_key, mapping={
            "online": "1",
            "rssi": str(random.randint(30, 95)),
            "pow_fail_dc": random.choice(["0", "0", "0", "1"]),  # 25%概率停电
            "pow_fail_ac": "0",
        })
        r.expire(device_key, 300)

        # 板级 + 通道级
        for board_idx in range(1, board_count + 1):
            board_key = f"{device_key}:board:{board_idx}"
            status_hex = gen_board_status_hex(board_idx)
            r.hset(board_key, mapping={"status_hex": status_hex})

            # 解析6个通道
            for ch in range(6):
                ch_hex = status_hex[ch * 8 : (ch + 1) * 8]
                byte1 = int(ch_hex[0:2], 16)
                ch_key = f"{board_key}:ch:{ch}"
                r.hset(ch_key, mapping={
                    "on": "1" if byte1 & 0x02 else "0",
                    "full": "1" if byte1 & 0x04 else "0",
                    "fault": "1" if byte1 & 0x08 else "0",
                    "hex": ch_hex,
                })

        name = dev.get("name") or imei
        print(f"  ✓ {name} ({imei}) — {board_count}板")

    print(f"\n完成！已填充 {len(devices)} 个设备，300秒后自动过期")


def clear_mock_data():
    """清除所有模拟数据"""
    r = redis.Redis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, db=REDIS_DB, decode_responses=True)
    devices = get_devices()
    count = 0
    for dev in devices:
        imei = dev["imei"]
        keys = r.keys(f"device:{imei}*")
        if keys:
            r.delete(*keys)
            count += len(keys)
    print(f"已清除 {count} 个 Redis key")


if __name__ == "__main__":
    if "--clear" in sys.argv:
        clear_mock_data()
    else:
        fill_mock_data()