windows-troll 0.1.0

Modular Windows prank library
//! Screen capture helpers, ported from the `screen-messy-fun` project.

use windows::Win32::Foundation::GetLastError;
use windows::Win32::Graphics::Gdi::{
    BitBlt, CreateCompatibleBitmap, CreateCompatibleDC, DeleteDC, DeleteObject, GetWindowDC,
    ReleaseDC, SelectObject, HBITMAP, HDC, HGDIOBJ, ROP_CODE, SRCCOPY,
};
use windows::Win32::UI::WindowsAndMessaging::GetDesktopWindow;

fn last_error() -> String {
    format!("Windows error {}", unsafe { GetLastError().0 })
}

/// Captures a rectangular region of the desktop and can blit it back anywhere
/// on the screen (i.e. move the captured pixels around).
pub struct ScreenCapture {
    hdc: HDC,
    h_compatible_dc: HDC,
    h_bitmap: HBITMAP,
    width: i32,
    height: i32,
}

impl ScreenCapture {
    pub fn new(x: i32, y: i32, width: i32, height: i32) -> Result<Self, String> {
        unsafe {
            let desktop = GetDesktopWindow();
            let hdc = GetWindowDC(Some(desktop));
            if hdc.0.is_null() {
                return Err(format!("Failed to get window DC: {}", last_error()));
            }

            let h_compatible_dc = CreateCompatibleDC(Some(hdc));
            if h_compatible_dc.0.is_null() {
                ReleaseDC(Some(desktop), hdc);
                return Err(format!("Failed to create compatible DC: {}", last_error()));
            }

            let h_bitmap = CreateCompatibleBitmap(hdc, width, height);
            if h_bitmap.0.is_null() {
                let _ = DeleteDC(h_compatible_dc);
                ReleaseDC(Some(desktop), hdc);
                return Err(format!("Failed to create compatible bitmap: {}", last_error()));
            }

            let old_obj = SelectObject(h_compatible_dc, HGDIOBJ(h_bitmap.0));

            if BitBlt(
                h_compatible_dc,
                0,
                0,
                width,
                height,
                Some(hdc),
                x,
                y,
                ROP_CODE(SRCCOPY.0),
            )
            .is_err()
            {
                let _ = DeleteObject(HGDIOBJ(h_bitmap.0));
                let _ = DeleteDC(h_compatible_dc);
                ReleaseDC(Some(desktop), hdc);
                return Err(format!("Failed to capture screen: {}", last_error()));
            }

            SelectObject(h_compatible_dc, old_obj);

            Ok(ScreenCapture {
                hdc,
                h_compatible_dc,
                h_bitmap,
                width,
                height,
            })
        }
    }

    /// Draws this captured region onto the desktop at `(x, y)`.
    pub fn render_at(&self, x: i32, y: i32) -> Result<(), String> {
        unsafe {
            let old = SelectObject(self.h_compatible_dc, HGDIOBJ(self.h_bitmap.0));

            if BitBlt(
                self.hdc,
                x,
                y,
                self.width,
                self.height,
                Some(self.h_compatible_dc),
                0,
                0,
                ROP_CODE(SRCCOPY.0),
            )
            .is_err()
            {
                SelectObject(self.h_compatible_dc, old);
                return Err(format!(
                    "Failed to render screen at ({}, {}): {}",
                    x,
                    y,
                    last_error()
                ));
            }

            SelectObject(self.h_compatible_dc, old);
            Ok(())
        }
    }
}

impl Drop for ScreenCapture {
    fn drop(&mut self) {
        unsafe {
            let _ = DeleteObject(HGDIOBJ(self.h_bitmap.0));
            let _ = DeleteDC(self.h_compatible_dc);
            ReleaseDC(Some(GetDesktopWindow()), self.hdc);
        }
    }
}