Re-implement the ws client

This commit is contained in:
2024-10-03 12:49:33 +02:00
parent 25df2642e9
commit fd4d43596f
6 changed files with 179 additions and 115 deletions

View File

@ -1,14 +1,21 @@
use cidr::IpCidr;
use hickory_resolver::Name;
use rocket::futures::{stream::Next, SinkExt, StreamExt};
use rocket_ws::Message;
use std::{collections::HashMap, net::IpAddr, str::FromStr};
use crate::{worker::detection::detect_scanner_from_name, DbConn, Scanner};
use crate::worker::{
detection::detect_scanner_from_name,
modules::{Network, WorkerMessages},
};
use crate::{DbConn, Scanner};
pub struct Server {
pub clients: HashMap<u32, Worker>,
pub struct Server<'a> {
pub clients: HashMap<u32, Worker<'a>>,
pub new_scanners: HashMap<String, IpAddr>,
}
impl Server {
impl<'a> Server<'a> {
pub async fn commit(&mut self, conn: &mut DbConn) -> &Server {
for (name, query_address) in self.new_scanners.clone() {
let scanner_name = Name::from_str(name.as_str()).unwrap();
@ -44,18 +51,81 @@ impl Server {
}
}
#[derive(Debug, Clone)]
pub struct Worker {
pub authenticated: bool,
pub login: Option<String>,
pub struct Worker<'a> {
authenticated: bool,
login: Option<String>,
stream: &'a mut rocket_ws::stream::DuplexStream,
}
impl Worker {
pub fn initial() -> Worker {
impl<'a> Worker<'a> {
pub fn initial(stream: &mut rocket_ws::stream::DuplexStream) -> Worker {
info!("New worker");
Worker {
authenticated: false,
login: None,
stream,
}
}
pub async fn send(&mut self, msg: Message) -> Result<(), rocket_ws::result::Error> {
self.stream.send(msg).await
}
pub fn next(&mut self) -> Next<'_, rocket_ws::stream::DuplexStream> {
self.stream.next()
}
pub async fn poll(&mut self) -> Result<bool, ()> {
let message = self.next();
match message.await {
Some(Ok(message)) => {
match message {
rocket_ws::Message::Text(_) => match self.on_message(&message).await {
Ok(_) => {}
Err(err) => error!("Processing error: {err}"),
},
rocket_ws::Message::Binary(data) => {
// Handle Binary message
info!("Received Binary message: {:?}", data);
}
rocket_ws::Message::Close(close_frame) => {
// Handle Close message
info!("Received Close message: {:?}", close_frame);
let close_frame = rocket_ws::frame::CloseFrame {
code: rocket_ws::frame::CloseCode::Normal,
reason: "Client disconected".to_string().into(),
};
let _ = self.stream.close(Some(close_frame)).await;
return Ok(true);
}
rocket_ws::Message::Ping(ping_data) => {
match self.send(rocket_ws::Message::Pong(ping_data)).await {
Ok(_) => {}
Err(err) => error!("Processing error: {err}"),
}
}
rocket_ws::Message::Pong(pong_data) => {
// Handle Pong message
info!("Received Pong message: {:?}", pong_data);
}
_ => {
info!("Received other message: {:?}", message);
}
};
Ok(false)
}
Some(Err(err)) => {
info!("Connection closed");
let close_frame = rocket_ws::frame::CloseFrame {
code: rocket_ws::frame::CloseCode::Normal,
reason: "Client disconected".to_string().into(),
};
let _ = self.stream.close(Some(close_frame)).await;
// The connection is closed by the client
Ok(true)
}
None => Ok(false),
}
}
@ -77,63 +147,43 @@ impl Worker {
self.authenticated = true;
self
}
}
/*
impl ws2::Handler for Server {
fn on_open(&mut self, ws: &WebSocket) -> Pod {
info!("New client: {ws}");
let worker = Worker::initial();
// Add the client
self.clients.insert(ws.id(), worker);
Ok(())
}
fn on_close(&mut self, ws: &WebSocket) -> Pod {
info!("Client /quit: {ws}");
// Drop the client
self.clients.remove(&ws.id());
Ok(())
}
fn on_message(&mut self, ws: &WebSocket, msg: String) -> Pod {
let client = self.clients.get_mut(&ws.id());
if client.is_none() {
// Impossible, close in case
return ws.close();
}
let worker: &mut Worker = client.unwrap();
info!("on message: {msg}, {ws}");
pub async fn on_message(&mut self, msg: &Message) -> Result<(), String> {
info!("on message: {msg}");
let mut worker_reply: Option<WorkerMessages> = None;
let worker_request: WorkerMessages = msg.clone().into();
let worker_request: WorkerMessages = match msg.clone().try_into() {
Ok(data) => data,
Err(err) => return Err(err),
};
let result = match worker_request {
WorkerMessages::AuthenticateRequest { login } => {
if !worker.is_authenticated() {
worker.authenticate(login);
if !self.is_authenticated() {
self.authenticate(login);
return Ok(());
} else {
error!("Already authenticated: {ws}");
error!("Already authenticated");
return Ok(());
}
}
WorkerMessages::ScannerFoundResponse { name, address } => {
info!("Detected {name} for {address}");
self.new_scanners.insert(name, address);
//self.new_scanners.insert(name, address);
Ok(())
}
WorkerMessages::GetWorkRequest {} => {
let net = IpCidr::from_str("52.189.78.0/24").unwrap();
worker_reply = Some(WorkerMessages::DoWorkRequest {
neworks: vec![Network(IpCidr::from_str("52.189.78.0/24")?)],
neworks: vec![Network(net)],
});
Ok(())
}
WorkerMessages::DoWorkRequest { .. } | WorkerMessages::Invalid { .. } => {
error!("Unable to understand: {msg}, {ws}");
error!("Unable to understand: {msg}");
// Unable to understand, close the connection
return ws.close();
//return ws.close();
Err("Unable to understand: {msg}}")
} /*msg => {
error!("No implemented: {:#?}", msg);
Ok(())
@ -143,14 +193,14 @@ impl ws2::Handler for Server {
// it has a request to send
if let Some(worker_reply) = worker_reply {
let msg_string: String = worker_reply.to_string();
match ws.send(msg_string) {
match self.send(rocket_ws::Message::Text(msg_string)).await {
Ok(_) => match worker_reply {
WorkerMessages::DoWorkRequest { .. } => {}
msg => error!("No implemented: {:#?}", msg),
},
Err(err) => error!("Error sending reply to {ws}: {err}"),
Err(err) => error!("Error sending reply: {err}"),
}
}
result
Ok(result?)
}
}*/
}