netid_monitor.py

#!/usr/bin/env python3
"""
NetID Monitor — USB Display Network Identity Daemon
Listens for WeAct Studio USB display dongle and pushes hostname + IP.
Cross-platform: Debian/Ubuntu Linux & Windows.

Protocol based on WeActStudio.SystemMonitor / turing-smart-screen-python.
"""

import sys
import time
import socket
import struct
import logging
import threading
import argparse
from typing import Optional
from enum import IntEnum

try:
    import serial
    import serial.tools.list_ports
except ImportError:
    print("ERROR: pyserial not installed. Run: pip install pyserial pillow")
    sys.exit(1)

try:
    from PIL import Image, ImageDraw, ImageFont
except ImportError:
    print("ERROR: Pillow not installed. Run: pip install pyserial pillow")
    sys.exit(1)

# ─── Logging ──────────────────────────────────────────────────────────────────

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%H:%M:%S",
)
log = logging.getLogger("netid")

# ─── Display Protocol Constants ───────────────────────────────────────────────

class Cmd(IntEnum):
    RESET          = 0x01
    CLEAR          = 0x02
    DISPLAY_BITMAP = 0x2C
    SET_BRIGHTNESS = 0x74
    HELLO          = 0x6D

# Known VID/PID pairs for WeAct / Turing-compatible displays
KNOWN_DEVICES = [
    # WeAct Studio Display (STM32-based CDC)
    (0x0483, 0x5740),   # STM32 Virtual COM Port
    (0x1a86, 0x7523),   # CH340 (some variants)
    (0x0403, 0x6001),   # FTDI (some variants)
]

# Serial strings used to identify the device
DEVICE_STRINGS = [
    "USB35INCHIPS",
    "WeAct",
    "usbserial",
    "ttyACM",       # Linux fallback substring
]

BAUD_RATE   = 115200
DISPLAY_W   = 320   # 3.5" display
DISPLAY_H   = 480
DISPLAY_W_SMALL = 80   # 0.96" display
DISPLAY_H_SMALL = 160

# ─── Helpers ──────────────────────────────────────────────────────────────────

def get_hostname() -> str:
    return socket.gethostname()

def get_ip() -> str:
    """Best-effort: return the primary outbound IP address."""
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        s.connect(("8.8.8.8", 80))
        ip = s.getsockname()[0]
        s.close()
        return ip
    except Exception:
        pass
    try:
        return socket.gethostbyname(socket.gethostname())
    except Exception:
        return "?.?.?.?"

def find_display_port() -> Optional[str]:
    """Scan serial ports and return first likely WeAct display port."""
    ports = list(serial.tools.list_ports.comports())
    for p in ports:
        vid = p.vid
        pid = p.pid
        desc = (p.description or "").lower()
        mfg  = (p.manufacturer or "").lower()
        hwid = (p.hwid or "").upper()

        # Match by VID/PID
        if vid and pid and (vid, pid) in KNOWN_DEVICES:
            log.info(f"Found device by VID/PID: {p.device} ({p.description})")
            return p.device

        # Match by description keywords
        for kw in DEVICE_STRINGS:
            if kw.lower() in desc or kw.lower() in mfg or kw in hwid:
                log.info(f"Found device by keyword '{kw}': {p.device} ({p.description})")
                return p.device

    return None

# ─── Display Protocol ─────────────────────────────────────────────────────────

def _send_reg(ser: serial.Serial, cmd: int, x: int, y: int, ex: int, ey: int):
    """Send a 6-byte command register packet."""
    payload = bytearray([cmd, x >> 2, (x & 3) << 6 | y >> 4,
                         (y & 0xF) << 4 | ex >> 6,
                         (ex & 0x3F) << 2 | ey >> 8, ey & 0xFF])
    ser.write(payload)

def reset_display(ser: serial.Serial):
    """Send HELLO/reset sequence."""
    hello = bytearray([Cmd.HELLO] * 6)
    ser.write(hello)
    time.sleep(0.2)
    ser.reset_input_buffer()

def set_brightness(ser: serial.Serial, level: int = 100):
    """Set backlight brightness 0–255."""
    level = max(0, min(255, level))
    _send_reg(ser, Cmd.SET_BRIGHTNESS, level, 0, 0, 0)

