Broadcast work requests to each node

This commit is contained in:
2024-10-07 00:15:29 +02:00
parent e5c3b38121
commit f589d4c11e
3 changed files with 35 additions and 21 deletions

View File

@ -3,6 +3,7 @@ use chrono::{NaiveDateTime, Utc};
#[macro_use]
extern crate rocket;
use cidr::IpCidr;
use event_bus::{EventBusSubscriber, EventBusWriter};
use rocket::{fairing::AdHoc, futures::SinkExt, trace::error, Rocket, State};
use rocket_db_pools::{
@ -18,6 +19,8 @@ use rocket_db_pools::{
Connection,
};
use crate::worker::modules::{Network, WorkerMessages};
use rocket_db_pools::diesel::mysql::{Mysql, MysqlValue};
use rocket_db_pools::diesel::serialize::IsNull;
use rocket_db_pools::diesel::sql_types::Text;
@ -29,9 +32,9 @@ use rocket_ws::WebSocket;
use server::Server;
use worker::detection::{detect_scanner, get_dns_client, Scanners};
use std::path::PathBuf;
use std::{env, fmt};
use std::{io::Write, net::SocketAddr};
use std::{path::PathBuf, str::FromStr};
use uuid::Uuid;
use serde::{Deserialize, Deserializer, Serialize};
@ -220,17 +223,21 @@ enum MultiReply {
}
#[post("/scan", data = "<form>")]
async fn handle_scan(mut db: DbConn, form: Form<ScanParams<'_>>) -> MultiReply {
async fn handle_scan(
mut db: DbConn,
form: Form<ScanParams<'_>>,
event_bus_writer: &State<EventBusWriter>,
) -> MultiReply {
if form.username.len() < 4 {
return MultiReply::FormError(PlainText("Invalid username".to_string()));
}
let task_group_id: Uuid = Uuid::now_v7();
for ip in form.ips.lines() {
for cidr in form.ips.lines() {
let scan_task = ScanTask {
task_group_id: task_group_id.to_string(),
cidr: ip.to_string(),
cidr: cidr.to_string(),
created_by_username: form.username.to_string(),
created_at: Utc::now().naive_utc(),
updated_at: None,
@ -238,8 +245,19 @@ async fn handle_scan(mut db: DbConn, form: Form<ScanParams<'_>>) -> MultiReply {
still_processing_at: None,
ended_at: None,
};
let mut bus_tx = event_bus_writer.write();
match scan_task.save(&mut db).await {
Ok(_) => error!("Added {}", ip.to_string()),
Ok(_) => {
info!("Added {}", cidr.to_string());
let net = IpCidr::from_str(cidr).unwrap();
let msg = WorkerMessages::DoWorkRequest {
neworks: vec![Network(net)],
}
.into();
let _ = bus_tx.send(msg).await;
}
Err(err) => error!("Not added: {:?}", err),
}
}
@ -460,11 +478,7 @@ async fn index() -> HtmlContents {
}
#[get("/ping")]
async fn pong(event_bus_writer: &State<EventBusWriter>) -> PlainText {
let mut bus_tx = event_bus_writer.write();
let _ = bus_tx
.send(rocket_ws::Message::Text("Groumpf!".to_string()))
.await;
async fn pong() -> PlainText {
PlainText("pong".to_string())
}

View File

@ -1,11 +1,10 @@
use cidr::IpCidr;
use rocket::futures::{stream::Next, SinkExt, StreamExt};
use rocket_ws::{frame::CloseFrame, Message};
use std::{pin::Pin, str::FromStr};
use std::pin::Pin;
use crate::{
event_bus::{EventBusEvent, EventBusWriter},
worker::modules::{Network, WorkerMessages},
worker::modules::WorkerMessages,
};
use rocket::futures::channel::mpsc as rocket_mpsc;
@ -21,10 +20,8 @@ impl WsChat {
use crate::rocket::futures::StreamExt;
use rocket::tokio;
let _ = bus_tx
.send(rocket_ws::Message::Text("I am new !".to_string()))
.await;
//SharedData::send_to_all(&workers, "I am new !");
let _ = bus_tx.send(rocket_ws::Message::Ping(vec![])).await;
let mut worker = Worker::initial(&mut stream);
let mut interval = rocket::tokio::time::interval(std::time::Duration::from_secs(60));
loop {
@ -279,10 +276,7 @@ impl<'a> Worker<'a> {
Ok(())
}
WorkerMessages::GetWorkRequest {} => {
let net = IpCidr::from_str("52.189.78.0/24").unwrap();
worker_reply = Some(WorkerMessages::DoWorkRequest {
neworks: vec![Network(net)],
});
worker_reply = Some(WorkerMessages::DoWorkRequest { neworks: vec![] });
Ok(())
}
WorkerMessages::DoWorkRequest { .. } | WorkerMessages::Invalid { .. } => {

View File

@ -67,6 +67,12 @@ impl Into<WorkerMessages> for String {
}
}
impl Into<RocketMessage> for WorkerMessages {
fn into(self) -> RocketMessage {
RocketMessage::Text(self.to_string())
}
}
impl TryInto<WorkerMessages> for RocketMessage {
type Error = String;