AnyBot/src/x11.rs

// Linux: seeing and using the screen through X11 (x11rb, pure Rust): screenshots with GetImage,
// mouse and keyboard with the XTEST extension, windows from the window manager's EWMH hints,
// idle time from the MIT-SCREEN-SAVER extension. The same functions as computer.rs's macOS part,
// so the computer_use tools work the same way.
//
// Wayland sessions go through wl.rs or portal.rs instead (desk.rs picks).

use std::sync::Mutex;
use std::time::Duration;
use x11rb::connection::Connection;
use x11rb::protocol::xproto::{self, AtomEnum, ConnectionExt as _, ImageFormat, Window};
use x11rb::protocol::xtest::ConnectionExt as _;
use x11rb::rust_connection::RustConnection;

struct X {
    conn: RustConnection,
    root: Window,
    width: u16,
    height: u16,
}

static CONN: Mutex<Option<X>> = Mutex::new(None);

fn with<T>(f: impl FnOnce(&X) -> Option<T>) -> Option<T> {
    let mut g = CONN.lock().unwrap_or_else(|e| e.into_inner());
    if g.is_none() {
        let (conn, n) = x11rb::connect(None).ok()?;
        let s = &conn.setup().roots[n];
        let (root, width, height) = (s.root, s.width_in_pixels, s.height_in_pixels);
        *g = Some(X { conn, root, width, height });
    }
    let r = f(g.as_ref()?);
    if r.is_none() {
        // A broken connection (the X server restarted): connect again next time.
        if g.as_ref().map(|x| x.conn.get_input_focus().ok().and_then(|c| c.reply().ok()).is_none()).unwrap_or(true) {
            *g = None;
        }
    }
    r
}

pub fn available() -> bool {
    with(|_| Some(())).is_some()
}

/// The whole screen as RGB.
pub fn capture() -> Result<image::RgbImage, String> {
    with(|x| {
        let reply = x.conn.get_image(ImageFormat::Z_PIXMAP, x.root, 0, 0, x.width, x.height, !0).ok()?.reply().ok()?;
        let setup = x.conn.setup();
        let bpp = setup.pixmap_formats.iter().find(|f| f.depth == reply.depth).map(|f| f.bits_per_pixel).unwrap_or(32);
        if bpp != 32 {
            return None;
        }
        let lsb = setup.image_byte_order == xproto::ImageOrder::LSB_FIRST;
        let (w, h) = (x.width as usize, x.height as usize);
        let mut rgb = Vec::with_capacity(w * h * 3);
        for px in reply.data.chunks_exact(4).take(w * h) {
            if lsb {
                rgb.extend_from_slice(&[px[2], px[1], px[0]]);
            } else {
                rgb.extend_from_slice(&[px[1], px[2], px[3]]);
            }
        }
        image::RgbImage::from_raw(w as u32, h as u32, rgb)
    })
    .ok_or_else(|| "讀不到螢幕畫面（需要 X11 桌面；Wayland 請在登入畫面選「Xorg」工作階段）".to_string())
}

pub fn screen_points() -> (f64, f64) {
    with(|x| Some((x.width as f64, x.height as f64))).unwrap_or((0.0, 0.0))
}

pub fn screen_pixels() -> (u32, u32) {
    with(|x| Some((x.width as u32, x.height as u32))).unwrap_or((0, 0))
}

fn fake(x: &X, kind: u8, detail: u8, px: i16, py: i16) {
    let _ = x.conn.xtest_fake_input(kind, detail, 0, x.root, px, py, 0);
}

fn flush(x: &X) {
    let _ = x.conn.flush();
    let _ = x.conn.get_input_focus().ok().and_then(|c| c.reply().ok());
}

fn sleep_secs(s: f64) {
    std::thread::sleep(Duration::from_secs_f64(s));
}

fn gap() {
    std::thread::sleep(Duration::from_millis(8));
}

pub fn move_to(px: f64, py: f64) {
    with(|x| {
        fake(x, xproto::MOTION_NOTIFY_EVENT, 0, px.round() as i16, py.round() as i16);
        flush(x);
        Some(())
    });
}

fn button_no(button: &str) -> u8 {
    match button {
        "right" => 3,
        "middle" => 2,
        _ => 1,
    }
}

