diff --git a/.vscode/settings.json b/.vscode/settings.json new file mode 100644 index 0000000..ebfde9a --- /dev/null +++ b/.vscode/settings.json @@ -0,0 +1,3 @@ +{ + "editor.formatOnSave": true, +} \ No newline at end of file diff --git a/src/main.rs b/src/main.rs index 0f8997b..61be97b 100644 --- a/src/main.rs +++ b/src/main.rs @@ -3,7 +3,9 @@ mod server; use log::info; fn main() { - pretty_env_logger::formatted_builder().filter_level(log::LevelFilter::Debug).init(); + pretty_env_logger::formatted_builder() + .filter_level(log::LevelFilter::Debug) + .init(); info!("Hello, world!"); let rt = tokio::runtime::Runtime::new().unwrap(); @@ -11,4 +13,3 @@ fn main() { server::run::().await.unwrap(); }) } - diff --git a/src/server/ce_common.rs b/src/server/ce_common.rs index 46f192c..3bfa491 100644 --- a/src/server/ce_common.rs +++ b/src/server/ce_common.rs @@ -36,7 +36,10 @@ pub fn write_i32_prefixed_string(buf: &mut dyn BufMut, value: String) { } pub fn cstring_to_string(with_zeroes: &[u8]) -> String { - let len = with_zeroes.iter().position(|i| *i == 0).unwrap_or(with_zeroes.len()); + let len = with_zeroes + .iter() + .position(|i| *i == 0) + .unwrap_or(with_zeroes.len()); let without_zeroes = &with_zeroes[0..len]; String::from_utf8(without_zeroes.to_vec()).unwrap() } diff --git a/src/server/commands_request.rs b/src/server/commands_request.rs index f7359a8..ce1bdd0 100644 --- a/src/server/commands_request.rs +++ b/src/server/commands_request.rs @@ -2,9 +2,9 @@ use std::fmt::Debug; use bytes::Buf; -use super::{command::*, commands_response::*, ce_common::*}; +use super::{ce_common::*, command::*, commands_response::*}; -pub trait CERequest : Debug + Send { +pub trait CERequest: Debug + Send { type Response: CEResponse + Debug + Send; const ID: Command; diff --git a/src/server/commands_response.rs b/src/server/commands_response.rs index 0831e22..e4cb9f8 100644 --- a/src/server/commands_response.rs +++ b/src/server/commands_response.rs @@ -28,8 +28,7 @@ impl CEResponse for Process32Response { writer.put_u32_le(1u32); writer.put_i32_le(entry.pid); write_i32_prefixed_string(writer, entry.process_name) - } - else { + } else { writer.put_u32_le(0u32); writer.put_i32_le(0); // pid writer.put_i32_le(0); // process name length @@ -49,8 +48,7 @@ impl CEResponse for Module32Response { writer.put_i64_le(entry.module_base); writer.put_i32_le(entry.module_size); write_i32_prefixed_string(writer, entry.module_name) - } - else { + } else { writer.put_u32_le(0u32); writer.put_i64_le(0); // base writer.put_i32_le(0); // size diff --git a/src/server/handler.rs b/src/server/handler.rs index f7650d0..8563a64 100644 --- a/src/server/handler.rs +++ b/src/server/handler.rs @@ -1,15 +1,15 @@ use super::commands_request::*; -pub trait HandlerFactory { -} +pub trait HandlerFactory {} -pub trait FullHandler - : Handler +pub trait FullHandler: + Handler + Handler + Handler + Handler + Handler - + Handler { + + Handler +{ fn create() -> Self; } diff --git a/src/server/handlers/mod.rs b/src/server/handlers/mod.rs index 540a574..4bc8fb8 100644 --- a/src/server/handlers/mod.rs +++ b/src/server/handlers/mod.rs @@ -1,4 +1,4 @@ mod windows; #[cfg(windows)] -pub use windows::WindowsHandler; \ No newline at end of file +pub use windows::WindowsHandler; diff --git a/src/server/handlers/windows.rs b/src/server/handlers/windows.rs index ab2821f..519ea32 100644 --- a/src/server/handlers/windows.rs +++ b/src/server/handlers/windows.rs @@ -2,17 +2,18 @@ use crate::server::{ce_common::*, commands_request::*, commands_response::*, han use winapi::um::{ handleapi::CloseHandle, tlhelp32::{ - CreateToolhelp32Snapshot, - Process32First, Process32Next, Module32First, Module32Next, - LPPROCESSENTRY32, LPMODULEENTRY32 + CreateToolhelp32Snapshot, Module32First, Module32Next, Process32First, Process32Next, + LPMODULEENTRY32, LPPROCESSENTRY32, }, - winnt::HANDLE + winnt::HANDLE, }; pub struct WindowsHandler; impl FullHandler for WindowsHandler { - fn create() -> WindowsHandler { WindowsHandler } + fn create() -> WindowsHandler { + WindowsHandler + } } impl Handler for WindowsHandler { @@ -21,7 +22,7 @@ impl Handler for WindowsHandler { let ret = CreateToolhelp32Snapshot(req.snapshot_flags, req.process_id); HandleResponse { - handle: ret as usize + handle: ret as usize, } } } @@ -29,76 +30,64 @@ impl Handler for WindowsHandler { impl Handler for WindowsHandler { fn handle(&self, req: Process32FirstRequest) -> Process32Response { - unsafe { - get_process_response(req.handle, |x, y| Process32First(x, y)) - } + unsafe { get_process_response(req.handle, |x, y| Process32First(x, y)) } } } impl Handler for WindowsHandler { fn handle(&self, req: Process32NextRequest) -> Process32Response { - unsafe { - get_process_response(req.handle, |x, y| Process32Next(x, y)) - } + unsafe { get_process_response(req.handle, |x, y| Process32Next(x, y)) } } } unsafe fn get_process_response(handle: usize, func: F) -> Process32Response - where F : FnOnce(HANDLE, LPPROCESSENTRY32) -> i32 { +where + F: FnOnce(HANDLE, LPPROCESSENTRY32) -> i32, +{ let mut entry = std::mem::MaybeUninit::uninit().assume_init(); let response = func(std::mem::transmute(handle), &mut entry); if response != 0 { Process32Response { - entry: Some( - CeProcessEntry { - pid: entry.th32ProcessID as i32, - process_name: cstring_to_string(std::mem::transmute(&entry.szExeFile[..])), - } - ) + entry: Some(CeProcessEntry { + pid: entry.th32ProcessID as i32, + process_name: cstring_to_string(std::mem::transmute(&entry.szExeFile[..])), + }), } } else { - Process32Response { - entry: None - } + Process32Response { entry: None } } } impl Handler for WindowsHandler { fn handle(&self, req: Module32FirstRequest) -> Module32Response { - unsafe { - get_module_response(req.handle, |x, y| Module32First(x, y)) - } + unsafe { get_module_response(req.handle, |x, y| Module32First(x, y)) } } } impl Handler for WindowsHandler { fn handle(&self, req: Module32NextRequest) -> Module32Response { - unsafe { - get_module_response(req.handle, |x, y| Module32Next(x, y)) - } + unsafe { get_module_response(req.handle, |x, y| Module32Next(x, y)) } } } unsafe fn get_module_response(handle: usize, func: F) -> Module32Response - where F : FnOnce(HANDLE, LPMODULEENTRY32) -> i32 { +where + F: FnOnce(HANDLE, LPMODULEENTRY32) -> i32, +{ let mut entry = std::mem::MaybeUninit::uninit().assume_init(); let response = func(std::mem::transmute(handle), &mut entry); if response != 0 { Module32Response { - entry: Some( - CeModuleEntry { - module_base: entry.modBaseAddr as i64, - module_size: entry.modBaseSize as i32, - module_name: cstring_to_string(std::mem::transmute(&entry.szModule[..])) - } - ) + entry: Some(CeModuleEntry { + module_base: entry.modBaseAddr as i64, + module_size: entry.modBaseSize as i32, + module_name: cstring_to_string(std::mem::transmute(&entry.szModule[..])), + }), } } else { - Module32Response { - entry: None - } + Module32Response { entry: None } } } @@ -107,9 +96,7 @@ impl Handler for WindowsHandler { unsafe { let response = CloseHandle(std::mem::transmute(req.handle)); - I32Response { - response - } + I32Response { response } } } } diff --git a/src/server/mod.rs b/src/server/mod.rs index 8fb41e9..2463d30 100644 --- a/src/server/mod.rs +++ b/src/server/mod.rs @@ -1,10 +1,10 @@ -mod server; +mod ce_common; mod command; mod commands_request; mod commands_response; mod handler; mod handlers; -mod ce_common; +mod server; +pub use handlers::WindowsHandler; pub use server::run; -pub use handlers::WindowsHandler; \ No newline at end of file diff --git a/src/server/server.rs b/src/server/server.rs index c36e65a..f2db9f9 100644 --- a/src/server/server.rs +++ b/src/server/server.rs @@ -1,9 +1,9 @@ use bytes::Buf; -use log::{info, debug, error}; +use log::{debug, error, info}; use tokio::net::{TcpListener, TcpStream}; use tokio::prelude::*; -use super::{handler::*, commands_request::*, commands_response::*}; +use super::{commands_request::*, commands_response::*, handler::*}; const DEFAULT_PORT: u16 = 52736; @@ -42,7 +42,7 @@ macro_rules! gen_request_dispatch { let mut buffer = Vec::new(); let request = <$request>::read(&mut $bytes); - + debug!("Received item {:?}", request); let response = $handler.handle(request); debug!("... responding with {:?}", response);