diff --git a/snow-scanner/src/main.rs b/snow-scanner/src/main.rs index 18182c1..4a5b6fb 100644 --- a/snow-scanner/src/main.rs +++ b/snow-scanner/src/main.rs @@ -19,7 +19,7 @@ use hickory_client::client::SyncClient; use hickory_client::rr::Name; use hickory_client::tcp::TcpClientConnection; -use dns_ptr_resolver::{get_ptr, ResolvedResult}; +use dns_ptr_resolver::{get_ptr, ResolvedResult, ResolvingError}; // Create alias for HMAC-SHA256 type HmacSha256 = Hmac; @@ -148,11 +148,15 @@ fn detect_scanner(ptr_result: &ResolvedResult) -> Result { } } -fn handle_ip2(conn: &Connection, ip: String) -> Result { +fn handle_ip2(conn: &Connection, ip: String) -> Result> { let query_address = ip.parse().expect(format!("To parse: {}", ip).as_str()); let client = get_dns_client(); - let ptr_result = get_ptr(query_address, client).unwrap(); + let ptr_result: ResolvedResult = if let Ok(res) = get_ptr(query_address, client) { + res + } else { + return Err(None); + }; match detect_scanner(&ptr_result) { Ok(scanner_name) => { @@ -175,15 +179,19 @@ fn handle_ip2(conn: &Connection, ip: String) -> Result Ok(scanner) } - Err(_) => Err(ptr_result), + Err(_) => Err(Some(ptr_result)), } } -fn handle_ip(conn: &Mutex, ip: String) -> Result { +fn handle_ip(conn: &Mutex, ip: String) -> Result> { let query_address = ip.parse().expect("To parse"); let client = get_dns_client(); - let ptr_result = get_ptr(query_address, client).unwrap(); + let ptr_result: ResolvedResult = if let Ok(res) = get_ptr(query_address, client) { + res + } else { + return Err(None); + }; match detect_scanner(&ptr_result) { Ok(scanner_name) => { @@ -206,7 +214,7 @@ fn handle_ip(conn: &Mutex, ip: String) -> Result Err(ptr_result), + Err(_) => Err(Some(ptr_result)), } } @@ -269,13 +277,23 @@ fn handle_report(conn: &Mutex, request: &Request) -> Response { match handle_ip(conn, data.ip.clone()) { Ok(scanner) => rouille::Response::html(match scanner.scanner_name { - Scanners::Binaryedge => format!("Reported an escaped ninja! {}.", scanner.ip), - Scanners::Strechoid => format!("Reported a stretchoid agent! {}.", scanner.ip), + Scanners::Binaryedge => format!( + "Reported an escaped ninja! {} known as {:?}.", + scanner.ip, scanner.ip_ptr + ), + Scanners::Strechoid => format!( + "Reported a stretchoid agent! {} known as {:?}.", + scanner.ip, scanner.ip_ptr + ), }), Err(ptr_result) => rouille::Response::html(format!( "The IP {} resolved as {:?} did not match known scanners patterns.", - data.ip, ptr_result.result + data.ip, + match ptr_result { + Some(res) => res.result, + None => None, + } )), } }