blekin/crates/ericrfb/examples/snapshot.rs

use std::env;
use std::fs::File;
use std::io::BufWriter;
use std::path::Path;

use ericrfb::handshake::Config;
use ericrfb::session::{ActiveSession, Event};

fn main() {
    let args: Vec<String> = env::args().collect();

    let host = args
        .iter()
        .position(|a| a == "--host")
        .and_then(|i| args.get(i + 1))
        .expect("usage: --host <ip> --applet-id <token> [--port <port>] [--output <file.png>]");

    let applet_id = args
        .iter()
        .position(|a| a == "--applet-id")
        .and_then(|i| args.get(i + 1))
        .expect("usage: --host <ip> --applet-id <token> [--port <port>] [--output <file.png>]");

    let port: u16 = args
        .iter()
        .position(|a| a == "--port")
        .and_then(|i| args.get(i + 1))
        .and_then(|s| s.parse().ok())
        .unwrap_or(443);

    let output = args
        .iter()
        .position(|a| a == "--output")
        .and_then(|i| args.get(i + 1).map(|s| s.as_str()))
        .unwrap_or("frame.png");

    let cfg = Config::new(host, port, applet_id);
    println!("Connecting to {}:{}...", cfg.host, cfg.port);

    // Only request Raw encoding (0) for Phase 3
    let mut session = ActiveSession::connect(&cfg, &[0]).expect("connect failed");
    println!(
        "Connected: {}x{}, waiting for first frame...",
        session.framebuffer.width, session.framebuffer.height
    );

    // Process messages until we get a FramebufferDirty event
    loop {
        match session.process_one() {
            Ok(Some(Event::FramebufferDirty)) => {
                println!("Got framebuffer update, saving to {output}");
                break;
            }
            Ok(Some(Event::Resize { width, height })) => {
                println!("Resized to {width}x{height}");
            }
            Ok(Some(Event::Debug(s))) => {
                eprintln!("[debug] {s}");
            }
            Ok(Some(Event::RfbCommand(k, v))) => {
                eprintln!("[rfb-cmd] {k}={v}");
            }
            Ok(Some(_)) => {}
            Ok(None) => {}
            Err(e) => {
                eprintln!("Error: {e}");
                std::process::exit(1);
            }
        }
    }

    // Write PNG
    let rgba = session.framebuffer.to_rgba();
    let w = session.framebuffer.width as u32;
    let h = session.framebuffer.height as u32;

    let path = Path::new(output);
    let file = File::create(path).expect("cannot create output file");
    let bw = BufWriter::new(file);

    let mut encoder = png::Encoder::new(bw, w, h);
    encoder.set_color(png::ColorType::Rgba);
    encoder.set_depth(png::BitDepth::Eight);
    let mut writer = encoder.write_header().expect("png header failed");
    writer.write_image_data(&rgba).expect("png write failed");

    println!("Saved {w}x{h} frame to {}", path.display());
}