diff options
| -rw-r--r-- | Cargo.toml | 1 | ||||
| -rw-r--r-- | src/main.rs | 10 |
2 files changed, 10 insertions, 1 deletions
@@ -7,5 +7,6 @@ edition = "2018" [dependencies] chrono = { version = "0.4.6", features = ["serde"] } clap = "2.33" +ctrlc = { version = "3.1.2", features = ["termination"] } serde = { version = "1.0.91", features = ["derive"] } serde_json = "1.0.39" diff --git a/src/main.rs b/src/main.rs index db6bb73..36c5cf9 100644 --- a/src/main.rs +++ b/src/main.rs @@ -21,6 +21,7 @@ use std::error::Error; use std::fs::{ File, remove_file }; use std::io::{ BufRead, BufReader, Write }; use std::os::unix::net::{ UnixListener, UnixStream }; +use std::process::exit; use std::result::Result; use std::string::ToString; use std::sync::Arc; @@ -28,6 +29,13 @@ use std::thread; use std::time::SystemTime; fn main() { + // Set signal handler to clean up on termination. + ctrlc::set_handler(move || { + // We don't care about errors in remove_file. + remove_file("/tmp/phisher").unwrap_or(()); + exit(0); + }).expect("Couldn't set signal handler!"); + let args = clap_app!(app => (name: crate_name!()) (version: crate_version!()) @@ -78,7 +86,7 @@ fn run_server(tank: Arc<phisher::PhishTank>) -> Result<(), Box<dyn Error>> { } } - Ok(remove_file("/tmp/phisher")?) + Ok(()) } // Handle a connected client. |