single tool

src/daemon.rs

@@ -0,0 +1,74 @@
+use crate::error::{WError, WResult};
+use std::{
+    io::Read,
+    net::{TcpListener, TcpStream},
+    process::{Command, Stdio},
+    time::Duration,
+};
+
+pub fn daemon() -> WResult<()> {
+    let address = "0.0.0.0:21367";
+    let listener = TcpListener::bind(address).map_err(WError::DaemonSocketBindError)?;
+    println!("Listening for TCP packets at {}", address);
+    for mut stream in listener.incoming() {
+        match &mut stream {
+            Ok(stream) => {
+                if let Err(e) = handle_connection(stream) {
+                    eprintln!("Failed to handle request {}", e);
+                }
+            }
+            Err(e) => eprintln!("{}", e),
+        }
+    }
+
+    Ok(())
+}
+
+fn handle_connection(stream: &mut TcpStream) -> WResult<()> {
+    // TODO: Slowloris
+    stream
+        .set_read_timeout(Some(Duration::from_secs(5)))
+        .map_err(WError::DaemonConnectionReadError)?;
+    let command = read_command(stream)?;
+    match command.as_str() {
+        "shutdown" => execute_shutdown_command(),
+        _ => Err(WError::DaemonInvalidRequest(command)),
+    }
+}
+
+fn execute_shutdown_command() -> WResult<()> {
+    let mut shutdown_cmd = build_shutdown_cmd();
+    shutdown_cmd.stdout(Stdio::piped());
+    shutdown_cmd.stderr(Stdio::piped());
+
+    let output = shutdown_cmd
+        .output()
+        .map_err(|e| WError::DaemonCommandError(e.to_string()))?;
+    let status_code = output
+        .status
+        .code()
+        .ok_or(WError::DaemonCommandError("No exit code returned!".to_string()))?;
+
+    if status_code != 0 {
+        return Err(WError::DaemonCommandExecutionError(output));
+    };
+
+    Ok(())
+}
+
+fn read_command(stream: &mut TcpStream) -> WResult<String> {
+    let mut buf = Vec::new();
+    let read = stream
+        .read_to_end(&mut buf)
+        .map_err(WError::DaemonConnectionReadError)?;
+    let command = String::from_utf8_lossy(&buf[..read]).to_string();
+    Ok(command)
+}
+
+fn build_shutdown_cmd() -> Command {
+    let mut cmd = Command::new("systemctl");
+    cmd.arg("-k");
+    cmd.arg("poweroff");
+
+    cmd
+}