332 lines
14 KiB
Python

"""High-level async controller for the Shining Mask.
Example::
import asyncio
from mask import ShiningMask
async def main():
async with ShiningMask() as mask:
await mask.set_brightness(80)
await mask.show_image(3)
await mask.set_text("HELLO", color=(0, 255, 0))
asyncio.run(main())
"""
from __future__ import annotations
import asyncio
import logging
from typing import Optional, Sequence, Tuple
import bitmap, protocol
from constants import TextMode, UploadKind
from exceptions import UploadError
from transport import MaskTransport
log = logging.getLogger("shiningmask.mask")
Color = Tuple[int, int, int]
class ShiningMask:
"""Async controller for one Shining Mask over BLE."""
def __init__(
self,
address: Optional[str] = None,
*,
name_prefix: str = "MASK",
adapter: Optional[str] = None,
write_with_response: Optional[bool] = None,
transport: Optional[MaskTransport] = None,
):
self.transport = transport or MaskTransport(
address,
name_prefix=name_prefix,
adapter=adapter,
write_with_response=write_with_response,
)
self._upload_lock = asyncio.Lock()
# -- lifecycle ------------------------------------------------------------
async def connect(self, timeout: float = 15.0, attempts: int = 3):
await self.transport.connect(timeout=timeout, attempts=attempts)
return self
async def disconnect(self):
await self.transport.disconnect()
async def __aenter__(self):
await self.connect()
return self
async def __aexit__(self, *exc):
await self.disconnect()
@property
def is_connected(self) -> bool:
return self.transport.is_connected
# -- simple control commands ---------------------------------------------
async def set_brightness(self, level: int):
await self.transport.write_command(protocol.cmd_brightness(level))
async def show_image(self, image_id: int):
"""Show a built-in image (IMAG)."""
await self.transport.write_command(protocol.cmd_image(image_id))
async def play_animation(self, anim_id: int):
"""Play a built-in animation (ANIM)."""
await self.transport.write_command(protocol.cmd_animation(anim_id))
async def play_diy(self, image_id: int):
"""Show one uploaded DIY image (PLAY)."""
await self.transport.write_command(protocol.cmd_play_diy(image_id))
async def play_diy_sequence(self, image_ids: Sequence[int]):
await self.transport.write_command(protocol.cmd_play_diy_sequence(list(image_ids)))
async def delete_diy(self, image_ids: Sequence[int]):
await self.transport.write_command(protocol.cmd_delete_diy(list(image_ids)))
async def get_diy_count(self, timeout: float = 5.0) -> Optional[int]:
"""Ask how many DIY images are stored (CHEC). Returns None if no answer."""
await self.transport.write_command(protocol.cmd_check_diy_count())
try:
_, raw = await self._collect_notification("CHEC", timeout=timeout)
except Exception:
return None
import crypto
buf = crypto.decrypt(raw)
# Response frame is [len]['CHEC'][count]; the count sits right after the
# 4-byte command name (NOT after the total length byte buf[0]).
idx = 1 + len("CHEC")
return buf[idx] if len(buf) > idx else None
# -- text styling ---------------------------------------------------------
async def set_text_mode(self, mode: int):
await self.transport.write_command(protocol.cmd_text_mode(mode))
async def set_text_speed(self, speed: int):
await self.transport.write_command(protocol.cmd_text_speed(speed))
async def set_text_color_effect(self, mode: int, enable: bool = True):
await self.transport.write_command(protocol.cmd_text_color_effect(mode, enable))
async def set_foreground_color(self, r: int, g: int, b: int, enable: bool = True):
await self.transport.write_command(protocol.cmd_text_fg_color(r, g, b, enable))
async def set_background_color(self, r: int, g: int, b: int, enable: bool = True):
await self.transport.write_command(protocol.cmd_text_bg_color(r, g, b, enable))
# -- text / bitmap upload -------------------------------------------------
async def set_text(
self,
text: str,
*,
color: Color = (255, 255, 255),
mode: Optional[int] = TextMode.SCROLL_LEFT,
speed: Optional[int] = None,
clear_effects: bool = True,
timeout: float = 10.0,
**render_kwargs,
):
"""Render ``text`` to a bitmap and upload + display it.
``color`` sets the text color. ``mode`` is the scroll style (default
scroll-left; pass ``TextMode.STEADY`` to hold it still). With
``clear_effects`` we disable the ``M`` special effect / animated
background so the mask stops cycling to other animations, and the
scroll ``mode``/``speed`` are pinned *after* the upload so they stick.
"""
bmp, color_array = bitmap.build_text_upload(text, color=color, **render_kwargs)
await self.upload_bitmap(bmp, color_array, timeout=timeout)
await self._apply_display_state(color, mode, speed, clear_effects)
async def _apply_display_state(self, color, mode, speed, clear_effects):
"""Pin a clean, steady display state after a bitmap upload."""
if clear_effects:
# Turn off the M special-effect/background-image carousel and force a
# black background so nothing animates behind the bitmap.
await self.set_text_color_effect(0, enable=False)
await self.set_background_color(0, 0, 0, enable=True)
# Color the foreground the way the app does (FC), in addition to the
# per-column color array already in the upload.
await self.set_foreground_color(*color, enable=True)
if mode is not None:
await self.set_text_mode(mode)
if speed is not None:
await self.set_text_speed(speed)
async def upload_bitmap(self, bitmap_bytes: bytes, color_array: bytes, *, timeout: float = 10.0):
"""Upload a raw mask bitmap + color array using the DATS/REOK/DATCP flow."""
payload = bytes(bitmap_bytes) + bytes(color_array)
await self._upload(
payload,
second_field=len(bitmap_bytes),
kind=UploadKind.BITMAP,
timeout=timeout,
)
async def clear_diy(self, timeout: float = 4.0) -> int:
"""Delete all stored DIY images. Returns how many were removed."""
n = await self.get_diy_count(timeout=timeout) or 0
ids = list(range(1, n + 1))
for i in range(0, len(ids), 10): # DELE takes at most 10 ids
await self.delete_diy(ids[i:i + 10])
await asyncio.sleep(0.25)
return n
async def upload_frame(self, data: bytes, slot: int, **kwargs):
"""Upload one DIY image frame to ``slot`` (see :meth:`upload_raw_image`)."""
await self.upload_raw_image(bytes(data), index=slot, **kwargs)
async def play_frame(self, slot: int):
"""Show a stored DIY image instantly (alias of :meth:`play_diy`)."""
await self.play_diy(slot)
async def upload_raw_image(self, data: bytes, *, index: int = 1,
timestamp: int = 0x07C5F9FF, chunk_delay: float = 0.05,
init_delay: float = 0.15):
"""Upload a raw DIY image/animation to the mask (fire-and-forget).
This is the method that actually works for full-color custom content
(verified against BishopFox/shining-mask): an image ``DATS`` (toggle
``0x01``) is sent, then every data chunk is streamed WITHOUT waiting for
``DATSOK``/``REOK``, then ``DATCP`` with a 4-byte timestamp. The content
displays immediately and is stored as DIY image(s) that :meth:`play_diy`
can replay smoothly (no upload logo).
``data`` must already be in the mask's internal image format (a
multi-frame blob is played as an animation). Generating that format for
arbitrary images is not yet solved -- see the project notes.
"""
if self._upload_lock.locked():
raise UploadError("an upload is already in progress")
async with self._upload_lock:
dats_args = len(data).to_bytes(2, "big") + int(index).to_bytes(2, "big") + b"\x01"
await self.transport.write_command(protocol.encode_command("DATS", dats_args))
await asyncio.sleep(init_delay)
for packet in protocol.iter_upload_packets(data): # unpadded, like the app
await self.transport.write_upload(packet)
await asyncio.sleep(chunk_delay)
await self.transport.write_command(
protocol.encode_command("DATCP", int(timestamp).to_bytes(4, "big"))
)
async def upload_image(self, rgb_bytes: bytes, image_index: int, *, timeout: float = 15.0):
"""Upload a full-color DIY image (raw RGB) to a storage slot.
Experimental: mask-go does not implement image upload, so this follows
ble-protocol.md only. ``rgb_bytes`` is 3 bytes per pixel; the mask is 16
rows tall, so length should be ``16 * width * 3``.
"""
await self._upload(
bytes(rgb_bytes),
second_field=image_index,
kind=UploadKind.IMAGE,
timeout=timeout,
)
async def _upload(self, payload: bytes, *, second_field: int, kind: int, timeout: float,
wait_per_packet: bool = True):
total_len = len(payload)
if total_len == 0:
raise UploadError("nothing to upload (empty payload)")
if self._upload_lock.locked():
raise UploadError("an upload is already in progress")
async with self._upload_lock:
log.info("upload start: %d bytes (kind=%d)", total_len, kind)
await self.transport.write_command(
protocol.cmd_dats(total_len, second_field, kind)
)
await self.transport.wait_for(protocol.RESP_DATS_OK, timeout=timeout)
# Image uploads use fixed 100-byte packets (ble-protocol.md); bitmap
# uploads are unpadded (matching mask-go).
pad_to = 100 if kind == UploadKind.IMAGE else None
for packet in protocol.iter_upload_packets(payload, pad_to=pad_to):
await self.transport.write_upload(packet)
# wait_per_packet=False streams packets without the per-packet
# REOK round-trip (much faster, used for animation). The REOKs
# still arrive and are drained by the DATCPOK wait below.
if wait_per_packet:
await self.transport.wait_for(protocol.RESP_REC_OK, timeout=timeout)
await self.transport.write_command(protocol.cmd_datcp())
await self.transport.wait_for(protocol.RESP_DATCP_OK, timeout=timeout)
log.info("upload complete")
# -- faces -----------------------------------------------------------------
async def setup_face_mode(self, brightness: int = 80, mono_color: Optional[Color] = None):
"""Put the mask in a clean state for drawing faces/images.
Sets a black background, disables the M effect, and selects steady
(non-scrolling) display. If ``mono_color`` is given, the whole face is
colored uniformly via FC (so frames can be uploaded as bitmap-only --
half the bytes, much faster/more reliable); otherwise FC is disabled so
per-column colors drive the pixels.
"""
await self.set_brightness(brightness)
await self.set_text_color_effect(0, enable=False) # M effect off
await self.set_background_color(0, 0, 0, enable=True) # black background
if mono_color is not None:
await self.set_foreground_color(*mono_color, enable=True)
else:
await self.set_foreground_color(255, 255, 255, enable=False)
await self.set_text_mode(TextMode.STEADY)
async def show_face(self, columns, color: Color = (0, 220, 255), *,
per_column_colors=None, mono: bool = False,
fast: bool = True, timeout: float = 8.0):
"""Display a face/bitmap given as columns (each 16 ints, 0/1).
``mono`` uploads just the bitmap (no color array) -- use it after
``setup_face_mode(mono_color=...)`` for the fastest, most reliable frames
(half the packets). Otherwise ``color`` (single) or ``per_column_colors``
(multi) build a color array. With ``fast`` the frame is streamed without
per-packet acks (for animation).
"""
bmp = bitmap.encode_bitmap(columns)
if mono:
payload = bytes(bmp) # color comes from FC
elif per_column_colors is not None:
payload = bytes(bmp) + bitmap.encode_color_array_per_column(per_column_colors)
else:
payload = bytes(bmp) + bitmap.encode_color_array(len(columns), color)
await self._upload(
payload,
second_field=len(bmp),
kind=UploadKind.BITMAP,
timeout=timeout,
wait_per_packet=not fast,
)
# -- audio visualization (experimental) -----------------------------------
async def send_audio_frame(self, mode: int, nibbles: Sequence[int]):
await self.transport.write_audio(protocol.cmd_audio_frame(mode, list(nibbles)))
# -- internals ------------------------------------------------------------
async def _collect_notification(self, *prefixes: str, timeout: float):
"""Wait for a matching notification and also return its raw bytes."""
loop = asyncio.get_running_loop()
deadline = loop.time() + timeout
q = self.transport._notify_queue # noqa: SLF001 - intra-package access
while True:
remaining = deadline - loop.time()
if remaining <= 0:
raise TimeoutError(prefixes)
token, raw = await asyncio.wait_for(q.get(), timeout=remaining)
if any(protocol.matches(token, p) for p in prefixes):
return token, raw