/// Hold the modifiers in `flags` (macOS flag bits) while `f` runs.
fn with_mods(flags: u64, f: impl FnOnce()) {
    let mods = modifier_keysyms(flags);
    let codes: Vec<u8> = mods.iter().filter_map(|k| keycode_for(*k)).collect();
    with(|x| {
        for c in &codes {
            fake(x, xproto::KEY_PRESS_EVENT, *c, 0, 0);
        }
        flush(x);
        Some(())
    });
    f();
    with(|x| {
        for c in codes.iter().rev() {
            fake(x, xproto::KEY_RELEASE_EVENT, *c, 0, 0);
        }
        flush(x);
        Some(())
    });
}

pub fn click(px: f64, py: f64, button: &str, count: i64, flags: u64) {
    move_to(px, py);
    gap();
    let b = button_no(button);
    with_mods(flags, || {
        for n in 1..=count {
            if n > 1 {
                sleep_secs(crate::computer::CLICK_REPEAT_GAP_SECS);
            }
            with(|x| {
                fake(x, xproto::BUTTON_PRESS_EVENT, b, 0, 0);
                flush(x);
                Some(())
            });
            sleep_secs(crate::computer::CLICK_HOLD_SECS);
            with(|x| {
                fake(x, xproto::BUTTON_RELEASE_EVENT, b, 0, 0);
                flush(x);
                Some(())
            });
        }
    });
}

pub fn drag(path: &[(f64, f64)], button: &str, flags: u64) {
    let (x0, y0) = path[0];
    move_to(x0, y0);
    gap();
    let b = button_no(button);
    with_mods(flags, || {
        with(|x| {
            fake(x, xproto::BUTTON_PRESS_EVENT, b, 0, 0);
            flush(x);
            Some(())
        });
        sleep_secs(crate::computer::DRAG_HOLD_SECS);
        for (px, py) in crate::computer::drag_points(path) {
            gap();
            move_to(px, py);
        }
        with(|x| {
            fake(x, xproto::BUTTON_RELEASE_EVENT, b, 0, 0);
            flush(x);
            Some(())
        });
    });
}

/// Wheel clicks: one per scroll step (`SCROLL_PIXELS_PER_STEP` pixels on macOS).
pub fn scroll(at: Option<(f64, f64)>, dy: i32, dx: i32, flags: u64) {
    if let Some((px, py)) = at {
        move_to(px, py);
        gap();
    }
    let step = crate::computer::SCROLL_PIXELS_PER_STEP;
    let (vy, vx) = (dy / step, dx / step);
    with_mods(flags, || {
        with(|x| {
            for (n, up, down) in [(vy, 4u8, 5u8), (vx, 6, 7)] {
                let b = if n > 0 { up } else { down };
                for _ in 0..n.abs() {
                    fake(x, xproto::BUTTON_PRESS_EVENT, b, 0, 0);
                    fake(x, xproto::BUTTON_RELEASE_EVENT, b, 0, 0);
                    flush(x);
                    std::thread::sleep(Duration::from_millis(12));
                }
            }
            Some(())
        });
    });
}

// ---------- keyboard ----------

const XK_SHIFT_L: u32 = 0xffe1;
const XK_CONTROL_L: u32 = 0xffe3;
const XK_ALT_L: u32 = 0xffe9;

/// macOS flag bits → X keysyms. ⌘ becomes Control on Linux (⌘C = Ctrl+C, ⌘V = Ctrl+V…).
pub fn modifier_keysyms(flags: u64) -> Vec<u32> {
    use crate::computer::{FLAG_CMD, FLAG_CONTROL, FLAG_OPTION, FLAG_SHIFT};
    let mut v = vec![];
    if flags & (FLAG_CONTROL | FLAG_CMD) != 0 {
        v.push(XK_CONTROL_L);
    }
    if flags & FLAG_SHIFT != 0 {
        v.push(XK_SHIFT_L);
    }
    if flags & FLAG_OPTION != 0 {
        v.push(XK_ALT_L);
    }
    v
}

