From 990758c27e388b131241ab0d51069f044e0c7b58 Mon Sep 17 00:00:00 2001 From: c0repwn3r Date: Thu, 23 Mar 2023 13:50:21 -0400 Subject: [PATCH] piuszgdpiugdszfoigfxliudxfg fixup thing and stuff and consistentifiy thhhhhhh --- tfclient/src/config.rs | 49 +++++++++++++++++++++++++++++++++- tfclient/src/daemon.rs | 2 +- tfclient/src/dirs.rs | 8 ++++++ tfclient/src/socketworker.rs | 51 ++++++++++++++++++++++++++++-------- 4 files changed, 97 insertions(+), 13 deletions(-) diff --git a/tfclient/src/config.rs b/tfclient/src/config.rs index d5b10e4..9341f6c 100644 --- a/tfclient/src/config.rs +++ b/tfclient/src/config.rs @@ -2,7 +2,7 @@ use std::error::Error; use std::fs; use log::{debug, info}; use serde::{Deserialize, Serialize}; -use crate::dirs::{get_config_dir, get_config_file}; +use crate::dirs::{get_cdata_dir, get_cdata_file, get_config_dir, get_config_file, get_data_dir}; pub const DEFAULT_PORT: u16 = 8157; fn default_port() -> u16 { DEFAULT_PORT } @@ -13,6 +13,11 @@ pub struct TFClientConfig { pub listen_port: u16 } +#[derive(Serialize, Deserialize, Clone)] +pub struct TFClientData { + pub host_id: Option +} + pub fn create_config(instance: &str) -> Result<(), Box> { info!("Creating config directory..."); fs::create_dir_all(get_config_dir(instance).ok_or("Unable to load config dir")?)?; @@ -39,4 +44,46 @@ pub fn load_config(instance: &str) -> Result> { let config: TFClientConfig = toml::from_str(&config_str)?; info!("Loaded config successfully"); Ok(config) +} + +pub fn create_cdata(instance: &str) -> Result<(), Box> { + info!("Creating data directory..."); + fs::create_dir_all(get_cdata_dir(instance).ok_or("Unable to load data dir")?)?; + info!("Copying default data file to config directory..."); + let config = TFClientData { host_id: None }; + let config_str = toml::to_string(&config)?; + fs::write(get_cdata_file(instance).ok_or("Unable to load data dir")?, config_str)?; + Ok(()) +} + +pub fn load_cdata(instance: &str) -> Result> { + info!("Loading cdata..."); + let config_file = get_cdata_file(instance).ok_or("Unable to load cdata dir")?; + + if !config_file.exists() { + create_cdata(instance)?; + } + + debug!("opening {}", config_file.as_path().display()); + let config_str = fs::read_to_string(config_file)?; + debug!("parsing cdata file"); + let config: TFClientData = toml::from_str(&config_str)?; + info!("Loaded cdata successfully"); + Ok(config) +} + +pub fn save_cdata(instance: &str, data: TFClientData) -> Result<(), Box> { + info!("Saving cdata..."); + let config_file = get_cdata_file(instance).ok_or("Unable to load cdata dir")?; + + if !config_file.exists() { + create_cdata(instance)?; + } + + debug!("serializing cdata file"); + let config: String = toml::to_string(&data)?; + debug!("writing to {}", config_file.as_path().display()); + fs::write(config_file, config)?; + info!("Saved cdata successfully"); + Ok(()) } \ No newline at end of file diff --git a/tfclient/src/daemon.rs b/tfclient/src/daemon.rs index 29b7aaa..c462e6d 100644 --- a/tfclient/src/daemon.rs +++ b/tfclient/src/daemon.rs @@ -86,7 +86,7 @@ pub fn daemon_main(name: String, server: String) { info!("Starting socket worker thread..."); let socket_thread = thread::spawn(move || { - socketworker_main(config, transmitter, rx_socket); + socketworker_main(config, name.clone(), transmitter, rx_socket); }); info!("Waiting for socket thread to exit..."); diff --git a/tfclient/src/dirs.rs b/tfclient/src/dirs.rs index 24d85ad..d7dd0c3 100644 --- a/tfclient/src/dirs.rs +++ b/tfclient/src/dirs.rs @@ -10,4 +10,12 @@ pub fn get_config_dir(instance: &str) -> Option { pub fn get_config_file(instance: &str) -> Option { get_config_dir(instance).map(|f| f.join("tfclient.toml")) +} + +pub fn get_cdata_dir(instance: &str) -> Option { + dirs::config_dir().map(|f| f.join("tfclient_data/").join(format!("{}/", instance))) +} + +pub fn get_cdata_file(instance: &str) -> Option { + get_cdata_dir(instance).map(|f| f.join("tfclient.toml")) } \ No newline at end of file diff --git a/tfclient/src/socketworker.rs b/tfclient/src/socketworker.rs index 71d0c80..a48c8b6 100644 --- a/tfclient/src/socketworker.rs +++ b/tfclient/src/socketworker.rs @@ -5,10 +5,11 @@ use std::{io, thread}; use std::io::{BufRead, BufReader, BufWriter, Read, Write}; use std::net::{IpAddr, Shutdown, SocketAddr, TcpListener, TcpStream}; use std::sync::mpsc::{Receiver, TryRecvError}; +use clap::builder::Str; use log::{debug, error, info, trace, warn}; use serde::{Deserialize, Serialize}; use crate::apiworker::APIWorkerMessage; -use crate::config::TFClientConfig; +use crate::config::{load_cdata, TFClientConfig}; use crate::daemon::ThreadMessageSender; use crate::nebulaworker::NebulaWorkerMessage; @@ -16,9 +17,9 @@ pub enum SocketWorkerMessage { Shutdown } -pub fn socketworker_main(config: TFClientConfig, transmitter: ThreadMessageSender, rx: Receiver) { +pub fn socketworker_main(config: TFClientConfig, instance: String, transmitter: ThreadMessageSender, rx: Receiver) { info!("socketworker_main called, entering realmain"); - match _main(config, transmitter, rx) { + match _main(config, instance, transmitter, rx) { Ok(_) => (), Err(e) => { error!("Error in socket thread: {}", e); @@ -26,7 +27,7 @@ pub fn socketworker_main(config: TFClientConfig, transmitter: ThreadMessageSende }; } -fn _main(config: TFClientConfig, transmitter: ThreadMessageSender, rx: Receiver) -> Result<(), Box> { +fn _main(config: TFClientConfig, instance: String, transmitter: ThreadMessageSender, rx: Receiver) -> Result<(), Box> { let listener = TcpListener::bind(SocketAddr::new(IpAddr::from([127, 0, 0, 1]), config.listen_port))?; listener.set_nonblocking(true)?; @@ -34,8 +35,10 @@ fn _main(config: TFClientConfig, transmitter: ThreadMessageSender, rx: Receiver< match listener.accept() { Ok(stream) => { let transmitter_clone = transmitter.clone(); + let config_clone = config.clone(); + let instance_clone = instance.clone(); thread::spawn(|| { - match handle_stream(stream, transmitter_clone) { + match handle_stream(stream, transmitter_clone, config_clone, instance_clone) { Ok(_) => (), Err(e) => { error!("Error in client thread: {}", e); @@ -71,9 +74,9 @@ fn _main(config: TFClientConfig, transmitter: ThreadMessageSender, rx: Receiver< Ok(()) } -fn handle_stream(stream: (TcpStream, SocketAddr), transmitter: ThreadMessageSender) -> Result<(), io::Error> { +fn handle_stream(stream: (TcpStream, SocketAddr), transmitter: ThreadMessageSender, config: TFClientConfig, instance: String) -> Result<(), io::Error> { info!("Incoming client"); - match handle_client(stream.0, transmitter) { + match handle_client(stream.0, transmitter, config, instance) { Ok(()) => (), Err(e) if e.kind() == io::ErrorKind::TimedOut => { warn!("Client timed out, connection aborted"); @@ -95,14 +98,16 @@ fn handle_stream(stream: (TcpStream, SocketAddr), transmitter: ThreadMessageSend Ok(()) } -fn handle_client(stream: TcpStream, transmitter: ThreadMessageSender) -> Result<(), io::Error> { +fn handle_client(stream: TcpStream, transmitter: ThreadMessageSender, config: TFClientConfig, instance: String) -> Result<(), io::Error> { info!("Handling connection from {}", stream.peer_addr()?); let mut client = Client { state: ClientState::WaitHello, reader: BufReader::new(&stream), writer: BufWriter::new(&stream), - stream: &stream + stream: &stream, + config, + instance, }; loop { @@ -137,7 +142,9 @@ struct Client<'a> { state: ClientState, reader: BufReader<&'a TcpStream>, writer: BufWriter<&'a TcpStream>, - stream: &'a TcpStream + stream: &'a TcpStream, + config: TFClientConfig, + instance: String } fn waithello_handle(client: &mut Client, _transmitter: &ThreadMessageSender, command: JsonMessage) -> Result { @@ -207,6 +214,20 @@ fn senthello_handle(client: &mut Client, transmitter: &ThreadMessageSender, comm error!("Error sending shutdown message to socket worker thread: {}", e); } } + }, + + JsonMessage::GetHostID {} => { + let data = match load_cdata(&client.instance) { + Ok(d) => d, + Err(e) => { + error!("Error loading cdata: {}", e); + panic!("{}", e); // TODO: Find a better way of handling this + } + }; + client.stream.write_all(&ctob(JsonMessage::HostID { + has_id: data.host_id.is_some(), + id: data.host_id + }))?; } _ => { @@ -245,10 +266,18 @@ enum JsonMessage { reason: DisconnectReason }, #[serde(rename = "shutdown")] - Shutdown {} + Shutdown {}, + #[serde(rename = "get_host_id")] + GetHostID {}, + #[serde(rename = "host_id")] + HostID { + has_id: bool, + id: Option + } } #[derive(Serialize, Deserialize, Debug)] +#[serde(tag = "type")] enum DisconnectReason { #[serde(rename = "unsupported_version")] UnsupportedVersion { expected: i32, got: i32 },