From 2e577d802b99af2b50a9a5a7af42b28d9e85948f Mon Sep 17 00:00:00 2001
From: b1-team <120295547+b1nhack@users.noreply.github.com>
Date: Mon, 19 Jun 2023 19:56:17 -0700
Subject: [PATCH] Optimized code style

---
 Cargo.toml    |  2 +-
 src/driver.rs | 29 +++++++++++++++++++++--------
 src/main.rs   | 11 +++++------
 3 files changed, 27 insertions(+), 15 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index a4899ee..9f4b19f 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "superman"
-version = "1.1.5"
+version = "1.1.6"
 edition = "2021"
 authors = ["b1n <b1nhack@proton.me>"]
 description = "Kill everything."
diff --git a/src/driver.rs b/src/driver.rs
index b66d3ce..197c2f6 100644
--- a/src/driver.rs
+++ b/src/driver.rs
@@ -25,6 +25,7 @@ use windows_sys::Win32::System::Services::{
 };
 use windows_sys::Win32::System::IO::DeviceIoControl;
 
+/// Entry structure, representing a driver and its operations
 pub struct Driver {
     path: PathBuf,
     service_name: CString,
@@ -41,6 +42,21 @@ impl Driver {
     pub fn service_name(&self) -> &CStr {
         &self.service_name
     }
+
+    /// Load and start driver
+    pub fn load_driver(&self) -> anyhow::Result<()> {
+        load_driver(self)
+    }
+
+    /// Unload and delete driver
+    pub fn unload_driver(&self) -> anyhow::Result<()> {
+        unload_driver(self)
+    }
+
+    /// Send ioctl to kill pid
+    pub fn kill_pid(&self, args: &Args, rx: Receiver<bool>) -> anyhow::Result<()> {
+        kill_pid(self, args, rx)
+    }
 }
 
 /// Make sure driver status
@@ -84,15 +100,14 @@ fn check_service_status(driver: &Driver) -> anyhow::Result<bool> {
                 Ok(true)
             }
             _ => {
-                unload_driver(driver)?;
+                driver.unload_driver()?;
                 Ok(false)
             }
         }
     }
 }
 
-/// Load and start driver
-pub fn load_driver(driver: &Driver) -> anyhow::Result<()> {
+fn load_driver(driver: &Driver) -> anyhow::Result<()> {
     if check_service_status(driver)? {
         return Ok(());
     }
@@ -144,8 +159,7 @@ pub fn load_driver(driver: &Driver) -> anyhow::Result<()> {
     Ok(())
 }
 
-/// Unload and delete driver
-pub fn unload_driver(driver: &Driver) -> anyhow::Result<()> {
+fn unload_driver(driver: &Driver) -> anyhow::Result<()> {
     let mut status: SERVICE_STATUS = unsafe { zeroed() };
 
     unsafe {
@@ -188,8 +202,7 @@ pub fn unload_driver(driver: &Driver) -> anyhow::Result<()> {
     Ok(())
 }
 
-/// Send ioctl to kill pid
-pub fn kill_pid(args: &Args, driver: &Driver, rx: Receiver<bool>) -> anyhow::Result<()> {
+fn kill_pid(driver: &Driver, args: &Args, rx: Receiver<bool>) -> anyhow::Result<()> {
     let initialize_ioctl_code: u32 = 0x9876C004u32;
     let terminate_process_ioctl_code: u32 = 0x9876C094u32;
     let device_name = CStr::from_bytes_with_nul(b"\\\\.\\superman\0")?;
@@ -246,7 +259,7 @@ pub fn kill_pid(args: &Args, driver: &Driver, rx: Receiver<bool>) -> anyhow::Res
                 // exit
                 if rx.try_recv().is_ok() {
                     CloseHandle(device);
-                    unload_driver(driver)?;
+                    driver.unload_driver()?;
                     process::exit(0i32);
                 }
 
diff --git a/src/main.rs b/src/main.rs
index 17fbbe7..5f5ac6a 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -4,7 +4,7 @@ mod greet;
 mod utils;
 
 use crate::args::Args;
-use crate::driver::{kill_pid, load_driver, unload_driver, Driver};
+use crate::driver::Driver;
 use crate::utils::check_pid;
 use anyhow::anyhow;
 use clap::Parser;
@@ -56,19 +56,18 @@ fn main() {
     if let Err(e) = try_main(&args, &driver) {
         eprintln!("{}", e);
     }
-    let _ = unload_driver(&driver);
+    let _ = driver.unload_driver();
 }
 
 fn try_main(args: &Args, driver: &Driver) -> anyhow::Result<()> {
-    let (sx, rx) = mpsc::sync_channel(1);
+    let (sx, rx) = mpsc::sync_channel(0);
     init_ctrlc(sx)?;
 
     if check_pid(args.pid).not() {
         return Err(anyhow!("[-]Process not exists!"));
     }
 
-    load_driver(driver)?;
-
-    kill_pid(args, driver, rx)?;
+    driver.load_driver()?;
+    driver.kill_pid(args, rx)?;
     Ok(())
 }