/// macOS virtual key code (what `computer::key_code` gives) → X keysym.
pub fn keysym_of_mac_code(code: u16) -> Option<u32> {
    let named: &[(u16, u32)] = &[
        (36, 0xff0d), (48, 0xff09), (49, 0x20), (51, 0xff08), (53, 0xff1b), (117, 0xffff), (115, 0xff50), (119, 0xff57), (116, 0xff55), (121, 0xff56),
        (123, 0xff51), (124, 0xff53), (125, 0xff54), (126, 0xff52), (122, 0xffbe), (120, 0xffbf), (99, 0xffc0), (118, 0xffc1), (96, 0xffc2), (97, 0xffc3),
        (98, 0xffc4), (100, 0xffc5), (101, 0xffc6), (109, 0xffc7), (103, 0xffc8), (111, 0xffc9),
    ];
    if let Some((_, k)) = named.iter().find(|(c, _)| *c == code) {
        return Some(*k);
    }
    // Letters, digits and punctuation: the character whose key code this is.
    "abcdefghijklmnopqrstuvwxyz0123456789-=[];',./\\`".chars().find(|c| crate::computer::key_code(&c.to_string()) == Some(code)).map(|c| c as u32)
}

/// The key that types a keysym in the current layout (first level only, so no Shift needed).
fn keycode_for(keysym: u32) -> Option<u8> {
    with(|x| {
        let setup = x.conn.setup();
        let (min, max) = (setup.min_keycode, setup.max_keycode);
        let m = x.conn.get_keyboard_mapping(min, max - min + 1).ok()?.reply().ok()?;
        let per = m.keysyms_per_keycode as usize;
        m.keysyms.chunks(per).position(|ks| ks.first() == Some(&keysym)).map(|i| min + i as u8)
    })
}

pub fn press(code: u16, flags: u64, hold_secs: f64) {
    let Some(sym) = keysym_of_mac_code(code) else { return };
    with_mods(flags, || {
        let Some(kc) = keycode_for(sym) else { return };
        with(|x| {
            fake(x, xproto::KEY_PRESS_EVENT, kc, 0, 0);
            flush(x);
            Some(())
        });
        sleep_secs(hold_secs);
        with(|x| {
            fake(x, xproto::KEY_RELEASE_EVENT, kc, 0, 0);
            flush(x);
            Some(())
        });
    });
}

/// The keysym that types a character.
pub fn keysym_for_char(c: char) -> u32 {
    match c {
        '\n' | '\r' => 0xff0d,
        '\t' => 0xff09,
        c if (' '..='~').contains(&c) || ('\u{a0}'..='\u{ff}').contains(&c) => c as u32,
        c => 0x0100_0000 + c as u32,
    }
}

/// Type any text (Chinese too): each character is put on a spare key for a moment and that key
/// is pressed — the way xdotool types characters the keyboard layout doesn't have.
pub fn type_text(text: &str) {
    with(|x| {
        let setup = x.conn.setup();
        let (min, max) = (setup.min_keycode, setup.max_keycode);
        let m = x.conn.get_keyboard_mapping(min, max - min + 1).ok()?.reply().ok()?;
        let per = m.keysyms_per_keycode as usize;
        // A key with nothing on it (from the top, where they usually are).
        let spare = m.keysyms.chunks(per).enumerate().rev().find(|(_, ks)| ks.iter().all(|k| *k == 0)).map(|(i, _)| min + i as u8)?;
        for c in text.chars() {
            let sym = keysym_for_char(c);
            let syms = vec![sym; per];
            // On any failure stop the loop and fall through to restoring the mapping: an
            // early `?` here left the typed character glued to a real key forever.
            if x.conn.change_keyboard_mapping(1, spare, per as u8, &syms).is_err() {
                break;
            }
            flush(x);
            std::thread::sleep(Duration::from_millis(12));
            fake(x, xproto::KEY_PRESS_EVENT, spare, 0, 0);
            flush(x);
            std::thread::sleep(Duration::from_millis(8));
            fake(x, xproto::KEY_RELEASE_EVENT, spare, 0, 0);
            flush(x);
            std::thread::sleep(Duration::from_millis(8));
        }
        let _ = x.conn.change_keyboard_mapping(1, spare, per as u8, &vec![0; per]);
        flush(x);
        Some(())
    });
}

// ---------- windows ----------

fn atom(x: &X, name: &str) -> Option<u32> {
    Some(x.conn.intern_atom(false, name.as_bytes()).ok()?.reply().ok()?.atom)
}

