diff --git a/snow-scanner/src/main.rs b/snow-scanner/src/main.rs index 39193a9..760e0a9 100644 --- a/snow-scanner/src/main.rs +++ b/snow-scanner/src/main.rs @@ -1,7 +1,7 @@ #[macro_use] extern crate rouille; -use chrono::{Utc, NaiveDateTime}; +use chrono::{NaiveDateTime, Utc}; use cidr::IpCidr; use hmac::{Hmac, Mac}; use rouille::{Request, Response, ResponseBody}; @@ -30,6 +30,17 @@ enum Scanners { Binaryedge, } +impl FromStr for Scanners { + type Err = (); + fn from_str(input: &str) -> Result { + match input { + "stretchoid" => Ok(Scanners::Stretchoid), + "binaryedge" => Ok(Scanners::Binaryedge), + _ => Err(()), + } + } +} + impl fmt::Display for Scanners { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { write!( @@ -301,7 +312,7 @@ fn handle_report(conn: &Mutex, request: &Request) -> Response { } } -fn handle_list_scanners(conn: &Mutex, scanner_name: String) -> Response { +fn handle_list_scanners(conn: &Mutex, scanner_name: Scanners) -> Response { let db = conn.lock().unwrap(); let mut stmt = db.prepare("SELECT ip FROM scanners WHERE scanner_name = :scanner_name ORDER BY ip_type, created_at").unwrap(); let mut rows = stmt @@ -373,7 +384,8 @@ fn handle_list_scan_tasks(conn: &Mutex) -> Response { {:#?} {:#?} - ", started_at, still_processing_at, ended_at + ", + started_at, still_processing_at, ended_at )); } @@ -553,7 +565,7 @@ fn main() -> Result<()> { rouille::Response::html(format!("Success! {}.", hex::encode(code_bytes))) }, - (GET) (/scanners/{scanner_name: String}) => { + (GET) (/scanners/{scanner_name: Scanners}) => { handle_list_scanners(&conn, scanner_name) }, (GET) (/{api_key: String}/scanners/{scanner_name: String}) => {