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 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<Scanners, ()> {
.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<Scanners, ()> {
}
}
// The HTML document of the home page.
static FORM: &str = r#"
<html>
<head>
@ -113,53 +110,19 @@ static FORM: &str = r#"
<p><input type="ip" name="ip" placeholder="An IPv4 or IPv6" /></p>
<p><button>Report this IP</button></p>
</form>
<form action="/scan" method="POST">
<p><textarea name="ips"></textarea></p>
<p><button>Scan</button></p>
</form>
</body>
</html>
"#;
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 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) => {
fn handle_report(
conn: &Connection,
client: SyncClient<TcpClientConnection>,
request: &Request,
) -> Response {
let data = try_or_400!(post_input!(request, {
ip: String,
}));
@ -184,24 +147,94 @@ fn main() -> Result<()> {
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! <b>{}</a> {:?}.", scanner.ip, ptr_result.result.unwrap()),
Scanners::Strechoid =>
format!("Reported a stretchoid agent! <b>{}</a> {:?}.", scanner.ip, ptr_result.result.unwrap())
rouille::Response::html(match scanner_name {
Scanners::Binaryedge => format!(
"Reported an escaped ninja! <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(_) =>
rouille::Response::html(
format!("The IP <b>{}</a> resolved as {:?} did not match known scanners patterns.", data.ip, ptr_result.result)
)
}
(GET) (/ping) => {
rouille::Response::text("pong")
},
(POST) (/report) => {handle_report(&conn, client, &request)},
(POST) (/register) => {
let data = try_or_400!(post_input!(request, {
email: String,
@ -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<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,
}
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")