fn prop_u32s(x: &X, win: Window, name: &str, kind: AtomEnum) -> Vec<u32> {
    let Some(a) = atom(x, name) else { return vec![] };
    x.conn.get_property(false, win, a, kind, 0, 4096).ok().and_then(|c| c.reply().ok()).and_then(|r| r.value32().map(|v| v.collect())).unwrap_or_default()
}

/// A window's app: (name, pid). The name is the WM_CLASS class ("Firefox", "Gedit"…), else the
/// process name.
fn app_of(x: &X, win: Window) -> (String, i32) {
    let pid = prop_u32s(x, win, "_NET_WM_PID", AtomEnum::CARDINAL).first().copied().unwrap_or(0) as i32;
    let class = x
        .conn
        .get_property(false, win, AtomEnum::WM_CLASS, AtomEnum::STRING, 0, 256)
        .ok()
        .and_then(|c| c.reply().ok())
        .map(|r| r.value.split(|b| *b == 0).filter(|s| !s.is_empty()).map(|s| String::from_utf8_lossy(s).to_string()).collect::<Vec<_>>())
        .and_then(|v| v.last().cloned())
        .unwrap_or_default();
    let name = if !class.is_empty() { class } else if pid > 0 { std::fs::read_to_string(format!("/proc/{pid}/comm")).unwrap_or_default().trim().to_string() } else { String::new() };
    (name, pid)
}

/// Top-level app windows, front first, with their place on screen.
fn windows(x: &X) -> Vec<(Window, (i32, i32, i32, i32))> {
    let mut list = prop_u32s(x, x.root, "_NET_CLIENT_LIST_STACKING", AtomEnum::WINDOW);
    if list.is_empty() {
        list = x.conn.query_tree(x.root).ok().and_then(|c| c.reply().ok()).map(|t| t.children).unwrap_or_default();
    }
    list.reverse();
    list.into_iter()
        .filter_map(|w| {
            let g = x.conn.get_geometry(w).ok()?.reply().ok()?;
            let t = x.conn.translate_coordinates(w, x.root, 0, 0).ok()?.reply().ok()?;
            let attrs = x.conn.get_window_attributes(w).ok()?.reply().ok()?;
            (attrs.map_state == xproto::MapState::VIEWABLE && g.width > 1 && g.height > 1).then_some((w, (t.dst_x as i32, t.dst_y as i32, g.width as i32, g.height as i32)))
        })
        .collect()
}

pub fn window_at(px: f64, py: f64) -> Option<(String, i32)> {
    with(|x| {
        let (w, _) = windows(x).into_iter().find(|(_, (wx, wy, ww, wh))| px >= *wx as f64 && py >= *wy as f64 && px < (*wx + *ww) as f64 && py < (*wy + *wh) as f64)?;
        Some(app_of(x, w)).filter(|a| !a.0.is_empty())
    })
}

pub fn front_app() -> Option<(String, i32)> {
    with(|x| {
        let active = prop_u32s(x, x.root, "_NET_ACTIVE_WINDOW", AtomEnum::WINDOW).first().copied().filter(|w| *w != 0);
        let w = active.or_else(|| windows(x).first().map(|w| w.0))?;
        Some(app_of(x, w)).filter(|a| !a.0.is_empty())
    })
}

/// Seconds since the last keyboard or mouse input.
pub fn idle_secs() -> Option<f64> {
    use x11rb::protocol::screensaver::ConnectionExt as _;
    with(|x| Some(x.conn.screensaver_query_info(x.root).ok()?.reply().ok()?.ms_since_user_input as f64 / 1000.0))
}

/// The X screensaver is on (blanked screen).
pub fn screensaver_on() -> bool {
    use x11rb::protocol::screensaver::ConnectionExt as _;
    with(|x| Some(x.conn.screensaver_query_info(x.root).ok()?.reply().ok()?.state == 1)).unwrap_or(false)
}

// ---------- watching the person's input (示範 recording) ----------

/// A mouse click or key press the person made, from the X RECORD extension.
#[derive(Debug, Clone)]
pub enum Input {
    Click { x: f64, y: f64, right: bool },
    /// `mods`: 1 shift, 4 control, 8 alt, 64 super (X's modifier bits).
    Key { text: String, keysym: u32, mods: u16 },
    /// +1 up / −1 down.
    Scroll { dy: i64 },
}

