clair/CVM-Sentry
1
0
Fork 0
Code Activity

YAY SCREENSHOTS WORK NOW

Browse Source
This commit is contained in:
Clair Delafuente
2025-10-13 13:30:09 -04:00
parent 906f26c220
commit 9727eba274
1 changed files with 88 additions and 70 deletions
Download Patch File Download Diff File Expand all files Collapse all files

158
cvmsentry.py
View File

@@ -63,9 +63,10 @@ async def periodic_snapshot_task():
log.info("Starting periodic snapshot task") log.info("Starting periodic snapshot task")
while True: while True:
try: try:
await asyncio.sleep(10) # Wait 10 seconds await asyncio.sleep(5)
log.debug("Running periodic framebuffer snapshot capture...") log.debug("Running periodic framebuffer snapshot capture...")
save_tasks = []
for vm_name, vm_data in vms.items(): for vm_name, vm_data in vms.items():
# Skip if VM doesn't have a framebuffer # Skip if VM doesn't have a framebuffer
if not vm_data.get("framebuffer"): if not vm_data.get("framebuffer"):
@@ -76,8 +77,8 @@ async def periodic_snapshot_task():
snapshot_dir = os.path.join("logs", "webp", vm_name, date_str) snapshot_dir = os.path.join("logs", "webp", vm_name, date_str)
os.makedirs(snapshot_dir, exist_ok=True) os.makedirs(snapshot_dir, exist_ok=True)
# Get current epoch timestamp # Get current epoch timestamp in milliseconds
epoch_timestamp = int(datetime.now().timestamp()) epoch_timestamp = int(datetime.now().timestamp() * 1000)
filename = f"{epoch_timestamp}.webp" filename = f"{epoch_timestamp}.webp"
filepath = os.path.join(snapshot_dir, filename) filepath = os.path.join(snapshot_dir, filename)
@@ -91,18 +92,37 @@ async def periodic_snapshot_task():
# Only save if the framebuffer has changed since last snapshot # Only save if the framebuffer has changed since last snapshot
if current_hash != vm_data.get("last_frame_hash"): if current_hash != vm_data.get("last_frame_hash"):
try: # Create a copy of the image to avoid race conditions
# Save as webp with good quality/size balance img_copy = framebuffer.copy()
framebuffer.save(filepath, format="WEBP", quality=65, method=6) # Create and store task for asynchronous saving
vm_data["last_frame_hash"] = current_hash save_tasks.append(asyncio.create_task(
log.debug(f"Saved snapshot of {vm_name} to {filepath}") save_image_async(img_copy, filepath, vm_name, vm_data, current_hash)
except Exception as e: ))
log.error(f"Failed to save snapshot for {vm_name}: {e}")
# Wait for all save tasks to complete
if save_tasks:
await asyncio.gather(*save_tasks)
except Exception as e: except Exception as e:
log.error(f"Error in periodic snapshot task: {e}") log.error(f"Error in periodic snapshot task: {e}")
# Continue running even if there's an error # Continue running even if there's an error
async def save_image_async(image, filepath, vm_name, vm_data, current_hash):
"""Save an image to disk asynchronously."""
try:
# Run the image saving in a thread pool to avoid blocking
await asyncio.to_thread(
image.save,
filepath,
format="WEBP",
quality=65,
method=6
)
vm_data["last_frame_hash"] = current_hash
log.info(f"Saved snapshot of {vm_name} to {filepath}")
except Exception as e:
log.error(f"Failed to save snapshot for {vm_name}: {e}")
async def connect(vm_name: str): async def connect(vm_name: str):
global STATE global STATE
@@ -310,68 +330,66 @@ async def connect(vm_name: str):
case ["size", "0", width, height]: case ["size", "0", width, height]:
log.debug(f"({STATE.name} - {vm_name}) !!! Framebuffer size update: {width}x{height} !!!") log.debug(f"({STATE.name} - {vm_name}) !!! Framebuffer size update: {width}x{height} !!!")
vms[vm_name]["size"] = (int(width), int(height)) vms[vm_name]["size"] = (int(width), int(height))
case ["png", "0", "0", "0", "0", initial_frame_b64]: case ["png", "0", "0", "0", "0", full_frame_b64]:
if STATE < CollabVMState.LOGGED_IN: try:
try: log.debug(f"({STATE.name} - {vm_name}) !!! Received full framebuffer update !!!")
log.debug(f"({STATE.name} - {vm_name}) !!! Initial framebuffer received !!!") expected_width, expected_height = vms[vm_name]["size"]
image_data = base64.b64decode(initial_frame_b64)
framebuffer = Image.open(BytesIO(image_data)) # Decode the base64 data to get the PNG image
vms[vm_name]["framebuffer"] = framebuffer frame_data = base64.b64decode(full_frame_b64)
log.debug(f"({STATE.name} - {vm_name}) Initial framebuffer loaded: {framebuffer.size}") frame_img = Image.open(BytesIO(frame_data))
except Exception as e:
log.error(f"({STATE.name} - {vm_name}) Failed to process initial framebuffer: {e}") # Validate image size and handle partial frames
elif STATE >= CollabVMState.LOGGED_IN: if expected_width > 0 and expected_height > 0:
try: if frame_img.size != (expected_width, expected_height):
image_data = base64.b64decode(initial_frame_b64) log.warning(f"({STATE.name} - {vm_name}) Partial framebuffer update: "
framebuffer = Image.open(BytesIO(image_data)) f"expected {expected_width}x{expected_height}, got {frame_img.size}")
# Get expected size from VM info # Create a new image of expected size if no framebuffer exists
expected_width, expected_height = vms[vm_name]["size"] if vms[vm_name]["framebuffer"] is None:
vms[vm_name]["framebuffer"] = Image.new('RGB', (expected_width, expected_height))
# Ensure the framebuffer matches the expected size
if framebuffer.size != (expected_width, expected_height) and expected_width > 0 and expected_height > 0: # Only update the portion that was received
framebuffer = framebuffer.resize((expected_width, expected_height)) if vms[vm_name]["framebuffer"]:
# Create a copy of the current framebuffer to modify
vms[vm_name]["framebuffer"] = framebuffer updated_img = vms[vm_name]["framebuffer"].copy()
# Paste the new partial frame at position (0,0)
log.debug(f"({STATE.name} - {vm_name}) Framebuffer updated: {framebuffer.size}") updated_img.paste(frame_img, (0, 0))
except Exception as e: # Use this as our new framebuffer
log.error(f"({STATE.name} - {vm_name}) Failed to process framebuffer update: {e}") frame_img = updated_img
# Update the framebuffer with the new image
vms[vm_name]["framebuffer"] = frame_img
log.debug(f"({STATE.name} - {vm_name}) Framebuffer updated with full frame, size: {frame_img.size}")
except Exception as e:
log.error(f"({STATE.name} - {vm_name}) Failed to process full framebuffer update: {e}")
case ["png", "0", "0", x, y, rect_b64]: case ["png", "0", "0", x, y, rect_b64]:
try:
if vms[vm_name]["framebuffer"] is not None: log.debug(f"({STATE.name} - {vm_name}) Received partial framebuffer update at position ({x}, {y})")
try: x, y = int(x), int(y)
# Check if we have a framebuffer with the right size
log.debug(f"({STATE.name} - {vm_name}) !!! Framebuffer rectangle update at ({x},{y}) !!!") # Decode the base64 data to get the PNG image fragment
expected_width, expected_height = vms[vm_name]["size"] frame_data = base64.b64decode(rect_b64)
if (vms[vm_name]["framebuffer"].size != (expected_width, expected_height) and fragment_img = Image.open(BytesIO(frame_data))
expected_width > 0 and expected_height > 0):
log.warning(f"({STATE.name} - {vm_name}) Framebuffer size mismatch: " # If we don't have a framebuffer yet or it's incompatible, create one
f"expected {expected_width}x{expected_height}, got {vms[vm_name]['framebuffer'].size}") if vms[vm_name]["framebuffer"] is None:
continue # drop
continue
# Decode and paste the rectangle update
x_pos, y_pos = int(x), int(y) # If we have a valid framebuffer, update it with the fragment
rect_data = base64.b64decode(rect_b64) if vms[vm_name]["framebuffer"]:
rect_img = Image.open(BytesIO(rect_data)) # Create a copy to modify
updated_img = vms[vm_name]["framebuffer"].copy()
# Create a copy of the framebuffer to work with # Paste the fragment at the specified position
framebuffer = vms[vm_name]["framebuffer"].copy() updated_img.paste(fragment_img, (x, y))
framebuffer.paste(rect_img, (x_pos, y_pos))
# Verify the framebuffer size is consistent
if rect_img.size[0] + x_pos > framebuffer.size[0] or rect_img.size[1] + y_pos > framebuffer.size[1]:
log.warning(f"({STATE.name} - {vm_name}) Rectangle update would exceed framebuffer bounds: "
f"rect size {rect_img.size} at position ({x_pos},{y_pos}) exceeds framebuffer {framebuffer.size}")
continue
# Update the framebuffer # Update the framebuffer
vms[vm_name]["framebuffer"] = framebuffer vms[vm_name]["framebuffer"] = updated_img
log.debug(f"({STATE.name} - {vm_name}) Updated framebuffer with fragment at ({x}, {y}), fragment size: {fragment_img.size}")
except Exception as e: else:
log.error(f"({STATE.name} - {vm_name}) Failed to process framebuffer rectangle update: {e}") log.warning(f"({STATE.name} - {vm_name}) Cannot update framebuffer - no base framebuffer exists")
else: except Exception as e:
log.debug(f"({STATE.name} - {vm_name}) Received rectangle update but framebuffer is not initialized") log.error(f"({STATE.name} - {vm_name}) Failed to process partial framebuffer update: {e}")
case ["turn", turn_time, count, current_turn, *queue]: case ["turn", turn_time, count, current_turn, *queue]:
if ( if (
queue == vms[vm_name]["turn_queue"] queue == vms[vm_name]["turn_queue"]