Extract the code into functions

This commit is contained in:
2024-06-22 20:32:59 +02:00
parent 642ae6e714
commit dfdc1db11a

View File

@ -3,13 +3,11 @@ extern crate rouille;
use chrono::Utc; use chrono::Utc;
use hmac::{Hmac, Mac}; use hmac::{Hmac, Mac};
use rouille::Response; use rouille::{Request, Response, ResponseBody};
use rouille::ResponseBody;
use rusqlite::types::ToSqlOutput; use rusqlite::types::ToSqlOutput;
use rusqlite::{named_params, Connection, OpenFlags, Result, ToSql}; use rusqlite::{named_params, Connection, OpenFlags, Result, ToSql};
use sha2::Sha256; use sha2::Sha256;
use std::fmt; use std::fmt;
use std::io;
use std::str::FromStr; use std::str::FromStr;
use hickory_client::client::SyncClient; use hickory_client::client::SyncClient;
@ -87,7 +85,7 @@ fn detect_scanner(ptr_result: &ResolvedResult) -> Result<Scanners, ()> {
.eq_case(&Name::from_str("binaryedge.ninja.").expect("Should parse")) => .eq_case(&Name::from_str("binaryedge.ninja.").expect("Should parse")) =>
{ {
Ok(Scanners::Binaryedge) Ok(Scanners::Binaryedge)
}, }
Some(ref x) Some(ref x)
if x.trim_to(2) if x.trim_to(2)
.eq_case(&Name::from_str("stretchoid.com").expect("Should parse")) => .eq_case(&Name::from_str("stretchoid.com").expect("Should parse")) =>
@ -98,7 +96,6 @@ fn detect_scanner(ptr_result: &ResolvedResult) -> Result<Scanners, ()> {
} }
} }
// The HTML document of the home page.
static FORM: &str = r#" static FORM: &str = r#"
<html> <html>
<head> <head>
@ -113,53 +110,19 @@ static FORM: &str = r#"
<p><input type="ip" name="ip" placeholder="An IPv4 or IPv6" /></p> <p><input type="ip" name="ip" placeholder="An IPv4 or IPv6" /></p>
<p><button>Report this IP</button></p> <p><button>Report this IP</button></p>
</form> </form>
<form action="/scan" method="POST">
<p><textarea name="ips"></textarea></p>
<p><button>Scan</button></p>
</form>
</body> </body>
</html> </html>
"#; "#;
fn main() -> Result<()> { fn handle_report(
println!("Now listening on localhost:8000"); conn: &Connection,
client: SyncClient<TcpClientConnection>,
let server = "1.1.1.1:53".parse().expect("To parse"); request: &Request,
let conn = TcpClientConnection::with_timeout(server, std::time::Duration::new(5, 0)).unwrap(); ) -> Response {
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();
router!(request,
(GET) (/) => {
rouille::Response::html(FORM)
},
(GET) (/ping) => {
rouille::Response::text("pong")
},
(POST) (/report) => {
let data = try_or_400!(post_input!(request, { let data = try_or_400!(post_input!(request, {
ip: String, ip: String,
})); }));
@ -184,24 +147,94 @@ fn main() -> Result<()> {
last_checked_at: Utc::now().to_string(), last_checked_at: Utc::now().to_string(),
}; };
save_scanner(&conn, &scanner).unwrap(); save_scanner(&conn, &scanner).unwrap();
rouille::Response::html( rouille::Response::html(match scanner_name {
match scanner_name { Scanners::Binaryedge => format!(
Scanners::Binaryedge => "Reported an escaped ninja! <b>{}</a> {:?}.",
format!("Reported an escaped ninja! <b>{}</a> {:?}.", scanner.ip, ptr_result.result.unwrap()), scanner.ip,
Scanners::Strechoid => ptr_result.result.unwrap()
format!("Reported a stretchoid agent! <b>{}</a> {:?}.", scanner.ip, ptr_result.result.unwrap()) ),
Scanners::Strechoid => format!(
"Reported a stretchoid agent! <b>{}</a> {:?}.",
scanner.ip,
ptr_result.result.unwrap()
),
})
} }
)
Err(_) => rouille::Response::html(format!(
"The IP <b>{}</a> 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<String> = 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");
let server = "1.1.1.1:53".parse().expect("To parse");
let conn = TcpClientConnection::with_timeout(server, std::time::Duration::new(5, 0)).unwrap();
rouille::start_server("localhost:8000", move |request| {
let client = SyncClient::new(conn);
let conn = get_connection();
router!(request,
(GET) (/) => {
rouille::Response::html(FORM)
}, },
Err(_) => (GET) (/ping) => {
rouille::Response::html( rouille::Response::text("pong")
format!("The IP <b>{}</a> resolved as {:?} did not match known scanners patterns.", data.ip, ptr_result.result)
)
}
}, },
(POST) (/report) => {handle_report(&conn, client, &request)},
(POST) (/register) => { (POST) (/register) => {
let data = try_or_400!(post_input!(request, { let data = try_or_400!(post_input!(request, {
email: String, email: String,
@ -227,19 +260,7 @@ fn main() -> Result<()> {
}, },
(GET) (/scanners/{scanner_name: String}) => { (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(); handle_list_scanners(&conn, scanner_name)
let mut rows = stmt.query(named_params! { ":scanner_name": scanner_name }).unwrap();
let mut ips: Vec<String> = 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,
}
}, },
(GET) (/{api_key: String}/scanners/{scanner_name: String}) => { (GET) (/{api_key: String}/scanners/{scanner_name: String}) => {
let mut mac = HmacSha256::new_from_slice(b"my secret and secure key") let mut mac = HmacSha256::new_from_slice(b"my secret and secure key")