Xory/skylink
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat: move ws to separate file & implement c2 commands

Browse source
This commit is contained in:
Xory 2025-12-06 20:23:26 +02:00
parent a521782a37
commit 4cd62f4f43
7 changed files with 272 additions and 252 deletions
Download patch file Download diff file Expand all files Collapse all files

87
src/main.rs
View file

@ -1,97 +1,14 @@
use futures_util::{SinkExt, stream::StreamExt};
use skylink::{ping_job, reconnect_websocket, WsTx};
use skylink::lib::logger::{LogLevel, log};
use skylink::{LOG_PATH, WS_URL};
use skylink::{WsTx, eval_command, websockets::ping_job, websockets::reconnect_websocket};
use std::sync::Arc;
use tokio::sync::Mutex;
use tokio_tungstenite::tungstenite::protocol::Message;
// Some parts of this function were generated by an LLM. I'm taking note of this in case a
// weird barely detectable bug pops up, as LLMs tend to generate.
async fn websocket_handler(ws_tx: WsTx) {
use std::time::Duration;
use tokio_tungstenite::connect_async;
loop {
let connection_result = connect_async(WS_URL).await;
let ws_stream = match connection_result {
Ok((stream, _)) => {
log(LogLevel::Info, LOG_PATH, "[ws] WebSocket connection established.".to_string()).await;
stream
}
Err(e) => {
log(LogLevel::Warning, LOG_PATH, format!("[ws] Failed to connect: {:?}. Retrying in 5s....", e)).await;
tokio::time::sleep(Duration::from_secs(5)).await;
continue;
}
};
let (ws_send, mut ws_recv) = ws_stream.split();
{
let mut unlocked = ws_tx.lock().await;
*unlocked = Some(ws_send);
}
let ws_tx_clone = ws_tx.clone();
tokio::spawn(async move {
loop {
if let Err(_) = ping_job(ws_tx_clone.clone()).await {
reconnect_websocket(ws_tx_clone.clone()).await;
}
}
});
while let Some(msg) = ws_recv.next().await {
match msg {
Ok(Message::Text(text)) => {
log(LogLevel::Debug, LOG_PATH, format!("[ws] received text: {}", &text)).await;
}
Ok(Message::Close(_)) => {
log(LogLevel::Warning, LOG_PATH, format!("[ws] received close frame, disconnecting.")).await;
break;
}
Ok(Message::Ping(h)) => {
log(LogLevel::Debug, LOG_PATH, format!("[ws] received ping, sending pong")).await;
let mut unlocked = ws_tx.lock().await;
match unlocked.as_mut() {
Some(v) => {
if let Err(e) = v.send(Message::Pong(h)).await {
log(LogLevel::Error, LOG_PATH, format!("[ws] failed to send pong: {e}, reconnecting.")).await;
break;
}
}
None => {
log(LogLevel::Error, LOG_PATH, format!("[ws] failed to respond: no sink, reconnecting.")).await;
break;
}
}
}
Err(e) => {
log(LogLevel::Error, LOG_PATH, format!("[ws] error receiving message: {:?}.", e)).await;
break;
}
_ => {}
}
}
{
let mut unlocked_ws_tx = ws_tx.lock().await;
*unlocked_ws_tx = None;
}
log(LogLevel::Error, LOG_PATH, format!("[ws] connection lost.")).await;
tokio::time::sleep(std::time::Duration::from_secs(5)).await;
}
}
#[tokio::main]
async fn main() -> anyhow::Result<()> {
log(LogLevel::Info, LOG_PATH, format!("[main] Skylink version 1.0.0 starting...")).await;
let ws_tx: WsTx = Arc::new(Mutex::new(None));
let ws_tx_for_handler = Arc::clone(&ws_tx);
websocket_handler(ws_tx_for_handler).await;