From dfdc1db11a5712ff946450ea3022a05909a52882 Mon Sep 17 00:00:00 2001 From: William Desportes Date: Sat, 22 Jun 2024 20:32:59 +0200 Subject: [PATCH] Extract the code into functions --- snow-scanner/src/main.rs | 189 ++++++++++++++++++++++----------------- 1 file changed, 105 insertions(+), 84 deletions(-) diff --git a/snow-scanner/src/main.rs b/snow-scanner/src/main.rs index 4d9493b..69494b5 100644 --- a/snow-scanner/src/main.rs +++ b/snow-scanner/src/main.rs @@ -3,13 +3,11 @@ extern crate rouille; use chrono::Utc; use hmac::{Hmac, Mac}; -use rouille::Response; -use rouille::ResponseBody; +use rouille::{Request, Response, ResponseBody}; use rusqlite::types::ToSqlOutput; use rusqlite::{named_params, Connection, OpenFlags, Result, ToSql}; use sha2::Sha256; use std::fmt; -use std::io; use std::str::FromStr; use hickory_client::client::SyncClient; @@ -87,7 +85,7 @@ fn detect_scanner(ptr_result: &ResolvedResult) -> Result { .eq_case(&Name::from_str("binaryedge.ninja.").expect("Should parse")) => { Ok(Scanners::Binaryedge) - }, + } Some(ref x) if x.trim_to(2) .eq_case(&Name::from_str("stretchoid.com").expect("Should parse")) => @@ -98,7 +96,6 @@ fn detect_scanner(ptr_result: &ResolvedResult) -> Result { } } -// The HTML document of the home page. static FORM: &str = r#" @@ -113,10 +110,110 @@ static FORM: &str = r#"

+
+

+

