From 976e203a4cca1a756287643d5a378accccbb5086 Mon Sep 17 00:00:00 2001 From: David Young Date: Sun, 6 Nov 2022 15:12:36 +0000 Subject: [PATCH 1/2] Convert most things to async --- .github/workflows/release.yml | 4 +-- Cargo.lock | 57 +++++++++++++++++++++++++++++++++++ Cargo.toml | 1 + src/main.rs | 43 ++++++++++++-------------- src/rdp/mod.rs | 12 ++++---- src/reporting.rs | 9 +++--- src/vnc2.rs | 57 +++++++++++++++++++++++++++++++++++ src/web/chrome.rs | 3 +- src/web/mod.rs | 4 +-- 9 files changed, 152 insertions(+), 38 deletions(-) create mode 100644 src/vnc2.rs diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 4ffca65..4d3687b 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -74,8 +74,8 @@ jobs: run: | cargo deb # Fix the tilde-name that happens on alpha builds - mv target/debian/scrying_*_amd64.deb \ - target/debian/scrying_${{ needs.create_new_release.outputs.version_num }}_amd64.deb + #mv target/debian/scrying_*_amd64.deb \ + # target/debian/scrying_${{ needs.create_new_release.outputs.version_num }}_amd64.deb - name: Zip binary run: | diff --git a/Cargo.lock b/Cargo.lock index 9a7f91c..b38e8ef 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -26,6 +26,12 @@ dependencies = [ "memchr", ] +[[package]] +name = "anyhow" +version = "1.0.66" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "216261ddc8289130e551ddcd5ce8a064710c0d064a4d2895c67151c92b5443f6" + [[package]] name = "arrayvec" version = "0.5.2" @@ -1351,6 +1357,29 @@ version = "3.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c1b04fb49957986fdce4d6ee7a65027d55d4b6d2265e5848bbb507b58ccfdb6f" +[[package]] +name = "parking_lot" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3742b2c103b9f06bc9fff0a37ff4912935851bee6d36f3c02bcc755bcfec228f" +dependencies = [ + "lock_api", + "parking_lot_core", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4dc9e0dc2adc1c69d09143aff38d3d30c5c3f0df0dad82e6d25547af174ebec0" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-sys 0.42.0", +] + [[package]] name = "percent-encoding" version = "2.2.0" @@ -1696,6 +1725,7 @@ dependencies = [ "tokio", "url", "vnc", + "vnc-rs", ] [[package]] @@ -1772,6 +1802,15 @@ dependencies = [ "lazy_static", ] +[[package]] +name = "signal-hook-registry" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e51e73328dc4ac0c7ccbda3a494dfa03df1de2f46018127f60c693f2648455b0" +dependencies = [ + "libc", +] + [[package]] name = "simplelog" version = "0.12.0" @@ -2015,10 +2054,14 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a9e03c497dc955702ba729190dc4aac6f2a0ce97f913e5b1b5912fc5039d9099" dependencies = [ "autocfg", + "bytes", "libc", + "memchr", "mio", "num_cpus", + "parking_lot", "pin-project-lite", + "signal-hook-registry", "socket2", "tokio-macros", "winapi", @@ -2051,6 +2094,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8ce8c33a8d48bd45d624a6e523445fd21ec13d3653cd51f681abf67418f54eb8" dependencies = [ "cfg-if", + "log 0.4.17", "pin-project-lite", "tracing-attributes", "tracing-core", @@ -2205,6 +2249,19 @@ dependencies = [ "log 0.3.9", ] +[[package]] +name = "vnc-rs" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aa575d4f8c5cb9073acd2ffb5dacf9e22f6abd5426065a6cf46cf80ee200b1c1" +dependencies = [ + "anyhow", + "flate2 1.0.24", + "thiserror", + "tokio", + "tracing", +] + [[package]] name = "wasi" version = "0.9.0+wasi-snapshot-preview1" diff --git a/Cargo.toml b/Cargo.toml index 6d367ed..b09b497 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -29,6 +29,7 @@ simplelog = "0.12" socks = "0.3" url = "2.1.1" vnc = "0.4" +vnc-rs = { version = "0.3", package = "vnc-rs" } [dependencies.chromiumoxide] version = "0.4" diff --git a/src/main.rs b/src/main.rs index a1230bf..c9f0599 100644 --- a/src/main.rs +++ b/src/main.rs @@ -19,8 +19,6 @@ use crate::argparse::Opts; use crate::reporting::ReportMessage; -//#[allow(unused)] -//use log::{debug, error, info, trace, warn}; use color_eyre::Result; use parsing::{generate_target_lists, InputLists}; use simplelog::{ @@ -30,8 +28,9 @@ use simplelog::{ use std::fs::{create_dir_all, File}; use std::path::Path; use std::sync::atomic::{AtomicBool, Ordering}; -use std::sync::{mpsc, Arc}; +use std::sync::Arc; use std::thread; +use tokio::sync::mpsc; use web::chrome_worker; //#[macro_use] @@ -43,8 +42,10 @@ mod rdp; mod reporting; mod util; mod vnc; +mod vnc2; mod web; +#[derive(Debug)] pub enum ThreadStatus { Complete, } @@ -152,10 +153,10 @@ async fn main() -> Result<()> { let (report_tx, report_rx): ( mpsc::Sender, mpsc::Receiver<_>, - ) = mpsc::channel(); + ) = mpsc::channel(10); let opts_clone = opts.clone(); let targets_clone = targets.clone(); - let reporting_handle = thread::spawn(move || { + let reporting_handle = tokio::task::spawn({ log::debug!("Starting report thread"); reporting::reporting_thread(report_rx, opts_clone, targets_clone) }); @@ -184,7 +185,7 @@ async fn main() -> Result<()> { let opts_clone = opts.clone(); let report_tx_clone = report_tx.clone(); let caught_ctrl_c_clone = caught_ctrl_c.clone(); - Some(thread::spawn(move || { + Some(tokio::task::spawn({ log::debug!("Starting VNC worker threads"); vnc_worker( targets_clone, @@ -192,7 +193,6 @@ async fn main() -> Result<()> { report_tx_clone, caught_ctrl_c_clone, ) - .unwrap() })) } else { None @@ -214,13 +214,13 @@ async fn main() -> Result<()> { // wait for the workers to complete if let Some(h) = rdp_handle { - h.join().unwrap().unwrap(); + h.join().unwrap()?; } if let Some(h) = vnc_handle { - h.join().unwrap(); + tokio::join!(h).0??; } - report_tx.send(ReportMessage::GenerateReport).unwrap(); - reporting_handle.join().unwrap().unwrap(); + report_tx.send(ReportMessage::GenerateReport).await.unwrap(); + tokio::join!(reporting_handle).0.unwrap().unwrap(); Ok(()) } @@ -236,10 +236,10 @@ fn rdp_worker( let mut num_workers: usize = 0; let mut targets_iter = targets.rdp_targets.iter(); let mut workers: Vec<_> = Vec::new(); - let (thread_status_tx, thread_status_rx): ( + let (thread_status_tx, mut thread_status_rx): ( Sender, Receiver, - ) = mpsc::channel(); + ) = mpsc::channel(10); while !caught_ctrl_c.load(Ordering::SeqCst) { // check for status messages // Turn off clippy's single_match warning here because match @@ -282,7 +282,7 @@ fn rdp_worker( Ok(()) } -fn vnc_worker( +async fn vnc_worker( targets: Arc, opts: Arc, report_tx: mpsc::Sender, @@ -293,16 +293,12 @@ fn vnc_worker( let mut num_workers: usize = 0; let mut targets_iter = targets.vnc_targets.iter(); let mut workers: Vec<_> = Vec::new(); - let (thread_status_tx, thread_status_rx): ( + let (thread_status_tx, mut thread_status_rx): ( Sender, Receiver, - ) = mpsc::channel(); + ) = mpsc::channel(10); while !caught_ctrl_c.load(Ordering::SeqCst) { // check for status messages - // Turn off clippy's single_match warning here because match - // matches the intuition for how try_recv is processed better - // than an if let. - #[allow(clippy::single_match)] match thread_status_rx.try_recv() { Ok(ThreadStatus::Complete) => { info!("VNC", "Thread complete, yay"); @@ -317,8 +313,9 @@ fn vnc_worker( let opts_clone = opts.clone(); let tx = thread_status_tx.clone(); let report_tx_clone = report_tx.clone(); - let handle = thread::spawn(move || { - vnc::capture(&target, &opts_clone, tx, &report_tx_clone) + let handle = tokio::task::spawn(async move { + vnc2::capture(&target, &opts_clone, tx, &report_tx_clone) + .await }); workers.push(handle); @@ -331,7 +328,7 @@ fn vnc_worker( debug!("VNC", "At the join part"); for w in workers { debug!("VNC", "Joining {:?}", w); - w.join().unwrap(); + tokio::join!(w).0.unwrap(); } Ok(()) diff --git a/src/rdp/mod.rs b/src/rdp/mod.rs index cb6397f..535a815 100644 --- a/src/rdp/mod.rs +++ b/src/rdp/mod.rs @@ -274,7 +274,7 @@ impl Write for SocketType { fn capture_worker( target: &Target, opts: &Opts, - report_tx: &mpsc::Sender, + report_tx: &tokio::sync::mpsc::Sender, ) -> Result<(), Error> { info!(target, "Connecting to {:?}", target); let addr = match target { @@ -362,7 +362,7 @@ fn capture_worker( relative_filepath.display().to_string(), ), }); - report_tx.send(report_message)?; + report_tx.blocking_send(report_message)?; } None => { warn!(target, @@ -444,8 +444,8 @@ fn bmp_thread( pub fn capture( target: &Target, opts: &Opts, - tx: mpsc::Sender, - report_tx: &mpsc::Sender, + tx: tokio::sync::mpsc::Sender, + report_tx: &tokio::sync::mpsc::Sender, ) { if let Err(e) = capture_worker(target, opts, report_tx) { warn!(target, "error: {}", e); @@ -470,9 +470,9 @@ pub fn capture( }), }; report_tx - .send(report_message) + .blocking_send(report_message) .expect("Reporting thread seems to have disconnected"); } - tx.send(ThreadStatus::Complete).unwrap(); + tx.blocking_send(ThreadStatus::Complete).unwrap(); } diff --git a/src/reporting.rs b/src/reporting.rs index a12598a..0342dbe 100644 --- a/src/reporting.rs +++ b/src/reporting.rs @@ -25,7 +25,8 @@ use askama::Template; use color_eyre::Result; use std::fs; use std::path::Path; -use std::sync::{mpsc, Arc}; +use std::sync::Arc; +use tokio::sync::mpsc; #[allow(unused)] use log::{debug, error, info, trace, warn}; @@ -74,8 +75,8 @@ pub enum FileError { Error(String), } -pub fn reporting_thread( - rx: mpsc::Receiver, +pub async fn reporting_thread( + mut rx: mpsc::Receiver, opts: Arc, targets: Arc, ) -> Result<()> { @@ -90,7 +91,7 @@ pub fn reporting_thread( let mut vnc_errors: Vec = Vec::new(); // Main loop listening on the channel - while let Ok(msg) = rx.recv() { + while let Some(msg) = rx.recv().await { use ReportMessage::*; debug!("Received message: {:?}", msg); match msg { diff --git a/src/vnc2.rs b/src/vnc2.rs new file mode 100644 index 0000000..9cf88d4 --- /dev/null +++ b/src/vnc2.rs @@ -0,0 +1,57 @@ +/* + * This file is part of NCC Group Scrying https://github.com/nccgroup/scrying + * Copyright 2020-2021 David Young + * Released as open source by NCC Group Plc - https://www.nccgroup.com + * + * Scrying is free software: you can redistribute it and/or modify + * it under the terms of the GNU General Public License as published by + * the Free Software Foundation, either version 3 of the License, or + * (at your option) any later version. + * + * Scrying is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * GNU General Public License for more details. + * + * You should have received a copy of the GNU General Public License + * along with Scrying. If not, see . +*/ + +use crate::argparse::Mode::Vnc; +use crate::argparse::Opts; +use crate::parsing::Target; +use crate::reporting::ReportMessageContent; +use crate::reporting::{FileError, ReportMessage}; +use crate::util::target_to_filename; +use crate::ThreadStatus; +#[allow(unused)] +use crate::{debug, error, info, trace, warn}; +use color_eyre::{eyre::eyre, Result}; +use image::{DynamicImage, ImageBuffer, Rgb}; +use std::cmp::min; +use std::convert::TryInto; +use std::path::Path; +use tokio::net::TcpStream; +use tokio::sync::mpsc::Sender; +use vnc_rs::{PixelFormat, Rect, VncConnector, VncEvent, X11Event}; + +async fn vnc_capture( + target: &Target, + opts: &Opts, + report_tx: &Sender, +) -> Result<()> { + todo!() +} + +pub async fn capture( + target: &Target, + opts: &Opts, + tx: Sender, + report_tx: &Sender, +) { + if let Err(e) = vnc_capture(target, opts, report_tx).await { + warn!(target, "VNC error: {}", e); + } + + tx.send(ThreadStatus::Complete).await.unwrap(); +} diff --git a/src/web/chrome.rs b/src/web/chrome.rs index 58a1b33..48294d7 100644 --- a/src/web/chrome.rs +++ b/src/web/chrome.rs @@ -10,8 +10,9 @@ use color_eyre::{eyre::eyre, Result}; use futures::StreamExt; use std::sync::{ atomic::{AtomicBool, Ordering}, - mpsc, Arc, + Arc, }; +use tokio::sync::mpsc; pub async fn chrome_worker( targets: Arc, diff --git a/src/web/mod.rs b/src/web/mod.rs index f62686e..772f90d 100644 --- a/src/web/mod.rs +++ b/src/web/mod.rs @@ -25,8 +25,8 @@ use crate::util::target_to_filename; use crate::{debug, error, info, trace, warn}; use color_eyre::Result; use std::path::Path; -use std::sync::mpsc; use std::{fs::File, io::Write}; +use tokio::sync::mpsc; pub use chrome::chrome_worker; mod chrome; @@ -51,7 +51,7 @@ pub fn save( target: target.to_string(), output: FileError::File(relative_filepath.display().to_string()), }); - report_tx.send(report_message)?; + report_tx.blocking_send(report_message)?; Ok(()) } From 34bbf7d6c13fe1334b2ccd62872903170beec079 Mon Sep 17 00:00:00 2001 From: David Young Date: Sun, 6 Nov 2022 15:52:24 +0000 Subject: [PATCH 2/2] Initialise vnc connection --- Cargo.lock | 1 + Cargo.toml | 1 + src/main.rs | 3 +-- src/vnc2.rs | 57 ++++++++++++++++++++++++++++++++++++++++++++++++++++- 4 files changed, 59 insertions(+), 3 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index b38e8ef..9e4b9fd 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1707,6 +1707,7 @@ checksum = "d29ab0c6d3fc0ee92fe66e2d99f700eab17a8d57d1c1d3b748380fb20baa78cd" name = "scrying" version = "0.9.0" dependencies = [ + "anyhow", "askama", "chromiumoxide", "clap", diff --git a/Cargo.toml b/Cargo.toml index b09b497..0d302ed 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -12,6 +12,7 @@ readme = "README.md" [features] [dependencies] +anyhow = "1" askama = "0.11" clap = { version = "3", features = ["cargo", "derive"] } color-eyre = "0.6" diff --git a/src/main.rs b/src/main.rs index c9f0599..8447971 100644 --- a/src/main.rs +++ b/src/main.rs @@ -313,9 +313,8 @@ async fn vnc_worker( let opts_clone = opts.clone(); let tx = thread_status_tx.clone(); let report_tx_clone = report_tx.clone(); - let handle = tokio::task::spawn(async move { + let handle = tokio::task::spawn({ vnc2::capture(&target, &opts_clone, tx, &report_tx_clone) - .await }); workers.push(handle); diff --git a/src/vnc2.rs b/src/vnc2.rs index 9cf88d4..93c8e02 100644 --- a/src/vnc2.rs +++ b/src/vnc2.rs @@ -17,6 +17,7 @@ * along with Scrying. If not, see . */ +use self::vnc::{PixelFormat, Rect, VncConnector, VncEvent, X11Event}; use crate::argparse::Mode::Vnc; use crate::argparse::Opts; use crate::parsing::Target; @@ -33,13 +34,55 @@ use std::convert::TryInto; use std::path::Path; use tokio::net::TcpStream; use tokio::sync::mpsc::Sender; -use vnc_rs::{PixelFormat, Rect, VncConnector, VncEvent, X11Event}; +use vnc_rs as vnc; async fn vnc_capture( target: &Target, opts: &Opts, report_tx: &Sender, ) -> Result<()> { + info!(target, "Connecting to {:?}", target); + let addr = match target { + Target::Address(sock_addr) => sock_addr, + Target::Url(_) => { + return Err(eyre!("Invalid VNC target: {target}",)); + } + }; + + async fn auth() -> anyhow::Result { + Ok(String::new()) + } + + let tcp = TcpStream::connect(addr).await?; + let vnc = VncConnector::new(tcp) + .set_auth_method(auth()) + .add_encoding(vnc::VncEncoding::Tight) + .add_encoding(vnc::VncEncoding::Zrle) + .add_encoding(vnc::VncEncoding::CopyRect) + .add_encoding(vnc::VncEncoding::Raw) + .allow_shared(true) + .set_pixel_format(PixelFormat::bgra()) + .build() + .unwrap() + .try_start() + .await + .unwrap() + .finish() + .unwrap(); + let (vnc_event_sender, mut vnc_event_receiver) = + tokio::sync::mpsc::channel(100); + let (x11_event_sender, x11_event_receiver) = + tokio::sync::mpsc::channel(100); + tokio::spawn(async move { + vnc.run(vnc_event_sender, x11_event_receiver).await.unwrap() + }); + let _ = x11_event_sender.send(X11Event::Refresh).await; + + let mut image = VncImage::new(); + while let Some(event) = vnc_event_receiver.recv().await { + image.handle_event(event)?; + } + todo!() } @@ -55,3 +98,15 @@ pub async fn capture( tx.send(ThreadStatus::Complete).await.unwrap(); } + +struct VncImage {} + +impl VncImage { + pub fn new() -> Self { + Self {} + } + + pub fn handle_event(&mut self, event: VncEvent) -> Result<()> { + Ok(()) + } +}