def send_image(ser: serial.Serial, img: Image.Image, x: int = 0, y: int = 0):
    """Send a PIL image as RGB565 bitmap to the display."""
    w, h = img.size
    _send_reg(ser, Cmd.DISPLAY_BITMAP, x, y, x + w - 1, y + h - 1)

    pix = img.convert("RGB").load()
    chunk = bytearray()
    CHUNK_SIZE = DISPLAY_W * 8  # send in chunks for stability

    for row in range(h):
        for col in range(w):
            r, g, b = pix[col, row]
            rgb565 = ((r >> 3) << 11) | ((g >> 2) << 5) | (b >> 3)
            chunk += struct.pack(">H", rgb565)
            if len(chunk) >= CHUNK_SIZE:
                ser.write(chunk)
                chunk = bytearray()
    if chunk:
        ser.write(chunk)
    time.sleep(0.01)

# ─── Screen Rendering ─────────────────────────────────────────────────────────

PALETTE = {
    "bg":        (10,  12,  20),
    "accent":    (0,   200, 255),
    "accent2":   (0,   120, 200),
    "hostname":  (255, 255, 255),
    "label":     (100, 140, 180),
    "ip":        (0,   220, 180),
    "border":    (0,   60,  100),
    "dim":       (40,  50,  70),
}

def _try_font(size: int, bold: bool = False):
    """Try to load a system monospace font, fall back to default."""
    candidates = []
    if sys.platform.startswith("win"):
        candidates = ["consola.ttf", "cour.ttf", "lucon.ttf"]
    else:
        candidates = [
            "/usr/share/fonts/truetype/dejavu/DejaVuSansMono-Bold.ttf",
            "/usr/share/fonts/truetype/dejavu/DejaVuSansMono.ttf",
            "/usr/share/fonts/truetype/liberation/LiberationMono-Bold.ttf",
            "/usr/share/fonts/truetype/liberation/LiberationMono-Regular.ttf",
            "/usr/share/fonts/truetype/ubuntu/UbuntuMono-B.ttf",
            "/usr/share/fonts/truetype/ubuntu/UbuntuMono-R.ttf",
        ]
    for path in candidates:
        try:
            return ImageFont.truetype(path, size)
        except (IOError, OSError):
            continue
    return ImageFont.load_default()

def render_screen(hostname: str, ip: str,
                  width: int = DISPLAY_W, height: int = DISPLAY_H) -> Image.Image:
    """Render the identification screen as a PIL Image."""
    img = Image.new("RGB", (width, height), PALETTE["bg"])
    draw = ImageDraw.Draw(img)

    # ── Background grid decoration ─────────────────────────────────────────
    grid_step = 32
    for gx in range(0, width, grid_step):
        draw.line([(gx, 0), (gx, height)], fill=(20, 28, 45), width=1)
    for gy in range(0, height, grid_step):
        draw.line([(0, gy), (width, gy)], fill=(20, 28, 45), width=1)

    # ── Top accent bar ─────────────────────────────────────────────────────
    draw.rectangle([(0, 0), (width, 4)], fill=PALETTE["accent"])
    draw.rectangle([(0, 5), (width, 6)], fill=PALETTE["accent2"])

    # ── Bottom accent bar ──────────────────────────────────────────────────
    draw.rectangle([(0, height - 4), (width, height)], fill=PALETTE["accent"])
    draw.rectangle([(0, height - 7), (width, height - 5)], fill=PALETTE["accent2"])

    # ── Central panel ─────────────────────────────────────────────────────
    pad = 12
    panel_y1 = height // 2 - 110
    panel_y2 = height // 2 + 110
    draw.rounded_rectangle(
        [(pad, panel_y1), (width - pad, panel_y2)],
        radius=10,
        fill=(15, 20, 35),
        outline=PALETTE["border"],
        width=2,
    )

    cx = width // 2

    # ── "NET ID" header ────────────────────────────────────────────────────
    font_hdr = _try_font(11, bold=True)
    label_y = panel_y1 + 14
    draw.text((cx, label_y), "◆  NET ID  ◆",
              font=font_hdr, fill=PALETTE["accent"], anchor="mm")

    # Divider
    div_y = label_y + 16
    draw.line([(pad + 16, div_y), (width - pad - 16, div_y)],
              fill=PALETTE["border"], width=1)

    # ── Hostname section ───────────────────────────────────────────────────
    font_label = _try_font(10)

    hn_label_y = div_y + 14
    draw.text((cx, hn_label_y), "HOSTNAME",
              font=font_label, fill=PALETTE["label"], anchor="mm")

    hn_y = hn_label_y + 22
    # Shrink font if hostname is long
    hn_font = _try_font(18, bold=True)
    hn_display = hostname
    if len(hostname) > 14:
        hn_font = _try_font(13, bold=True)

    draw.text((cx, hn_y), hn_display,
              font=hn_font, fill=PALETTE["hostname"], anchor="mm")

    # ── Divider ────────────────────────────────────────────────────────────
    mid_div_y = hn_y + 26
    draw.line([(pad + 24, mid_div_y), (width - pad - 24, mid_div_y)],
              fill=PALETTE["dim"], width=1)

    # ── IP section ─────────────────────────────────────────────────────────
    ip_label_y = mid_div_y + 12
    draw.text((cx, ip_label_y), "IP ADDRESS",
              font=font_label, fill=PALETTE["label"], anchor="mm")

    ip_y = ip_label_y + 22
    font_ip = _try_font(20, bold=True)
    draw.text((cx, ip_y), ip,
              font=font_ip, fill=PALETTE["ip"], anchor="mm")

    # ── Corner decorations ─────────────────────────────────────────────────
    corner = 8
    corners = [
        (pad + 2, panel_y1 + 2),
        (width - pad - 2, panel_y1 + 2),
        (pad + 2, panel_y2 - 2),
        (width - pad - 2, panel_y2 - 2),
    ]
    for cx2, cy2 in corners:
        draw.ellipse([(cx2 - corner, cy2 - corner),
                      (cx2 + corner, cy2 + corner)],
                     outline=PALETTE["accent"], width=2)

    # ── Footer ─────────────────────────────────────────────────────────────
    font_footer = _try_font(9)
    draw.text((cx, height - 14), "⬡  NETWORK IDENTITY BEACON  ⬡",
              font=font_footer, fill=PALETTE["dim"], anchor="mm")

    return img


