Create a unix socket server to check URLs.

Set up a simple unix domain socket server to handle requests to check
if URLs are phishes or not.

1 files changed, 56 insertions, 3 deletions

diff --git a/src/main.rs b/src/main.rs
index 17f443b..d2198e1 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -16,8 +16,14 @@
 
 use phisher;
 use clap::{clap_app, crate_name, crate_authors, crate_version, crate_description};
-use std::fs::File;
-use std::io::BufReader;
+use std::boxed::Box;
+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::result::Result;
+use std::sync::Arc;
+use std::thread;
 use std::time::SystemTime;
 
 fn main() {
@@ -32,13 +38,60 @@ fn main() {
     let start_time = SystemTime::now();
     let filename = args.value_of("PHISHTANK").unwrap();
     if let Ok(file) = File::open(filename) {
-        let tank = phisher::load_phistank(BufReader::new(file)).unwrap();
+        let tank = Arc::new(phisher::load_phistank(BufReader::new(file)).unwrap());
 
         println!("Loaded {} phishes in {} seconds!",
             tank.phishes.len(),
             start_time.elapsed().unwrap().as_secs());
+
+        match run_server(tank) {
+            Ok(_) => println!("Exiting."),
+            Err(e) => eprintln!("Received error: {}, terminating.", e.to_string())
+        };
     }
     else {
         eprintln!("Could not open file {}.\nIs the filename correct?", filename);
     }
 }
+
+fn run_server(tank: Arc<phisher::PhishTank>) -> Result<(), Box<dyn Error>> {
+    let listener = UnixListener::bind("/tmp/phisher")?;
+    for stream in listener.incoming() {
+        match stream {
+            Ok(stream) => {
+                let t = Arc::clone(&tank);
+                thread::spawn(move || handle_connection(stream, t));
+            }
+            Err(e) => {
+                eprintln!("Received error: {}", e.to_string());
+                break;
+            }
+        }
+    }
+
+    Ok(remove_file("/tmp/phisher")?)
+}
+
+// Handle a connected client.
+//
+// Once a client is connected, it can request the status of multiple
+// URLs before closing the connection. Each URL must be passed in on
+// a separate line, and the results are returned in the same order
+// as the URLs were passed in.
+//
+fn handle_connection(mut stream: UnixStream, tank: Arc<phisher::PhishTank>) {
+    // Clone the stream so we get another ref to the underlying socket,
+    // then wrap it in a BufReader, so we can use read_line on it.
+    let mut bufreader = BufReader::new(stream.try_clone().unwrap());
+
+    loop {
+        let mut buffer = String::new();
+        bufreader.read_line(&mut buffer).unwrap();
+        if tank.is_phish(buffer.trim_end()) {
+            stream.write(b"it's a phish!\n").unwrap();
+        }
+        else {
+            stream.write(b"good url.\n").unwrap();
+        }
+    }
+}