chore: use new repl, catch more errors
This commit is contained in:
parent
66d1d7a1da
commit
1fbcd17441
2 changed files with 32 additions and 49 deletions
|
@ -1,9 +1,8 @@
|
||||||
use mlua::{Lua, MultiValue};
|
use mlua::Lua;
|
||||||
use tokio::{sync::{mpsc, broadcast}, net::{TcpListener, TcpStream}, io::{AsyncWriteExt, AsyncReadExt}};
|
use tokio::{sync::{mpsc, broadcast}, net::{TcpListener, TcpStream}, io::{AsyncWriteExt, AsyncReadExt}};
|
||||||
use tracing::{debug, error};
|
use tracing::{debug, error, warn};
|
||||||
|
|
||||||
use crate::{runtime::register_builtin_fn, helpers::pretty_lua};
|
|
||||||
|
|
||||||
|
use crate::runtime::{register_builtin_fn, VERSIONTEXT, repl::LuaRepl};
|
||||||
|
|
||||||
pub struct ControlChannel {
|
pub struct ControlChannel {
|
||||||
addr: String,
|
addr: String,
|
||||||
|
@ -50,57 +49,33 @@ impl ControlChannel {
|
||||||
|
|
||||||
async fn process(&mut self, mut stream: TcpStream) {
|
async fn process(&mut self, mut stream: TcpStream) {
|
||||||
let mut lua = Lua::new();
|
let mut lua = Lua::new();
|
||||||
if let Err(e) = register_builtin_fn(&mut lua, self.source.clone()) {
|
let mut repl = LuaRepl::new(self.source.clone().into());
|
||||||
error!("could not prepare Lua runtime: {}", e);
|
|
||||||
return;
|
let intro_text = format!(
|
||||||
|
"{} inside process #{}\n@> ",
|
||||||
|
VERSIONTEXT, std::process::id()
|
||||||
|
);
|
||||||
|
|
||||||
|
if let Err(e) = repl.write(intro_text) {
|
||||||
|
warn!("could not display version on repl: {}", e);
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Err(e) = register_builtin_fn(&mut lua, self.source.clone()) {
|
||||||
|
error!("could not prepare runtime environment: {}", e);
|
||||||
}
|
}
|
||||||
|
|
||||||
self.source.send(
|
|
||||||
format!("LuaJit 5.2 via rlua inside process #{}\n@> ", std::process::id())
|
|
||||||
).unwrap();
|
|
||||||
let mut cmd = String::new();
|
|
||||||
loop {
|
loop {
|
||||||
tokio::select! {
|
tokio::select! {
|
||||||
|
|
||||||
rx = stream.read_u8() => match rx { // FIXME is read_exact cancelable?
|
rx = stream.read_u8() => match rx { // TODO should be cancelable, but is it really?
|
||||||
Ok(c) => {
|
Ok(c) => {
|
||||||
// TODO move this "lua repl" code outside of here!
|
|
||||||
if !c.is_ascii() {
|
if !c.is_ascii() {
|
||||||
debug!("character '{}' is not ascii", c);
|
debug!("character '{}' is not ascii", c);
|
||||||
break;
|
break;
|
||||||
}
|
}
|
||||||
let ch : char = c as char;
|
let ch : char = c as char;
|
||||||
match ch {
|
if let Err(e) = repl.evaluate(&lua, ch) {
|
||||||
'\u{8}' => {
|
error!("could not evaluate input '{}' : {}", repl.buffer(), e);
|
||||||
if cmd.len() > 0 {
|
|
||||||
cmd.remove(cmd.len() - 1);
|
|
||||||
}
|
|
||||||
},
|
|
||||||
'\n' => {
|
|
||||||
match lua.load(&cmd).eval::<MultiValue>() {
|
|
||||||
Ok(values) => {
|
|
||||||
for val in values {
|
|
||||||
self.source.send(format!("=({}) {}", val.type_name(), pretty_lua(val))).unwrap();
|
|
||||||
}
|
|
||||||
self.source.send("\n@> ".into()).unwrap();
|
|
||||||
cmd = String::new();
|
|
||||||
},
|
|
||||||
Err(e) => {
|
|
||||||
match e {
|
|
||||||
mlua::Error::SyntaxError { message: _, incomplete_input: true } => {
|
|
||||||
self.source.send("@ ".into()).unwrap();
|
|
||||||
cmd.push(ch);
|
|
||||||
},
|
|
||||||
_ => {
|
|
||||||
self.source.send(format!("! {}\n@> ", e)).unwrap();
|
|
||||||
cmd = String::new();
|
|
||||||
},
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
},
|
|
||||||
'\0' => break,
|
|
||||||
_ => cmd.push(ch),
|
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
|
@ -110,12 +85,18 @@ impl ControlChannel {
|
||||||
},
|
},
|
||||||
|
|
||||||
tx = self.sink.recv() => match tx {
|
tx = self.sink.recv() => match tx {
|
||||||
Some(txt) => stream.write_all(txt.as_bytes()).await.unwrap(),
|
Some(txt) => {
|
||||||
|
if let Err(e) = stream.write_all(txt.as_bytes()).await {
|
||||||
|
error!("could not send output to remote console: {}", e);
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
None => {
|
None => {
|
||||||
error!("command sink closed, exiting processor");
|
error!("command sink closed, exiting processor");
|
||||||
break;
|
break;
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
10
src/lib.rs
10
src/lib.rs
|
@ -7,8 +7,7 @@ use tracing::{error, debug};
|
||||||
|
|
||||||
#[ctor::ctor]
|
#[ctor::ctor]
|
||||||
fn contructor() {
|
fn contructor() {
|
||||||
eprint!(" -[infected]- ");
|
std::thread::spawn(move || -> Result<(), std::io::Error> {
|
||||||
std::thread::spawn(move || {
|
|
||||||
tracing_subscriber::fmt()
|
tracing_subscriber::fmt()
|
||||||
.with_max_level(tracing::Level::DEBUG)
|
.with_max_level(tracing::Level::DEBUG)
|
||||||
.with_writer(std::io::stderr)
|
.with_writer(std::io::stderr)
|
||||||
|
@ -16,15 +15,18 @@ fn contructor() {
|
||||||
debug!("infected process");
|
debug!("infected process");
|
||||||
tokio::runtime::Builder::new_current_thread()
|
tokio::runtime::Builder::new_current_thread()
|
||||||
.enable_all()
|
.enable_all()
|
||||||
.build()
|
.build()?
|
||||||
.unwrap()
|
|
||||||
.block_on(main());
|
.block_on(main());
|
||||||
|
Ok(())
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
#[ctor::dtor]
|
#[ctor::dtor]
|
||||||
fn destructor() {}
|
fn destructor() {}
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
async fn main() {
|
async fn main() {
|
||||||
let mut handle = ControlChannel::run("127.0.0.1:13337".into());
|
let mut handle = ControlChannel::run("127.0.0.1:13337".into());
|
||||||
|
|
||||||
|
|
Loading…
Reference in a new issue