From a143b79c9730453cde8c3c98ed6414eefcdaecc4 Mon Sep 17 00:00:00 2001 From: Michael Zhang Date: Mon, 10 Sep 2018 03:15:42 -0500 Subject: [PATCH] yay capturing works --- .gdb_history | 3 --- Cargo.lock | 70 ++++++++++++++++++++++++++++++++++++++++++++++++++ Cargo.toml | 2 ++ src/capture.rs | 7 +++-- src/errors.rs | 21 +++++++++++++++ src/gui.rs | 57 ++++++++++++++++++++++++---------------- src/image.rs | 55 ++++++++++++++++++++++++++++++++++++--- src/lib.rs | 12 +++------ src/main.rs | 18 ++++++------- 9 files changed, 195 insertions(+), 50 deletions(-) delete mode 100644 .gdb_history diff --git a/.gdb_history b/.gdb_history deleted file mode 100644 index 8785128..0000000 --- a/.gdb_history +++ /dev/null @@ -1,3 +0,0 @@ -run -run -o shiet.png fullscreen -bt diff --git a/Cargo.lock b/Cargo.lock index 8880252..0013c1b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1,3 +1,8 @@ +[[package]] +name = "adler32" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" + [[package]] name = "ansi_term" version = "0.11.0" @@ -42,6 +47,11 @@ name = "bitflags" version = "1.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" +[[package]] +name = "byteorder" +version = "1.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" + [[package]] name = "cc" version = "1.0.15" @@ -66,6 +76,15 @@ dependencies = [ "vec_map 0.8.1 (registry+https://github.com/rust-lang/crates.io-index)", ] +[[package]] +name = "deflate" +version = "0.7.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +dependencies = [ + "adler32 1.0.3 (registry+https://github.com/rust-lang/crates.io-index)", + "byteorder 1.2.6 (registry+https://github.com/rust-lang/crates.io-index)", +] + [[package]] name = "failure" version = "0.1.1" @@ -85,16 +104,57 @@ dependencies = [ "synstructure 0.6.1 (registry+https://github.com/rust-lang/crates.io-index)", ] +[[package]] +name = "inflate" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +dependencies = [ + "adler32 1.0.3 (registry+https://github.com/rust-lang/crates.io-index)", +] + [[package]] name = "libc" version = "0.2.40" source = "registry+https://github.com/rust-lang/crates.io-index" +[[package]] +name = "num-integer" +version = "0.1.39" +source = "registry+https://github.com/rust-lang/crates.io-index" +dependencies = [ + "num-traits 0.2.5 (registry+https://github.com/rust-lang/crates.io-index)", +] + +[[package]] +name = "num-iter" +version = "0.1.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +dependencies = [ + "num-integer 0.1.39 (registry+https://github.com/rust-lang/crates.io-index)", + "num-traits 0.2.5 (registry+https://github.com/rust-lang/crates.io-index)", +] + +[[package]] +name = "num-traits" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" + [[package]] name = "pkg-config" version = "0.3.11" source = "registry+https://github.com/rust-lang/crates.io-index" +[[package]] +name = "png" +version = "0.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +dependencies = [ + "bitflags 1.0.3 (registry+https://github.com/rust-lang/crates.io-index)", + "deflate 0.7.18 (registry+https://github.com/rust-lang/crates.io-index)", + "inflate 0.4.3 (registry+https://github.com/rust-lang/crates.io-index)", + "num-iter 0.1.37 (registry+https://github.com/rust-lang/crates.io-index)", +] + [[package]] name = "proc-macro2" version = "0.4.19" @@ -139,6 +199,8 @@ name = "screenshot" version = "0.3.0" dependencies = [ "failure 0.1.1 (registry+https://github.com/rust-lang/crates.io-index)", + "libc 0.2.40 (registry+https://github.com/rust-lang/crates.io-index)", + "png 0.12.0 (registry+https://github.com/rust-lang/crates.io-index)", "structopt 0.2.10 (registry+https://github.com/rust-lang/crates.io-index)", "time 0.1.40 (registry+https://github.com/rust-lang/crates.io-index)", "x11 2.18.1 (registry+https://github.com/rust-lang/crates.io-index)", @@ -282,18 +344,26 @@ dependencies = [ ] [metadata] +"checksum adler32 1.0.3 (registry+https://github.com/rust-lang/crates.io-index)" = "7e522997b529f05601e05166c07ed17789691f562762c7f3b987263d2dedee5c" "checksum ansi_term 0.11.0 (registry+https://github.com/rust-lang/crates.io-index)" = "ee49baf6cb617b853aa8d93bf420db2383fab46d314482ca2803b40d5fde979b" "checksum atty 0.2.10 (registry+https://github.com/rust-lang/crates.io-index)" = "2fc4a1aa4c24c0718a250f0681885c1af91419d242f29eb8f2ab28502d80dbd1" "checksum backtrace 0.3.7 (registry+https://github.com/rust-lang/crates.io-index)" = "8ea58cd16fd6c9d120b5bcb01d63883ae4cc7ba2aed35c1841b862a3c7ef6639" "checksum backtrace-sys 0.1.16 (registry+https://github.com/rust-lang/crates.io-index)" = "44585761d6161b0f57afc49482ab6bd067e4edef48c12a152c237eb0203f7661" "checksum bitflags 1.0.3 (registry+https://github.com/rust-lang/crates.io-index)" = "d0c54bb8f454c567f21197eefcdbf5679d0bd99f2ddbe52e84c77061952e6789" +"checksum byteorder 1.2.6 (registry+https://github.com/rust-lang/crates.io-index)" = "90492c5858dd7d2e78691cfb89f90d273a2800fc11d98f60786e5d87e2f83781" "checksum cc 1.0.15 (registry+https://github.com/rust-lang/crates.io-index)" = "0ebb87d1116151416c0cf66a0e3fb6430cccd120fd6300794b4dfaa050ac40ba" "checksum cfg-if 0.1.3 (registry+https://github.com/rust-lang/crates.io-index)" = "405216fd8fe65f718daa7102ea808a946b6ce40c742998fbfd3463645552de18" "checksum clap 2.31.2 (registry+https://github.com/rust-lang/crates.io-index)" = "f0f16b89cbb9ee36d87483dc939fe9f1e13c05898d56d7b230a0d4dff033a536" +"checksum deflate 0.7.18 (registry+https://github.com/rust-lang/crates.io-index)" = "32c8120d981901a9970a3a1c97cf8b630e0fa8c3ca31e75b6fd6fd5f9f427b31" "checksum failure 0.1.1 (registry+https://github.com/rust-lang/crates.io-index)" = "934799b6c1de475a012a02dab0ace1ace43789ee4b99bcfbf1a2e3e8ced5de82" "checksum failure_derive 0.1.1 (registry+https://github.com/rust-lang/crates.io-index)" = "c7cdda555bb90c9bb67a3b670a0f42de8e73f5981524123ad8578aafec8ddb8b" +"checksum inflate 0.4.3 (registry+https://github.com/rust-lang/crates.io-index)" = "6f53b811ee8e2057ccf9643ca6b4277de90efaf5e61e55fd5254576926bb4245" "checksum libc 0.2.40 (registry+https://github.com/rust-lang/crates.io-index)" = "6fd41f331ac7c5b8ac259b8bf82c75c0fb2e469bbf37d2becbba9a6a2221965b" +"checksum num-integer 0.1.39 (registry+https://github.com/rust-lang/crates.io-index)" = "e83d528d2677f0518c570baf2b7abdcf0cd2d248860b68507bdcb3e91d4c0cea" +"checksum num-iter 0.1.37 (registry+https://github.com/rust-lang/crates.io-index)" = "af3fdbbc3291a5464dc57b03860ec37ca6bf915ed6ee385e7c6c052c422b2124" +"checksum num-traits 0.2.5 (registry+https://github.com/rust-lang/crates.io-index)" = "630de1ef5cc79d0cdd78b7e33b81f083cbfe90de0f4b2b2f07f905867c70e9fe" "checksum pkg-config 0.3.11 (registry+https://github.com/rust-lang/crates.io-index)" = "110d5ee3593dbb73f56294327fe5668bcc997897097cbc76b51e7aed3f52452f" +"checksum png 0.12.0 (registry+https://github.com/rust-lang/crates.io-index)" = "f54b9600d584d3b8a739e1662a595fab051329eff43f20e7d8cc22872962145b" "checksum proc-macro2 0.4.19 (registry+https://github.com/rust-lang/crates.io-index)" = "ffe022fb8c8bd254524b0b3305906c1921fa37a84a644e29079a9e62200c3901" "checksum quote 0.3.15 (registry+https://github.com/rust-lang/crates.io-index)" = "7a6e920b65c65f10b2ae65c831a81a073a89edd28c7cce89475bff467ab4167a" "checksum quote 0.6.8 (registry+https://github.com/rust-lang/crates.io-index)" = "dd636425967c33af890042c483632d33fa7a18f19ad1d7ea72e8998c6ef8dea5" diff --git a/Cargo.toml b/Cargo.toml index 18d05b6..e13db52 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -9,6 +9,8 @@ crate-type = ["dylib", "rlib"] [dependencies] failure = "0.1" +png = "0.12" +libc = "0.2" structopt = "0.2" time = "0.1" x11 = { version = "2.18", features = ["xlib"] } diff --git a/src/capture.rs b/src/capture.rs index 1ed8c6b..6a76d8d 100644 --- a/src/capture.rs +++ b/src/capture.rs @@ -5,16 +5,16 @@ use image::Image; use options::{Options, Region}; pub fn capture(opt: &Options) -> Result { - let gui = GUI::new(); + let gui = GUI::new()?; let window_to_capture = match opt.region { Region::Fullscreen | Region::Selection => gui.get_root_window(), Region::ActiveWindow => gui.get_active_window(), }; - println!("Got window_to_capture: {:?}", window_to_capture); + println!("capturing window: {}", window_to_capture); let mut capture = gui.window_capture(window_to_capture)?; - println!("Captured the window."); + println!("captured the window"); // let final_capture = match opt.region { // Region::Fullscreen | Region::ActiveWindow => window_capture, @@ -24,7 +24,6 @@ pub fn capture(opt: &Options) -> Result { if let Region::Selection = opt.region { let region = gui.interactive_select(&capture)?; capture.apply_region(®ion); - println!("Subimaged."); }; Ok(capture) diff --git a/src/errors.rs b/src/errors.rs index 9dc70d2..dc63dbd 100644 --- a/src/errors.rs +++ b/src/errors.rs @@ -1,5 +1,26 @@ #[derive(Debug, Fail)] pub enum ScreenshotError { + #[fail(display = "x11 error: {}", message)] + XError { message: String }, + + #[fail(display = "io error")] + IOError(#[cause] ::std::io::Error), + + #[fail(display = "png encoding error")] + PngEncodingError(#[cause] ::png::EncodingError), + #[fail(display = "error")] Error, } + +impl From<::std::io::Error> for ScreenshotError { + fn from(err: ::std::io::Error) -> Self { + ScreenshotError::IOError(err) + } +} + +impl From<::png::EncodingError> for ScreenshotError { + fn from(err: ::png::EncodingError) -> Self { + ScreenshotError::PngEncodingError(err) + } +} diff --git a/src/gui.rs b/src/gui.rs index 9242db1..22ba62a 100644 --- a/src/gui.rs +++ b/src/gui.rs @@ -1,5 +1,6 @@ use std::ffi::CString; +use libc; use x11::xlib::{self, *}; use errors::ScreenshotError; @@ -11,50 +12,62 @@ pub struct GUI { } impl GUI { - pub fn new() -> Self { + pub fn new() -> Result { let display_str = CString::new(":0").unwrap(); let display = unsafe { xlib::XOpenDisplay(display_str.as_ptr()) }; + if display.is_null() { + return Err(ScreenshotError::XError { + message: format!("failed to open display"), + }); + } unsafe { XGrabServer(display) }; - GUI { display } + Ok(GUI { display }) } - fn get_window_attributes(&self, window: *mut Window) -> *mut XWindowAttributes { - let attr: *mut XWindowAttributes = unsafe { ::std::mem::uninitialized() }; - unsafe { XGetWindowAttributes(self.display, *window, attr) }; - attr + fn get_window_attributes( + &self, + window: Window, + ) -> Result<*mut XWindowAttributes, ScreenshotError> { + let attr_size = ::std::mem::size_of::(); + let attr = unsafe { libc::malloc(attr_size) as *mut XWindowAttributes }; + let result = unsafe { XGetWindowAttributes(self.display, window, attr) }; + if result == 0 { + return Err(ScreenshotError::XError { + message: format!("failed to get window attributes"), + }); + } + Ok(attr) } /// Captures the window and produces a DynamicImage. - pub fn window_capture(&self, window: *mut Window) -> Result { - println!("Getting window attributes."); - let attr = unsafe { *self.get_window_attributes(window) }; - println!("Got window attributes."); + pub fn window_capture(&self, window: Window) -> Result { + let attr = self.get_window_attributes(window)?; + println!("got window attributes"); let image = unsafe { xlib::XGetImage( self.display, - *window, - attr.x, - attr.y, - attr.width as u32, - attr.height as u32, + window, + (*attr).x, + (*attr).y, + (*attr).width as u32, + (*attr).height as u32, 0xffffffff, ZPixmap, ) }; - Ok(Image::from(image)) + Ok(Image::from(self.display, image)) } /// Get the full screen. - pub fn get_root_window(&self) -> *mut Window { - println!("Getting root window..."); - unsafe { xlib::XRootWindow(self.display, 0) as *mut Window } + pub fn get_root_window(&self) -> Window { + unsafe { xlib::XRootWindow(self.display, 0) as Window } } /// Get the active window. - pub fn get_active_window(&self) -> *mut Window { - let window: *mut Window = unsafe { ::std::mem::uninitialized() }; + pub fn get_active_window(&self) -> Window { + let mut window: Window = unsafe { ::std::mem::uninitialized() }; let mut revert_to_return: i32 = unsafe { ::std::mem::uninitialized() }; - unsafe { xlib::XGetInputFocus(self.display, window, &mut revert_to_return) }; + unsafe { xlib::XGetInputFocus(self.display, &mut window, &mut revert_to_return) }; window } diff --git a/src/image.rs b/src/image.rs index 11c3fa7..ae23142 100644 --- a/src/image.rs +++ b/src/image.rs @@ -1,15 +1,62 @@ -use x11::xlib::XImage; +use std::fs::File; +use std::io::BufWriter; +use std::path::Path; +use png::{self, Encoder, HasParameters}; +use x11::xlib::*; + +use errors::ScreenshotError; use Rectangle; +#[allow(dead_code)] pub struct Image { - _inner: *mut XImage, + display: *mut Display, + inner: *mut XImage, } impl Image { - pub fn from(inner: *mut XImage) -> Self { - Image { _inner: inner } + pub fn from(display: *mut Display, inner: *mut XImage) -> Self { + Image { inner, display } } pub fn apply_region(&mut self, _region: &Rectangle) {} + + /// Converts the image buffer into RGB(A). + fn to_data_buf(&self) -> Vec { + let im = unsafe { *self.inner }; + let size = 4usize * im.width as usize * im.height as usize; + let mut buf = vec![1; size]; + let sbuf = unsafe { ::std::slice::from_raw_parts(im.data, size) }; // source buffer + let mut sx = 0usize; // source idx + let mut dx = 0usize; // dest idx + if im.byte_order == LSBFirst { + while dx < size { + buf[dx] = sbuf[sx + 2] as u8; + buf[dx + 1] = sbuf[sx + 1] as u8; + buf[dx + 2] = sbuf[sx] as u8; + buf[dx + 3] = if im.depth == 32 { + sbuf[sx + 3] as u8 + } else { + 255u8 + }; + sx += 4; + dx += 4; + } + } + buf + } + + pub fn write_png(&self, out: impl AsRef) -> Result<(), ScreenshotError> { + let file = File::create(out.as_ref())?; + let ref mut out = BufWriter::new(file); + + let (width, height) = unsafe { ((*self.inner).width as u32, (*self.inner).height as u32) }; + let mut encoder = Encoder::new(out, width, height); + encoder.set(png::ColorType::RGBA).set(png::BitDepth::Eight); + let mut writer = encoder.write_header()?; + + let data = self.to_data_buf(); + writer.write_image_data(data.as_slice())?; + Ok(()) + } } diff --git a/src/lib.rs b/src/lib.rs index 34cb934..dace8a4 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,5 +1,7 @@ #[macro_use] extern crate failure; +extern crate libc; +extern crate png; #[macro_use] extern crate structopt; extern crate time; @@ -11,10 +13,10 @@ mod gui; mod image; mod options; -use failure::Error; use structopt::StructOpt; -use options::Options; +pub use capture::capture; +pub use options::Options; pub struct Rectangle { pub x: u32, @@ -22,9 +24,3 @@ pub struct Rectangle { pub width: u32, pub height: u32, } - -pub fn run() -> Result<(), Error> { - let opt = Options::from_args(); - capture::capture(&opt)?; - Ok(()) -} diff --git a/src/main.rs b/src/main.rs index 6f02a3a..09a8ff2 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,13 +1,13 @@ +extern crate failure; extern crate screenshot; +extern crate structopt; -use std::process::exit; +use failure::Error; +use screenshot::{capture, Options}; +use structopt::StructOpt; -fn main() { - match screenshot::run() { - Ok(()) => (), - Err(error) => { - eprintln!("Error: {}", error); - exit(1); - } - } +pub fn main() -> Result<(), Error> { + let opt = Options::from_args(); + let image = capture(&opt)?; + image.write_png(opt.outfile).map_err(|err| err.into()) }