feat: added the ability to send commands

This commit is contained in:
Andrew Rioux
2023-09-05 19:05:26 -04:00
parent 47b2191335
commit 52538ac7d7
5 changed files with 188 additions and 23 deletions

View File

@@ -33,7 +33,7 @@ impl Connection {
let signature: [u8; 64] = data[..64].try_into().unwrap();
self.foreign_sign_pubkey
.verify(data, &Signature::from(signature))?;
.verify(&data[64..], &Signature::from(signature))?;
let data = ecies_ed25519::decrypt(&self.config.enc_privkey, &data[64..])?;
@@ -67,7 +67,7 @@ impl Connection {
bail!("received packet from wrong computer");
}
self.decrypt_and_verify(&buffer[..read])
dbg!(self.decrypt_and_verify(&buffer[..read]))
}
}

View File

@@ -1,12 +1,21 @@
use std::{
io::{self, Read, Write},
io::{self, stdin, Read, Stdin, Write},
os::fd::AsRawFd,
path::PathBuf,
sync::Arc,
sync::{
mpsc::{channel, TryRecvError},
Arc,
},
thread::{self, scope},
};
use sparse_05_common::messages::Capabilities;
use sparse_05_common::messages::{Capabilities, Command, Response};
use structopt::StructOpt;
use tokio::{
io::{stderr, stdout, AsyncWriteExt},
runtime::Handle,
sync::mpsc,
};
use super::{commands, Connection};
@@ -49,7 +58,83 @@ fn convert_termios_raw(attrs: &mut libc::termios) -> anyhow::Result<()> {
Ok(())
}
async fn run_command(connection: Arc<Connection>) -> anyhow::Result<()> {
async fn run_command(command: String, connection: Arc<Connection>) -> anyhow::Result<()> {
connection
.send_command(Command::RunCommand(command))
.await?;
let id = loop {
let resp = connection.get_response().await?;
if let Response::AckRunCommand(id) = resp {
break id;
}
};
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>),
Remote(Response),
}
let Some(event) = tokio::select! {
v = connection.get_response() => v.ok().map(Event::Remote),
v = handle_send_message.recv() => v.map(Event::Stdin)
} else {
continue;
};
match event {
Event::Remote(Response::SendStdout(bytes, cid)) if cid == id => {
stdout().write(&bytes).await?;
}
Event::Remote(Response::SendStderr(bytes, cid)) if cid == id => {
stderr().write(&bytes).await?;
}
Event::Remote(Response::CommandDone(cid, code)) if cid == id => break,
Event::Stdin(stdin) => {
let _ = connection
.send_command(Command::SendStdin(stdin, id))
.await?;
}
_ => {}
}
}
let _ = kill.send(());
let _ = stdin_thread.join();
Ok(())
}
@@ -81,7 +166,7 @@ pub(super) async fn shell(
);
stdout.flush().unwrap();
let mut cmd = [0u8; 256];
let mut cmd = [0u8; 1024];
let amount = stdin.read(&mut cmd)?;
if amount == 0 {
@@ -113,7 +198,11 @@ pub(super) async fn shell(
(Ok(SparseCommands::Exit), _) => {
break;
}
_ => {}
_ => {
if let Err(e) = run_command(input.to_string(), Arc::clone(&connection)).await {
eprintln!("{e:?}");
};
}
}
}