diff --git a/sparse-05/sparse-05-client/src/commands/connect/commands/download.rs b/sparse-05/sparse-05-client/src/commands/connect/commands/download.rs new file mode 100644 index 0000000..e69de29 diff --git a/sparse-05/sparse-05-client/src/commands/connect/commands/help.rs b/sparse-05/sparse-05-client/src/commands/connect/commands/help.rs index 7151191..015700f 100644 --- a/sparse-05/sparse-05-client/src/commands/connect/commands/help.rs +++ b/sparse-05/sparse-05-client/src/commands/connect/commands/help.rs @@ -6,13 +6,29 @@ use crate::commands::connect::shell::SparseCommands; pub fn print_help(arg: Option) { match arg { Some(SparseCommands::Exit) => println!( - "Exits from the shell and disconnects from the binary" + "exit\n\ + \n\ + Exits from the shell and disconnects from the remote server" ), Some(SparseCommands::SysInfo) => println!( - "Prints system information from the system you are connecting to" + "#sysinfo\n\ + \n\ + Prints system information from the system you are connecting to" ), Some(SparseCommands::Cd { .. }) => println!( - "Changes the current working directory you are in" + "cd [directory]\n\ + \n\ + Changes the current working directory you are in" + ), + Some(SparseCommands::UploadFile { .. }) => println!( + "#uploadfile [local file path] [remote file path]\n\ + \n\ + Uploads the file at the local file path, overwriting the file at the remote file path" + ), + Some(SparseCommands::DownloadFile { .. }) => println!( + "#uploadfile [local file path] [remote file path]\n\ + \n\ + Uploads the file at the local file path, overwriting the file at the remote file path" ), None => println!( "\n{}{}\n\ @@ -21,7 +37,7 @@ pub fn print_help(arg: Option) { \n\ - #sysinfo\t\tprint information about the system you are connecting to - #help\t\tprints this help page, or alternatively prints info about a command passed as an argument\n\ - - #exit\t\texit from the shell and disconnect from the binary\n\ + - #exit\t\texit from the shell and disconnect from the remote server\n\ ", Style::new().bold().paint("SHELL COMMANDS"), Style::new().paint(""), diff --git a/sparse-05/sparse-05-client/src/commands/connect/commands/mod.rs b/sparse-05/sparse-05-client/src/commands/connect/commands/mod.rs index 35bdc0c..9de8e6c 100644 --- a/sparse-05/sparse-05-client/src/commands/connect/commands/mod.rs +++ b/sparse-05/sparse-05-client/src/commands/connect/commands/mod.rs @@ -1,2 +1,4 @@ +pub mod download; pub mod help; pub mod sysinfo; +pub mod upload; diff --git a/sparse-05/sparse-05-client/src/commands/connect/commands/upload.rs b/sparse-05/sparse-05-client/src/commands/connect/commands/upload.rs new file mode 100644 index 0000000..121051b --- /dev/null +++ b/sparse-05/sparse-05-client/src/commands/connect/commands/upload.rs @@ -0,0 +1,11 @@ +use std::{path::PathBuf, sync::Arc}; + +use crate::commands::connect::Connection; + +pub async fn upload_file( + conn: Arc, + local_path: PathBuf, + remote_path: PathBuf, +) -> anyhow::Result<()> { + Ok(()) +} diff --git a/sparse-05/sparse-05-client/src/commands/connect/mod.rs b/sparse-05/sparse-05-client/src/commands/connect/mod.rs index 70991b3..bbd02be 100644 --- a/sparse-05/sparse-05-client/src/commands/connect/mod.rs +++ b/sparse-05/sparse-05-client/src/commands/connect/mod.rs @@ -17,7 +17,7 @@ use crate::configs::ClientConfig; mod commands; mod shell; -struct Connection { +pub struct Connection { config: ClientConfig, foreign_sign_pubkey: ed25519_dalek::PublicKey, foreign_enc_pubkey: ecies_ed25519::PublicKey, @@ -26,7 +26,7 @@ struct Connection { } impl Connection { - fn decrypt_and_verify(&self, data: &[u8]) -> anyhow::Result { + pub fn decrypt_and_verify(&self, data: &[u8]) -> anyhow::Result { if data.len() < 65 { bail!("unable to parse out signature from message"); } @@ -40,7 +40,7 @@ impl Connection { Ok(rmp_serde::from_slice(&data)?) } - fn encrypt_and_sign(&self, data: &[u8]) -> anyhow::Result> { + pub fn encrypt_and_sign(&self, data: &[u8]) -> anyhow::Result> { let mut rng = rand::thread_rng(); let data = ecies_ed25519::encrypt(&self.foreign_enc_pubkey, data, &mut rng)?; @@ -49,7 +49,7 @@ impl Connection { Ok([&signature.to_bytes(), &*data].concat()) } - async fn send_command(&self, command: Command) -> anyhow::Result<()> { + pub async fn send_command(&self, command: Command) -> anyhow::Result<()> { let cmd = rmp_serde::to_vec(&command)?; self.socket @@ -59,7 +59,7 @@ impl Connection { Ok(()) } - async fn get_response(&self) -> anyhow::Result { + pub async fn get_response(&self) -> anyhow::Result { let mut buffer = [0u8; 2000]; let (read, from) = self.socket.recv_from(&mut buffer).await?; diff --git a/sparse-05/sparse-05-client/src/commands/connect/shell.rs b/sparse-05/sparse-05-client/src/commands/connect/shell.rs index 1f9428c..1e23c89 100644 --- a/sparse-05/sparse-05-client/src/commands/connect/shell.rs +++ b/sparse-05/sparse-05-client/src/commands/connect/shell.rs @@ -1,20 +1,17 @@ use std::{ - io::{self, stdin, Read, Stdin, Write}, + io::{self, Read, Write}, os::fd::AsRawFd, path::PathBuf, - sync::{ - mpsc::{channel, TryRecvError}, - Arc, - }, + sync::Arc, thread::{self, scope}, }; use sparse_05_common::messages::{Capabilities, Command, Response}; use structopt::StructOpt; use tokio::{ - io::{stderr, stdout, AsyncWriteExt}, + io::{stderr, stdout, AsyncReadExt, AsyncWriteExt}, runtime::Handle, - sync::mpsc, + sync::mpsc::{channel, Receiver}, }; use super::{commands, Connection}; @@ -29,6 +26,16 @@ pub enum SparseCommands { Cd { folder: PathBuf, }, + #[structopt(name = "#upload")] + UploadFile { + local_file: PathBuf, + remote_path: PathBuf, + }, + #[structopt(name = "#download")] + DownloadFile { + remote_file: PathBuf, + local_path: PathBuf, + }, } macro_rules! libc_try { @@ -58,7 +65,11 @@ fn convert_termios_raw(attrs: &mut libc::termios) -> anyhow::Result<()> { Ok(()) } -async fn run_command(command: String, connection: Arc) -> anyhow::Result<()> { +async fn run_command( + stdin: &mut Receiver>, + command: String, + connection: Arc, +) -> anyhow::Result<()> { connection .send_command(Command::RunCommand(command)) .await?; @@ -70,38 +81,6 @@ async fn run_command(command: String, connection: Arc) -> anyhow::Re } }; - let (kill, handle_kill) = channel(); - let (send_message, mut handle_send_message) = mpsc::channel(16); - let handle = Handle::current(); - - let stdin_thread = thread::spawn({ - let mut stdin = stdin(); - let connection = Arc::clone(&connection); - - move || { - let mut stdin_buffer = [0u8; 1024]; - - loop { - let Ok(amount) = stdin.read(&mut stdin_buffer) else { - continue; - }; - - handle.spawn({ - let message = stdin_buffer[..amount].to_vec(); - let send_message = send_message.clone(); - async move { send_message.send(message).await } - }); - - match handle_kill.try_recv() { - Ok(()) | Err(TryRecvError::Disconnected) => { - break; - } - Err(_) => {} - } - } - } - }); - loop { enum Event { Stdin(Vec), @@ -110,7 +89,7 @@ async fn run_command(command: String, connection: Arc) -> anyhow::Re let Some(event) = tokio::select! { v = connection.get_response() => v.ok().map(Event::Remote), - v = handle_send_message.recv() => v.map(Event::Stdin) + v = stdin.recv() => v.map(Event::Stdin) } else { continue; }; @@ -122,8 +101,9 @@ async fn run_command(command: String, connection: Arc) -> anyhow::Re Event::Remote(Response::SendStderr(bytes, cid)) if cid == id => { stderr().write(&bytes).await?; } - Event::Remote(Response::CommandDone(cid, code)) if cid == id => break, + Event::Remote(Response::CommandDone(cid, _)) if cid == id => break, Event::Stdin(stdin) => { + println!("here"); let _ = connection .send_command(Command::SendStdin(stdin, id)) .await?; @@ -132,20 +112,17 @@ async fn run_command(command: String, connection: Arc) -> anyhow::Re } } - let _ = kill.send(()); - let _ = stdin_thread.join(); - Ok(()) } pub(super) async fn shell( connection: Arc, - mut capabilities: Capabilities, + capabilities: Capabilities, ) -> anyhow::Result<()> { println!("Source code available at https://github.com/r-a303931/sparse (feel free to give it a star!)"); println!("Type #help to view a list of sparse commands\n"); - let mut stdin = io::stdin(); + let mut stdin = tokio::io::stdin(); let mut stdout = io::stdout(); let backup_term_attrs = get_term_attrs(&stdin)?; let mut raw_term_attrs = get_term_attrs(&stdin)?; @@ -153,6 +130,24 @@ pub(super) async fn shell( let mut cwd = "/".to_string(); + let (stdin_sender, mut stdin_receiver) = channel(64); + + let handle = Handle::current(); + thread::spawn(move || { + let mut stdin_buf = [0u8; 1024]; + loop { + let mut stdin = std::io::stdin(); + let Ok(amount) = stdin.read(&mut stdin_buf) else { + continue; + }; + let stdin_buf = stdin_buf[..amount].to_vec(); + let stdin_sender = stdin_sender.clone(); + handle.spawn(async move { + _ = stdin_sender.send(stdin_buf).await; + }); + } + }); + loop { print!( "{}@{}:{} {} ", @@ -166,14 +161,15 @@ pub(super) async fn shell( ); stdout.flush().unwrap(); - let mut cmd = [0u8; 1024]; - let amount = stdin.read(&mut cmd)?; + let Some(cmd) = stdin_receiver.recv().await else { + break; + }; - if amount == 0 { + if cmd.is_empty() { break; } - let Ok(input) = std::str::from_utf8(&cmd[..amount]) else { + let Ok(input) = std::str::from_utf8(&cmd) else { continue; }; @@ -198,8 +194,28 @@ pub(super) async fn shell( (Ok(SparseCommands::Exit), _) => { break; } + ( + Ok(SparseCommands::UploadFile { + local_file, + remote_path, + }), + _, + ) => { + if let Err(e) = + commands::upload::upload_file(Arc::clone(&connection), local_file, remote_path) + .await + { + eprintln!("{e:?}") + } + } _ => { - if let Err(e) = run_command(input.to_string(), Arc::clone(&connection)).await { + if let Err(e) = run_command( + &mut stdin_receiver, + input.to_string(), + Arc::clone(&connection), + ) + .await + { eprintln!("{e:?}"); }; } diff --git a/sparse-05/sparse-05-common/src/lib.rs b/sparse-05/sparse-05-common/src/lib.rs index c5d1bc6..347becc 100644 --- a/sparse-05/sparse-05-common/src/lib.rs +++ b/sparse-05/sparse-05-common/src/lib.rs @@ -24,10 +24,10 @@ pub mod messages { SendTTYSignal(u64, u64), StartUploadFile(PathBuf, u64), - SendFileSegment(u64, u64, Vec), + SendFileSegment(u64, u64, Vec), StartDownloadFile(PathBuf), - DownloadFileStatus(Result<(), Vec>), + DownloadFileStatus(u64, Result<(), Vec>), Disconnect, } diff --git a/sparse-05/sparse-05-server/src/connection.rs b/sparse-05/sparse-05-server/src/connection.rs index a30d802..c38b11d 100644 --- a/sparse-05/sparse-05-server/src/connection.rs +++ b/sparse-05/sparse-05-server/src/connection.rs @@ -3,7 +3,7 @@ use std::{ net::{Ipv4Addr, UdpSocket}, sync::{ mpsc::{channel, Receiver, Sender}, - Arc, + Arc, Mutex, }, thread, }; @@ -217,6 +217,7 @@ fn authenticate( } mod command; +mod upload_file; fn handle_full_connection( capabilities: Arc, @@ -229,9 +230,9 @@ where { use packets::*; - let mut commands = HashMap::new(); - /*let mut uploaded_files = HashMap::new(); - let mut downloaded_files = HashMap::new();*/ + let commands = Arc::new(Mutex::new(HashMap::new())); + let uploaded_files = Arc::new(Mutex::new(HashMap::new())); + /*let mut downloaded_files = HashMap::new();*/ std::thread::scope(|s| -> anyhow::Result<()> { loop { @@ -251,17 +252,31 @@ where match data { Command::RunCommand(comm) => { - let handler = match command::spawn_command(&s, comm, conninfo.clone()) { - Ok(handler) => handler, - Err(e) => { - eprintln!("error spawning command: {e:?}"); - continue; - } + let commands_clone = commands.clone(); + let Ok(mut lock) = commands.lock() else { + continue; }; - commands.insert(handler.id, handler); + let handler = + match command::spawn_command(&s, comm, conninfo.clone(), commands_clone) { + Ok(handler) => handler, + Err(e) => { + eprintln!("error spawning command: {e:?}"); + continue; + } + }; + + lock.insert(handler.id, handler); + } + Command::SendStdin(bytes, id) => { + let Ok(lock) = commands.lock() else { + continue; + }; + + if let Some(handler) = lock.get(&id) { + let _ = handler.data_sender.send(bytes); + } } - Command::SendStdin(_, _) => {} Command::Cd(_) => {} Command::Ls(_) => {} @@ -271,10 +286,40 @@ where Command::SendTTYData(_, _) => {} Command::SendTTYSignal(_, _) => {} - Command::StartUploadFile(_, _) => {} - Command::SendFileSegment(_, _, _) => {} + Command::StartUploadFile(path, packet_count) => { + let uploaded_files_clone = uploaded_files.clone(); + let Ok(mut lock) = uploaded_files.lock() else { + continue; + }; + + let handler = match upload_file::start_file_upload( + &s, + path, + packet_count, + conninfo.clone(), + uploaded_files_clone, + ) { + Ok(handler) => handler, + Err(e) => { + eprintln!("error starting file upload: {e:?}"); + continue; + } + }; + + lock.insert(handler.id, handler); + } + Command::SendFileSegment(id, number, bytes) => { + let Ok(lock) = uploaded_files.lock() else { + continue; + }; + + if let Some(handler) = lock.get(&id) { + let _ = handler.data_sender.send((number, bytes)); + } + } + Command::StartDownloadFile(_) => {} - Command::DownloadFileStatus(_) => {} + Command::DownloadFileStatus(_, _) => {} Command::Disconnect => { break; diff --git a/sparse-05/sparse-05-server/src/connection/command.rs b/sparse-05/sparse-05-server/src/connection/command.rs index ac0c19e..7415b2c 100644 --- a/sparse-05/sparse-05-server/src/connection/command.rs +++ b/sparse-05/sparse-05-server/src/connection/command.rs @@ -1,10 +1,11 @@ use std::{ + collections::HashMap, io::{Read, Write}, process::{Command, Stdio}, sync::{ atomic::{AtomicBool, AtomicU64, Ordering}, mpsc::{channel, Sender}, - Arc, + Arc, Mutex, }, thread::{scope, Scope}, }; @@ -13,24 +14,41 @@ use sparse_05_common::messages::Response; use super::ConnectionInformation; -const CURRENT_COMMAND_ID: AtomicU64 = AtomicU64::new(0); +static CURRENT_COMMAND_ID: AtomicU64 = AtomicU64::new(0); pub(super) struct CommandHandler { pub id: u64, - data_sender: Sender>, + pub data_sender: Sender>, } pub(super) fn spawn_command<'a, 'b: 'a>( s: &'a Scope<'a, 'b>, command: String, conninfo: ConnectionInformation, + command_map: Arc>>, ) -> anyhow::Result { let (data_sender, data_receiver) = channel(); - let id = CURRENT_COMMAND_ID.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + let id = CURRENT_COMMAND_ID.fetch_add(1, Ordering::Relaxed); - dbg!(&command); + /*let mut command = command.split(" "); + let bin_name = command.next(); + + let Some(bin_name) = bin_name else { + let resp1 = Response::AckRunCommand(id); + let resp1 = conninfo.encrypt_and_sign_resp(resp1)?; + conninfo.send(resp1)?; + let resp2 = Response::CommandDone(id, -1); + let resp2 = conninfo.encrypt_and_sign_resp(resp2)?; + conninfo.send(resp2)?; + + bail!("could not get binary name from command"); + }; + let bin_name = bin_name.trim(); + + let mut command = Command::new(bin_name) + .args(&command.collect::>())*/ let mut command = Command::new("sh") .arg("-c") .arg(&command) @@ -154,6 +172,12 @@ pub(super) fn spawn_command<'a, 'b: 'a>( if let Some(thread) = stdin_thread { _ = thread.join(); } + + let Ok(mut lock) = command_map.lock() else { + return Ok(()); + }; + lock.remove(&id); + Ok(()) }) }); diff --git a/sparse-05/sparse-05-server/src/connection/upload_file.rs b/sparse-05/sparse-05-server/src/connection/upload_file.rs new file mode 100644 index 0000000..5fcfb24 --- /dev/null +++ b/sparse-05/sparse-05-server/src/connection/upload_file.rs @@ -0,0 +1,33 @@ +use std::{ + collections::HashMap, + path::PathBuf, + sync::{ + atomic::{AtomicU64, Ordering}, + mpsc::{channel, Sender}, + Arc, Mutex, + }, + thread::Scope, +}; + +use super::ConnectionInformation; + +static CURRENT_FILE_UPLOAD_ID: AtomicU64 = AtomicU64::new(0); + +pub(super) struct UploadFileHandler { + pub id: u64, + pub data_sender: Sender<(u64, Vec)>, +} + +pub(super) fn start_file_upload<'a, 'b: 'a>( + s: &'a Scope<'a, 'b>, + file_path: PathBuf, + packet_count: u64, + conninfo: ConnectionInformation, + upload_file_map: Arc>>, +) -> anyhow::Result { + let (data_sender, data_receiver) = channel(); + + let id = CURRENT_FILE_UPLOAD_ID.fetch_add(1, Ordering::Relaxed); + + Ok(UploadFileHandler { id, data_sender }) +} diff --git a/sparse-05/sparse-05-server/src/interface.rs b/sparse-05/sparse-05-server/src/interface.rs index f285019..d2615cd 100644 --- a/sparse-05/sparse-05-server/src/interface.rs +++ b/sparse-05/sparse-05-server/src/interface.rs @@ -1,6 +1,7 @@ use std::{ net::{Ipv4Addr, SocketAddrV4, UdpSocket}, sync::Arc, + thread, }; use anyhow::{anyhow, bail}; @@ -25,6 +26,8 @@ impl Interface { .ok_or(anyhow!("could not get an ethernet interface"))?; let interface = loop { + thread::sleep(std::time::Duration::from_millis(250)); + macro_rules! retry { ($e:expr) => {{ match $e { @@ -147,7 +150,9 @@ impl InterfaceReceiver { let mut buf = [0u8; 2000]; - let Ok((count, from)) = interf.recv_from(&mut buf) else { continue; }; + let Ok((count, from)) = interf.recv_from(&mut buf) else { + continue; + }; let udp_packet = UDPPacket::construct(from.port(), *port, &buf[..count]); let ip_packet = IPv4Packet::construct(