fix: fixed stdin for processes

started work on upload file command
This commit is contained in:
Andrew Rioux 2023-09-06 00:07:15 -04:00
parent 9bb31ee6fa
commit 17e6056a03
Signed by: andrew.rioux
GPG Key ID: 9B8BAC47C17ABB94
11 changed files with 236 additions and 84 deletions

View File

@ -6,13 +6,29 @@ use crate::commands::connect::shell::SparseCommands;
pub fn print_help(arg: Option<SparseCommands>) {
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<SparseCommands>) {
\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(""),

View File

@ -1,2 +1,4 @@
pub mod download;
pub mod help;
pub mod sysinfo;
pub mod upload;

View File

@ -0,0 +1,11 @@
use std::{path::PathBuf, sync::Arc};
use crate::commands::connect::Connection;
pub async fn upload_file(
conn: Arc<Connection>,
local_path: PathBuf,
remote_path: PathBuf,
) -> anyhow::Result<()> {
Ok(())
}

View File

@ -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<Response> {
pub fn decrypt_and_verify(&self, data: &[u8]) -> anyhow::Result<Response> {
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<Vec<u8>> {
pub fn encrypt_and_sign(&self, data: &[u8]) -> anyhow::Result<Vec<u8>> {
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<Response> {
pub async fn get_response(&self) -> anyhow::Result<Response> {
let mut buffer = [0u8; 2000];
let (read, from) = self.socket.recv_from(&mut buffer).await?;

View File

@ -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<Connection>) -> anyhow::Result<()> {
async fn run_command(
stdin: &mut Receiver<Vec<u8>>,
command: String,
connection: Arc<Connection>,
) -> anyhow::Result<()> {
connection
.send_command(Command::RunCommand(command))
.await?;
@ -70,38 +81,6 @@ async fn run_command(command: String, connection: Arc<Connection>) -> 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<u8>),
@ -110,7 +89,7 @@ async fn run_command(command: String, connection: Arc<Connection>) -> 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<Connection>) -> 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<Connection>) -> anyhow::Re
}
}
let _ = kill.send(());
let _ = stdin_thread.join();
Ok(())
}
pub(super) async fn shell(
connection: Arc<Connection>,
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:?}");
};
}

View File

@ -24,10 +24,10 @@ pub mod messages {
SendTTYSignal(u64, u64),
StartUploadFile(PathBuf, u64),
SendFileSegment(u64, u64, Vec<u64>),
SendFileSegment(u64, u64, Vec<u8>),
StartDownloadFile(PathBuf),
DownloadFileStatus(Result<(), Vec<u64>>),
DownloadFileStatus(u64, Result<(), Vec<u64>>),
Disconnect,
}

View File

@ -3,7 +3,7 @@ use std::{
net::{Ipv4Addr, UdpSocket},
sync::{
mpsc::{channel, Receiver, Sender},
Arc,
Arc, Mutex,
},
thread,
};
@ -217,6 +217,7 @@ fn authenticate<F: Fn()>(
}
mod command;
mod upload_file;
fn handle_full_connection<F>(
capabilities: Arc<Capabilities>,
@ -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,7 +252,13 @@ where
match data {
Command::RunCommand(comm) => {
let handler = match command::spawn_command(&s, comm, conninfo.clone()) {
let commands_clone = commands.clone();
let Ok(mut lock) = commands.lock() else {
continue;
};
let handler =
match command::spawn_command(&s, comm, conninfo.clone(), commands_clone) {
Ok(handler) => handler,
Err(e) => {
eprintln!("error spawning command: {e:?}");
@ -259,9 +266,17 @@ where
}
};
commands.insert(handler.id, handler);
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;

View File

@ -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<Vec<u8>>,
pub data_sender: Sender<Vec<u8>>,
}
pub(super) fn spawn_command<'a, 'b: 'a>(
s: &'a Scope<'a, 'b>,
command: String,
conninfo: ConnectionInformation,
command_map: Arc<Mutex<HashMap<u64, CommandHandler>>>,
) -> anyhow::Result<CommandHandler> {
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::<Vec<_>>())*/
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(())
})
});

View File

@ -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<u8>)>,
}
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<Mutex<HashMap<u64, UploadFileHandler>>>,
) -> anyhow::Result<UploadFileHandler> {
let (data_sender, data_receiver) = channel();
let id = CURRENT_FILE_UPLOAD_ID.fetch_add(1, Ordering::Relaxed);
Ok(UploadFileHandler { id, data_sender })
}

View File

@ -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(