wdes/security
1
0
Fork 0
Code Actions Packages Releases Activity

Fixups and implement DB_FILE

Browse Source
This commit is contained in:
William Desportes
2024-06-23 13:43:50 +02:00
parent 3f3e069aab
commit 5d6297b7b5
1 changed files with 30 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

44
snow-scanner/src/main.rs
View File

@ -247,6 +247,15 @@ fn handle_scan(conn: &Mutex<Connection>, request: &Request) -> Response {
ips: String, ips: String,
})); }));
if data.username.len() < 4 {
return Response {
status_code: 422,
headers: vec![("Content-Type".into(), "text/plain; charset=utf-8".into())],
data: ResponseBody::from_string("Invalid username"),
upgrade: None,
};
}
let db = conn.lock().unwrap(); let db = conn.lock().unwrap();
let task_group_id = uuid7::uuid7(); let task_group_id = uuid7::uuid7();
@ -317,10 +326,9 @@ fn handle_list_scanners(conn: &Mutex<Connection>, scanner_name: String) -> Respo
} }
} }
fn get_connection() -> Connection { fn get_connection(db_path: &str) -> Connection {
let path = "./snow-scanner.sqlite";
let conn = Connection::open_with_flags( let conn = Connection::open_with_flags(
path, db_path,
OpenFlags::SQLITE_OPEN_READ_WRITE OpenFlags::SQLITE_OPEN_READ_WRITE
| OpenFlags::SQLITE_OPEN_CREATE | OpenFlags::SQLITE_OPEN_CREATE
| OpenFlags::SQLITE_OPEN_FULL_MUTEX, | OpenFlags::SQLITE_OPEN_FULL_MUTEX,
@ -328,9 +336,9 @@ fn get_connection() -> Connection {
.unwrap(); .unwrap();
conn.execute( conn.execute(
"CREATE TABLE IF NOT EXISTS scanners ( "CREATE TABLE IF NOT EXISTS scanners (
ip VARCHAR(255)NOT NULL, ip VARCHAR(255) NOT NULL,
ip_type TINYINT(1) NOT NULL, ip_type TINYINT(1) NOT NULL,
scanner_name VARCHAR(255), scanner_name VARCHAR(255) NOT NULL,
ip_ptr VARCHAR(255) NULL, ip_ptr VARCHAR(255) NULL,
created_at DATETIME NOT NULL, created_at DATETIME NOT NULL,
updated_at DATETIME NULL, updated_at DATETIME NULL,
@ -343,12 +351,12 @@ fn get_connection() -> Connection {
.unwrap(); .unwrap();
conn.execute( conn.execute(
"CREATE TABLE IF NOT EXISTS scan_tasks ( "CREATE TABLE IF NOT EXISTS scan_tasks (
task_group_id VARCHAR(255)NOT NULL, task_group_id VARCHAR(255) NOT NULL,
cidr VARCHAR(255)NOT NULL, cidr VARCHAR(255) NOT NULL,
created_by_username VARCHAR(255)NOT NULL, created_by_username VARCHAR(255) NOT NULL,
created_at DATETIMENOT NULL, created_at DATETIME NOT NULL,
updated_at DATETIME NULL, updated_at DATETIME NULL,
started_at DATETIME NOT NULL, started_at DATETIME NULL,
still_processing_at DATETIME NULL, still_processing_at DATETIME NULL,
ended_at DATETIME NULL, ended_at DATETIME NULL,
PRIMARY KEY (task_group_id, cidr) PRIMARY KEY (task_group_id, cidr)
@ -377,15 +385,23 @@ fn main() -> Result<()> {
println!("Now listening on {}", server_address); println!("Now listening on {}", server_address);
let conn = Mutex::new(get_connection()); let db_file: String = if let Ok(env) = env::var("DB_FILE") {
env
} else {
"./snow-scanner.sqlite".to_string()
};
println!("Database will be saved at: {}", db_file);
let conn = Mutex::new(get_connection(db_file.as_str()));
conn.lock() conn.lock()
.unwrap() .unwrap()
.execute("SELECT 0 WHERE 0;", named_params! {}) .execute("SELECT 0 WHERE 0;", named_params! {})
.expect("Failed to initialize database"); .expect("Failed to initialize database");
thread::spawn(|| loop { thread::spawn(move || loop {
let conn = get_connection(); let conn = get_connection(db_file.as_str());
let mut stmt = conn.prepare("SELECT task_group_id, cidr FROM scan_tasks WHERE started_at = 'NULL' ORDER BY created_at ASC").unwrap(); let mut stmt = conn.prepare("SELECT task_group_id, cidr FROM scan_tasks WHERE started_at IS NULL ORDER BY created_at ASC").unwrap();
let mut rows = stmt.query(named_params! {}).unwrap(); let mut rows = stmt.query(named_params! {}).unwrap();
println!("Waiting for jobs"); println!("Waiting for jobs");
while let Some(row) = rows.next().unwrap() { while let Some(row) = rows.next().unwrap() {