cool-mist
75 lines
2.0 KiB
Rust
75 lines
2.0 KiB
Rust
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
|
|
}
|