/// Watch the person's clicks and keys until `keep_going` says stop. Runs on its own thread
/// with its own connections (RECORD needs one that only listens).
pub fn record(mut each: impl FnMut(Input), keep_going: impl Fn() -> bool + Send + 'static) -> Result<(), String> {
    use x11rb::protocol::record::{self, ConnectionExt as _};
    let (ctrl, _) = x11rb::connect(None).map_err(|e| format!("連不上 X11：{e}"))?;
    let (data, _) = x11rb::connect(None).map_err(|e| format!("連不上 X11：{e}"))?;
    let ctx = ctrl.generate_id().map_err(|e| e.to_string())?;
    let range = record::Range {
        device_events: record::Range8 { first: xproto::KEY_PRESS_EVENT, last: xproto::BUTTON_PRESS_EVENT },
        ..Default::default()
    };
    ctrl.record_create_context(ctx, 0, &[record::CS::ALL_CLIENTS.into()], &[range]).map_err(|e| e.to_string())?.check().map_err(|e| format!("這個 X 伺服器不支援錄製輸入（RECORD）：{e}"))?;
    // The keyboard map, to turn key codes into characters.
    let setup = ctrl.setup();
    let (min, max) = (setup.min_keycode, setup.max_keycode);
    let map = ctrl.get_keyboard_mapping(min, max - min + 1).map_err(|e| e.to_string())?.reply().map_err(|e| e.to_string())?;
    let per = map.keysyms_per_keycode as usize;
    let mut mods: u16 = 0;
    let mod_bit = |sym: u32| -> u16 {
        match sym {
            0xffe1 | 0xffe2 => 1,
            0xffe3 | 0xffe4 => 4,
            0xffe9 | 0xffea | 0xff7e => 8,
            0xffeb | 0xffec => 64,
            _ => 0,
        }
    };
    let stop = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
    let stopper = std::thread::spawn({
        let ctx = ctx;
        let stop = stop.clone();
        move || {
            // Disabling the context ends the data stream below.
            while keep_going() && !stop.load(std::sync::atomic::Ordering::Relaxed) {
                std::thread::sleep(Duration::from_millis(150));
            }
            let _ = ctrl.record_disable_context(ctx);
            let _ = ctrl.flush();
            std::thread::sleep(Duration::from_millis(200));
            let _ = ctrl.record_free_context(ctx);
            let _ = ctrl.flush();
        }
    });
    let stream = match data.record_enable_context(ctx) {
        Ok(s) => s,
        // A failed start must stop the watchdog thread too — it used to spin forever
        // (holding an X connection) on this early return.
        Err(e) => {
            stop.store(true, std::sync::atomic::Ordering::Relaxed);
            let _ = stopper.join();
            return Err(e.to_string());
        }
    };
    for reply in stream {
        let Ok(reply) = reply else { break };
        if reply.category != 0 {
            continue;
        }
        for ev in reply.data.chunks_exact(32) {
            let (kind, detail) = (ev[0] & 0x7f, ev[1]);
            let rx = i16::from_ne_bytes([ev[20], ev[21]]) as f64;
            let ry = i16::from_ne_bytes([ev[22], ev[23]]) as f64;
            let sym_at = |level: usize| -> u32 {
                let i = (detail.saturating_sub(min)) as usize * per + level;
                map.keysyms.get(i).copied().unwrap_or(0)
            };
            match kind {
                xproto::KEY_PRESS_EVENT => {
                    let base = sym_at(0);
                    let bit = mod_bit(base);
                    if bit != 0 {
                        mods |= bit;
                        continue;
                    }
                    let shifted = sym_at(1);
                    let sym = if mods & 1 != 0 && shifted != 0 { shifted } else { base };
                    let text = char_of_keysym(sym).map(String::from).unwrap_or_default();
                    each(Input::Key { text, keysym: sym, mods });
                }
                xproto::KEY_RELEASE_EVENT => {
                    mods &= !mod_bit(sym_at(0));
                }
                xproto::BUTTON_PRESS_EVENT => match detail {
                    1 | 3 => each(Input::Click { x: rx, y: ry, right: detail == 3 }),
                    4 => each(Input::Scroll { dy: 1 }),
                    5 => each(Input::Scroll { dy: -1 }),
                    _ => {}
                },
                _ => {}
            }
        }
    }
    // The stream broke (the X server died, the context was dropped): tell the watchdog to
    // stop first — waiting on a thread that keeps polling forever is what froze the old code.
    stop.store(true, std::sync::atomic::Ordering::Relaxed);
    let _ = stopper.join();
    Ok(())
}