def render_small_screen(hostname: str, ip: str) -> Image.Image:
    """Render for the 0.96" 80x160 display."""
    img = Image.new("RGB", (DISPLAY_W_SMALL, DISPLAY_H_SMALL), PALETTE["bg"])
    draw = ImageDraw.Draw(img)

    cx = DISPLAY_W_SMALL // 2
    font = _try_font(8)

    draw.rectangle([(0, 0), (DISPLAY_W_SMALL, 2)], fill=PALETTE["accent"])
    draw.rectangle([(0, DISPLAY_H_SMALL - 2), (DISPLAY_W_SMALL, DISPLAY_H_SMALL)],
                   fill=PALETTE["accent"])

    draw.text((cx, 12), "NET ID", font=font, fill=PALETTE["accent"], anchor="mm")
    draw.line([(4, 20), (76, 20)], fill=PALETTE["border"], width=1)

    draw.text((cx, 32), "HOST:", font=font, fill=PALETTE["label"], anchor="mm")

    # wrap hostname if long
    hn_font = _try_font(8)
    hn = hostname if len(hostname) <= 10 else hostname[:10] + "…"
    draw.text((cx, 44), hn, font=hn_font, fill=PALETTE["hostname"], anchor="mm")

    draw.line([(4, 56), (76, 56)], fill=PALETTE["dim"], width=1)
    draw.text((cx, 68), "IP:", font=font, fill=PALETTE["label"], anchor="mm")

    # split IP into two lines for small screen
    ip_parts = ip.split(".")
    if len(ip_parts) == 4:
        line1 = ".".join(ip_parts[:2])
        line2 = ".".join(ip_parts[2:])
    else:
        line1 = ip[:10]
        line2 = ip[10:]

    ip_font = _try_font(8)
    draw.text((cx, 80), line1, font=ip_font, fill=PALETTE["ip"], anchor="mm")
    draw.text((cx, 92), line2, font=ip_font, fill=PALETTE["ip"], anchor="mm")

    return img


# ─── Main Daemon ──────────────────────────────────────────────────────────────

