249 lines
7.9 KiB
Rust
249 lines
7.9 KiB
Rust
use std::net::IpAddr;
|
|
|
|
use crate::{DbConn, Scanners};
|
|
use chrono::{NaiveDateTime, Utc};
|
|
use hickory_resolver::Name;
|
|
use rocket_db_pools::diesel::{dsl::insert_into, prelude::*, result::Error as DieselError};
|
|
|
|
use crate::schema::scan_tasks::dsl::scan_tasks;
|
|
use crate::schema::scanners::dsl::scanners;
|
|
|
|
#[derive(Queryable, Selectable, Debug)]
|
|
#[diesel(table_name = crate::schema::scanners)]
|
|
#[diesel(check_for_backend(diesel::mysql::Mysql))]
|
|
pub struct Scanner {
|
|
pub ip: String,
|
|
pub ip_type: u8,
|
|
pub scanner_name: Scanners,
|
|
pub ip_ptr: Option<String>,
|
|
pub created_at: NaiveDateTime,
|
|
pub updated_at: Option<NaiveDateTime>,
|
|
pub last_seen_at: Option<NaiveDateTime>,
|
|
pub last_checked_at: Option<NaiveDateTime>,
|
|
}
|
|
|
|
impl Scanner {
|
|
pub async fn find_or_new(
|
|
query_address: IpAddr,
|
|
scanner_name: Scanners,
|
|
ptr: Option<Name>,
|
|
conn: &mut DbConn,
|
|
) -> Result<Scanner, DieselError> {
|
|
let ip_type = if query_address.is_ipv6() { 6 } else { 4 };
|
|
let scanner_row_result = Scanner::find(query_address.to_string(), ip_type, conn).await;
|
|
let scanner_row = match scanner_row_result {
|
|
Ok(scanner_row) => scanner_row,
|
|
Err(err) => return Err(err),
|
|
};
|
|
|
|
let scanner = if let Some(mut scanner) = scanner_row {
|
|
scanner.last_seen_at = Some(Utc::now().naive_utc());
|
|
scanner.last_checked_at = Some(Utc::now().naive_utc());
|
|
scanner.updated_at = Some(Utc::now().naive_utc());
|
|
scanner
|
|
} else {
|
|
Scanner {
|
|
ip: query_address.to_string(),
|
|
ip_type: ip_type,
|
|
scanner_name: scanner_name.clone(),
|
|
ip_ptr: match ptr {
|
|
Some(ptr) => Some(ptr.to_string()),
|
|
None => None,
|
|
},
|
|
created_at: Utc::now().naive_utc(),
|
|
updated_at: None,
|
|
last_seen_at: None,
|
|
last_checked_at: None,
|
|
}
|
|
};
|
|
match scanner.save(conn).await {
|
|
Ok(scanner) => Ok(scanner),
|
|
Err(err) => Err(err),
|
|
}
|
|
}
|
|
|
|
pub async fn find(
|
|
ip_address: String,
|
|
ip_type: u8,
|
|
conn: &mut DbConn,
|
|
) -> Result<Option<Scanner>, DieselError> {
|
|
use crate::schema::scanners;
|
|
scanners
|
|
.select(Scanner::as_select())
|
|
.filter(scanners::ip.eq(ip_address))
|
|
.filter(scanners::ip_type.eq(ip_type))
|
|
.order((scanners::ip_type.desc(), scanners::created_at.desc()))
|
|
.first(conn)
|
|
.await
|
|
.optional()
|
|
}
|
|
|
|
pub async fn list_names(
|
|
scanner_name: Scanners,
|
|
conn: &mut DbConn,
|
|
) -> Result<Vec<String>, DieselError> {
|
|
use crate::schema::scanners;
|
|
use crate::schema::scanners::ip;
|
|
|
|
scanners
|
|
.select(ip)
|
|
.filter(scanners::scanner_name.eq(scanner_name.to_string()))
|
|
.order((scanners::ip_type.desc(), scanners::created_at.desc()))
|
|
.load::<String>(conn)
|
|
.await
|
|
}
|
|
|
|
pub async fn save(self: Scanner, conn: &mut DbConn) -> Result<Scanner, DieselError> {
|
|
use crate::schema::scanners;
|
|
|
|
let new_scanner: NewScanner = NewScanner::from_scanner(&self).await;
|
|
match insert_into(scanners::table)
|
|
.values(&new_scanner)
|
|
.on_conflict(diesel::dsl::DuplicatedKeys)
|
|
.do_update()
|
|
.set(&new_scanner)
|
|
.execute(conn)
|
|
.await
|
|
{
|
|
Ok(_) => Ok(self),
|
|
Err(err) => Err(err),
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Insertable, AsChangeset)]
|
|
#[diesel(table_name = crate::schema::scanners)]
|
|
#[diesel(check_for_backend(diesel::mysql::Mysql))]
|
|
pub struct NewScanner {
|
|
pub ip: String,
|
|
pub ip_type: u8,
|
|
pub scanner_name: String,
|
|
pub ip_ptr: Option<String>,
|
|
pub created_at: NaiveDateTime,
|
|
pub updated_at: Option<NaiveDateTime>,
|
|
pub last_seen_at: Option<NaiveDateTime>,
|
|
pub last_checked_at: Option<NaiveDateTime>,
|
|
}
|
|
|
|
impl NewScanner {
|
|
pub async fn from_scanner<'x>(scanner: &Scanner) -> NewScanner {
|
|
NewScanner {
|
|
ip: scanner.ip.to_string(),
|
|
ip_type: scanner.ip_type,
|
|
scanner_name: scanner.scanner_name.to_string(),
|
|
ip_ptr: scanner.ip_ptr.to_owned(),
|
|
created_at: scanner.created_at,
|
|
updated_at: scanner.updated_at,
|
|
last_seen_at: scanner.last_seen_at,
|
|
last_checked_at: scanner.last_checked_at,
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Queryable, Selectable, Debug)]
|
|
#[diesel(table_name = crate::schema::scan_tasks)]
|
|
#[diesel(check_for_backend(diesel::mysql::Mysql))]
|
|
pub struct ScanTask {
|
|
pub task_group_id: String,
|
|
pub cidr: String,
|
|
pub created_by_username: String,
|
|
pub created_at: NaiveDateTime,
|
|
pub updated_at: Option<NaiveDateTime>,
|
|
pub started_at: Option<NaiveDateTime>,
|
|
pub still_processing_at: Option<NaiveDateTime>,
|
|
pub ended_at: Option<NaiveDateTime>,
|
|
}
|
|
|
|
#[derive(Selectable, Debug, Queryable)]
|
|
#[diesel(table_name = crate::schema::scan_tasks)]
|
|
#[diesel(check_for_backend(diesel::mysql::Mysql))]
|
|
pub struct ScanTaskitem {
|
|
pub task_group_id: String,
|
|
pub cidr: String,
|
|
pub created_by_username: String,
|
|
pub started_at: Option<NaiveDateTime>,
|
|
pub still_processing_at: Option<NaiveDateTime>,
|
|
pub ended_at: Option<NaiveDateTime>,
|
|
}
|
|
|
|
impl ScanTask {
|
|
pub async fn list_not_started(mut conn: DbConn) -> Result<Vec<ScanTaskitem>, DieselError> {
|
|
use crate::schema::scan_tasks;
|
|
|
|
let res = scan_tasks
|
|
.select(ScanTaskitem::as_select())
|
|
.filter(scan_tasks::started_at.is_null())
|
|
.order((scan_tasks::created_at.asc(),))
|
|
.load::<ScanTaskitem>(&mut conn)
|
|
.await;
|
|
match res {
|
|
Ok(rows) => Ok(rows),
|
|
Err(err) => Err(err),
|
|
}
|
|
}
|
|
|
|
pub async fn list(conn: &mut DbConn) -> Result<Vec<ScanTaskitem>, DieselError> {
|
|
use crate::schema::scan_tasks;
|
|
|
|
let res = scan_tasks
|
|
.select(ScanTaskitem::as_select())
|
|
.order((
|
|
scan_tasks::created_at.desc(),
|
|
scan_tasks::task_group_id.asc(),
|
|
))
|
|
.load::<ScanTaskitem>(conn)
|
|
.await;
|
|
match res {
|
|
Ok(rows) => Ok(rows),
|
|
Err(err) => Err(err),
|
|
}
|
|
}
|
|
|
|
pub async fn save(self: &ScanTask, conn: &mut DbConn) -> Result<(), DieselError> {
|
|
use crate::schema::scan_tasks;
|
|
|
|
let new_scan_task: NewScanTask = NewScanTask::from_scan_task(self).await;
|
|
|
|
match insert_into(scan_tasks::table)
|
|
.values(&new_scan_task)
|
|
.on_conflict(diesel::dsl::DuplicatedKeys)
|
|
.do_update()
|
|
.set(&new_scan_task)
|
|
.execute(conn)
|
|
.await
|
|
{
|
|
Ok(_) => Ok(()),
|
|
Err(err) => Err(err),
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Insertable, AsChangeset)]
|
|
#[diesel(table_name = crate::schema::scan_tasks)]
|
|
#[diesel(check_for_backend(diesel::mysql::Mysql))]
|
|
pub struct NewScanTask {
|
|
pub task_group_id: String,
|
|
pub cidr: String,
|
|
pub created_by_username: String,
|
|
pub created_at: NaiveDateTime,
|
|
pub updated_at: Option<NaiveDateTime>,
|
|
pub started_at: Option<NaiveDateTime>,
|
|
pub still_processing_at: Option<NaiveDateTime>,
|
|
pub ended_at: Option<NaiveDateTime>,
|
|
}
|
|
|
|
impl NewScanTask {
|
|
pub async fn from_scan_task<'x>(scan_task: &ScanTask) -> NewScanTask {
|
|
NewScanTask {
|
|
task_group_id: scan_task.task_group_id.to_string(),
|
|
cidr: scan_task.cidr.to_owned(),
|
|
created_by_username: scan_task.created_by_username.to_owned(),
|
|
created_at: scan_task.created_at,
|
|
updated_at: scan_task.updated_at,
|
|
started_at: scan_task.started_at,
|
|
still_processing_at: scan_task.still_processing_at,
|
|
ended_at: scan_task.ended_at,
|
|
}
|
|
}
|
|
}
|