pub fn char_of_keysym(sym: u32) -> Option<char> {
    match sym {
        0x20..=0x7e | 0xa0..=0xff => char::from_u32(sym),
        0x0100_0000..=0x0110_ffff => char::from_u32(sym - 0x0100_0000),
        _ => None,
    }
}

pub fn keysym_name(sym: u32) -> Option<&'static str> {
    Some(match sym {
        0xff0d => "return",
        0xff09 => "tab",
        0xff08 => "delete",
        0xff1b => "escape",
        0xffff => "forwarddelete",
        0xff51 => "left",
        0xff53 => "right",
        0xff52 => "up",
        0xff54 => "down",
        0xff50 => "home",
        0xff57 => "end",
        0xff55 => "pageup",
        0xff56 => "pagedown",
        _ => return None,
    })
}


#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn keys_map_to_x() {
        assert_eq!(keysym_of_mac_code(crate::computer::key_code("return").unwrap()), Some(0xff0d));
        assert_eq!(keysym_of_mac_code(crate::computer::key_code("a").unwrap()), Some('a' as u32));
        assert_eq!(keysym_of_mac_code(crate::computer::key_code("/").unwrap()), Some('/' as u32));
        assert_eq!(keysym_of_mac_code(crate::computer::key_code("f5").unwrap()), Some(0xffc2));
        assert_eq!(keysym_for_char('中'), 0x0100_4e2d);
        assert_eq!(char_of_keysym(0x0100_4e2d), Some('中'));
        assert_eq!(modifier_keysyms(crate::computer::FLAG_CMD | crate::computer::FLAG_SHIFT), vec![XK_CONTROL_L, XK_SHIFT_L]);
    }

    /// Needs a display: `Xvfb :99 & DISPLAY=:99 cargo test x11_live -- --ignored`.
    #[test]
    #[ignore]
    fn x11_live() {
        assert!(available(), "no X display");
        let (w, h) = screen_pixels();
        let shot = capture().expect("screenshot");
        assert_eq!((shot.width(), shot.height()), (w as u32, h as u32));
        // Record while clicking and typing through XTEST: the recorder must see both.
        let got = std::sync::Arc::new(Mutex::new(Vec::new()));
        let g = got.clone();
        let stop = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
        let s2 = stop.clone();
        let t = std::thread::spawn(move || record(move |i| g.lock().unwrap().push(format!("{i:?}")), move || !s2.load(std::sync::atomic::Ordering::SeqCst)));
        std::thread::sleep(Duration::from_millis(500));
        click(30.0, 40.0, "left", 1, 0);
        // Real keys (what a person presses); type_text's borrowed spare key has no symbol by the
        // time the recorder looks it up.
        press(crate::computer::key_code("a").unwrap(), 0, 0.03);
        press(crate::computer::key_code("z").unwrap(), crate::computer::FLAG_SHIFT, 0.03);
        press(crate::computer::key_code("c").unwrap(), crate::computer::FLAG_CONTROL, 0.03);
        scroll(None, -crate::computer::SCROLL_PIXELS_PER_STEP, 0, 0);
        std::thread::sleep(Duration::from_millis(500));
        stop.store(true, std::sync::atomic::Ordering::SeqCst);
        t.join().unwrap().unwrap();
        let got = got.lock().unwrap().join("\n");
        assert!(got.contains("Click { x: 30.0, y: 40.0, right: false }"), "{got}");
        assert!(got.contains("text: \"a\"") && got.contains("text: \"Z\""), "{got}");
        assert!(got.contains("keysym: 99, mods: 4"), "ctrl+c: {got}");
        assert!(got.contains("Scroll { dy: -1 }"), "{got}");
        println!("{got}");
    }
}