class NetIDDaemon:
    def __init__(self, port: Optional[str] = None,
                 small: bool = False,
                 poll_interval: float = 10.0,
                 brightness: int = 80,
                 layout_file: Optional[str] = None):
        self.port          = port
        self.small         = small
        self.poll_interval = poll_interval
        self.brightness    = brightness
        self.layout_file   = layout_file
        self._ser: Optional[serial.Serial] = None
        self._stop         = threading.Event()

    def _open_port(self, port: str) -> bool:
        try:
            self._ser = serial.Serial(port, BAUD_RATE, timeout=2, rtscts=True)
            log.info(f"Opened serial port: {port}")
            return True
        except serial.SerialException as e:
            log.debug(f"Cannot open {port}: {e}")
            return False

    def _close_port(self):
        if self._ser and self._ser.is_open:
            try:
                self._ser.close()
            except Exception:
                pass
        self._ser = None

    def _push(self):
        hostname = get_hostname()
        ip       = get_ip()
        log.info(f"Pushing: {hostname} / {ip}")

        if self.layout_file:
            from pathlib import Path
            from layout_renderer import build_context, load_layout, render_layout
            layout = load_layout(self.layout_file)
            ctx = build_context(hostname=hostname, ip=ip)
            img = render_layout(layout, ctx, base_dir=Path(self.layout_file).parent)
            log.info(f"Rendered using layout: {self.layout_file}")
        elif self.small:
            img = render_small_screen(hostname, ip)
        else:
            img = render_screen(hostname, ip)

        reset_display(self._ser)
        time.sleep(0.1)
        set_brightness(self._ser, self.brightness)
        send_image(self._ser, img)
        log.info("Screen updated successfully.")

    def run(self):
        log.info("NetID Monitor started. Watching for USB display…")
        while not self._stop.is_set():
            # Determine port
            port = self.port or find_display_port()

            if port is None:
                log.debug("No display found, retrying in 3s…")
                self._stop.wait(3)
                continue

            if self._ser is None or not self._ser.is_open:
                if not self._open_port(port):
                    self._stop.wait(3)
                    continue

            try:
                self._push()
                self._stop.wait(self.poll_interval)
            except serial.SerialException as e:
                log.warning(f"Serial error: {e} — will reconnect.")
                self._close_port()
                self._stop.wait(2)
            except Exception as e:
                log.error(f"Unexpected error: {e}")
                self._close_port()
                self._stop.wait(2)

        self._close_port()
        log.info("NetID Monitor stopped.")

    def stop(self):
        self._stop.set()


# ─── Entry Point ──────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(
        description="NetID Monitor — push hostname+IP to WeAct USB display",
    )
    parser.add_argument("--port", "-p", default=None,
                        help="Serial port (e.g. COM3, /dev/ttyACM0). Auto-detect if omitted.")
    parser.add_argument("--small", action="store_true",
                        help="Use layout for 0.96\" 80×160 display.")
    parser.add_argument("--interval", "-i", type=float, default=10.0,
                        help="Refresh interval in seconds (default: 10).")
    parser.add_argument("--brightness", "-b", type=int, default=80,
                        help="Backlight brightness 0–255 (default: 80).")
    parser.add_argument("--list-ports", action="store_true",
                        help="List available serial ports and exit.")
    parser.add_argument("--once", action="store_true",
                        help="Push once and exit (no daemon loop).")
    parser.add_argument("--layout", default=None,
                        help="Path to layout JSON file for dynamic rendering (e.g. layouts/screen.json).")
    parser.add_argument("--verbose", "-v", action="store_true",
                        help="Enable debug logging.")
    args = parser.parse_args()

    if args.verbose:
        logging.getLogger().setLevel(logging.DEBUG)

    if args.list_ports:
        print("Available serial ports:")
        for p in serial.tools.list_ports.comports():
            vid = f"{p.vid:04X}" if p.vid else "????"
            pid = f"{p.pid:04X}" if p.pid else "????"
            print(f"  {p.device:20s}  VID={vid} PID={pid}  {p.description}")
        return

    daemon = NetIDDaemon(
        port          = args.port,
        small         = args.small,
        poll_interval = args.interval,
        brightness    = args.brightness,
        layout_file   = args.layout,
    )

    if args.once:
        port = args.port or find_display_port()
        if port is None:
            log.error("No display found. Use --list-ports or --port.")
            sys.exit(1)
        daemon._open_port(port)
        daemon._push()
        daemon._close_port()
        return

    try:
        daemon.run()
    except KeyboardInterrupt:
        print()
        daemon.stop()


if __name__ == "__main__":
    main()