+
"#; +fn handle_report( + conn: &Connection, + client: SyncClient, + request: &Request, +) -> Response { + let data = try_or_400!(post_input!(request, { + ip: String, + })); + + // We just print what was received on stdout. Of course in a real application + // you probably want to process the data, eg. store it in a database. + println!("Received data: {:?}", data); + let query_address = data.ip.parse().expect("To parse"); + + let ptr_result = get_ptr(query_address, client).unwrap(); + + match detect_scanner(&ptr_result) { + Ok(scanner_name) => { + let ip_type = if data.ip.contains(':') { 6 } else { 4 }; + let scanner = Scanner { + ip: data.ip, + ip_type: ip_type, + scanner_name: scanner_name.clone(), + created_at: Utc::now().to_string(), + updated_at: Utc::now().to_string(), + last_seen_at: Utc::now().to_string(), + last_checked_at: Utc::now().to_string(), + }; + save_scanner(&conn, &scanner).unwrap(); + rouille::Response::html(match scanner_name { + Scanners::Binaryedge => format!( + "Reported an escaped ninja! {} {:?}.", + scanner.ip, + ptr_result.result.unwrap() + ), + Scanners::Strechoid => format!( + "Reported a stretchoid agent! {} {:?}.", + scanner.ip, + ptr_result.result.unwrap() + ), + }) + } + + Err(_) => rouille::Response::html(format!( + "The IP {} resolved as {:?} did not match known scanners patterns.", + data.ip, ptr_result.result + )), + } +} + +fn handle_list_scanners(conn: &Connection, scanner_name: String) -> Response { + let mut stmt = conn.prepare("SELECT ip FROM scanners WHERE scanner_name = :scanner_name ORDER BY ip_type, created_at").unwrap(); + let mut rows = stmt + .query(named_params! { ":scanner_name": scanner_name }) + .unwrap(); + let mut ips: Vec = vec![]; + while let Some(row) = rows.next().unwrap() { + ips.push(row.get(0).unwrap()); + } + + Response { + status_code: 200, + headers: vec![("Content-Type".into(), "text/plain; charset=utf-8".into())], + data: ResponseBody::from_string(ips.join("\n")), + upgrade: None, + } +} + +fn get_connection() -> Connection { + let path = "./snow-scanner.sqlite"; + let conn = Connection::open_with_flags( + path, + OpenFlags::SQLITE_OPEN_READ_WRITE + | OpenFlags::SQLITE_OPEN_CREATE + | OpenFlags::SQLITE_OPEN_FULL_MUTEX, + ) + .unwrap(); + conn.execute( + "CREATE TABLE IF NOT EXISTS scanners ( + ip VARCHAR(255), + ip_type TINYINT(1), + scanner_name VARCHAR(255), + created_at DATETIME, + updated_at DATETIME, + last_seen_at DATETIME, + last_checked_at DATETIME, + PRIMARY KEY (ip, ip_type) + )", + (), // empty list of parameters. + ) + .unwrap(); + conn.pragma_update_and_check(None, "journal_mode", &"WAL", |_| Ok(())) + .unwrap(); + conn +} + fn main() -> Result<()> { println!("Now listening on localhost:8000"); @@ -125,30 +222,7 @@ fn main() -> Result<()> { rouille::start_server("localhost:8000", move |request| { let client = SyncClient::new(conn); - let path = "./snow-scanner.sqlite"; - let conn = Connection::open_with_flags( - path, - OpenFlags::SQLITE_OPEN_READ_WRITE - | OpenFlags::SQLITE_OPEN_CREATE - | OpenFlags::SQLITE_OPEN_FULL_MUTEX, - ) - .unwrap(); - conn.execute( - "CREATE TABLE IF NOT EXISTS scanners ( - ip VARCHAR(255), - ip_type TINYINT(1), - scanner_name VARCHAR(255), - created_at DATETIME, - updated_at DATETIME, - last_seen_at DATETIME, - last_checked_at DATETIME, - PRIMARY KEY (ip, ip_type) - )", - (), // empty list of parameters. - ) - .unwrap(); - conn.pragma_update_and_check(None, "journal_mode", &"WAL", |_| Ok(())) - .unwrap(); + let conn = get_connection(); router!(request, (GET) (/) => { @@ -159,48 +233,7 @@ fn main() -> Result<()> { rouille::Response::text("pong") }, - (POST) (/report) => { - let data = try_or_400!(post_input!(request, { - ip: String, - })); - - // We just print what was received on stdout. Of course in a real application - // you probably want to process the data, eg. store it in a database. - println!("Received data: {:?}", data); - let query_address = data.ip.parse().expect("To parse"); - - let ptr_result = get_ptr(query_address, client).unwrap(); - - match detect_scanner(&ptr_result) { - Ok(scanner_name) => { - let ip_type = if data.ip.contains(':') {6} else {4}; - let scanner = Scanner { - ip: data.ip, - ip_type: ip_type, - scanner_name: scanner_name.clone(), - created_at: Utc::now().to_string(), - updated_at: Utc::now().to_string(), - last_seen_at: Utc::now().to_string(), - last_checked_at: Utc::now().to_string(), - }; - save_scanner(&conn, &scanner).unwrap(); - rouille::Response::html( - match scanner_name { - Scanners::Binaryedge => - format!("Reported an escaped ninja! {} {:?}.", scanner.ip, ptr_result.result.unwrap()), - Scanners::Strechoid => - format!("Reported a stretchoid agent! {} {:?}.", scanner.ip, ptr_result.result.unwrap()) - } - ) - - }, - - Err(_) => - rouille::Response::html( - format!("The IP {} resolved as {:?} did not match known scanners patterns.", data.ip, ptr_result.result) - ) - } - }, + (POST) (/report) => {handle_report(&conn, client, &request)}, (POST) (/register) => { let data = try_or_400!(post_input!(request, { @@ -227,19 +260,7 @@ fn main() -> Result<()> { }, (GET) (/scanners/{scanner_name: String}) => { - let mut stmt = conn.prepare("SELECT ip FROM scanners WHERE scanner_name = :scanner_name ORDER BY ip_type, created_at").unwrap(); - let mut rows = stmt.query(named_params! { ":scanner_name": scanner_name }).unwrap(); - let mut ips: Vec = vec!(); - while let Some(row) = rows.next().unwrap() { - ips.push(row.get(0).unwrap()); - } - - Response { - status_code: 200, - headers: vec![("Content-Type".into(), "text/plain; charset=utf-8".into())], - data: ResponseBody::from_string(ips.join("\n")), - upgrade: None, - } + handle_list_scanners(&conn, scanner_name) }, (GET) (/{api_key: String}/scanners/{scanner_name: String}) => { let mut mac = HmacSha256::new_from_slice(b"my secret and secure key")