Move commands to separate file

This commit is contained in:
William Brawner 2022-07-27 21:42:08 -06:00
parent 92428e8429
commit e0e2082341
2 changed files with 94 additions and 97 deletions

79
src/command.rs Normal file
View file

@ -0,0 +1,79 @@
use std::{fs::{self, DirEntry}, path::Path, env, error::Error, process::Command, net::TcpStream, time::Duration, io::{Read, Write}};
use telnet::{Telnet, Event};
pub fn cat(path: &str) -> String {
let path = Path::new(path);
String::from_utf8(fs::read(path).unwrap()).unwrap()
}
pub fn cd(destination: &str) -> Result<String, std::io::Error> {
// TODO: Change path per-worker instead of globally
let path = Path::new(destination);
match env::set_current_dir(path) {
Ok(_) => Ok(pwd()),
Err(e) => Err(e),
}
}
pub fn get(url: &str, destination: &str) -> Result<(), Box<dyn Error>> {
let response = minreq::get(url).send()?;
let path = Path::new(destination);
match fs::write(path, response.as_bytes()) {
Ok(_) => Ok(()),
Err(e) => Err(Box::new(e)),
}
}
pub fn ls(path: &str) -> String {
let mut files: Vec<String> = Vec::new();
let mut paths: Vec<DirEntry> = fs::read_dir(path).unwrap().map(|r| r.unwrap()).collect();
paths.sort_by(|a, b| a.file_name().cmp(&b.file_name()));
for file in paths {
files.push(file.file_name().into_string().unwrap());
}
files.join("\n")
}
pub fn pwd() -> String {
env::current_dir().unwrap().to_string_lossy().to_string()
}
pub fn shell<'a, T>(command: &str, args: T) -> String
where
T: Iterator<Item = &'a str>,
{
String::from_utf8(Command::new(command).args(args).output().unwrap().stdout).unwrap()
}
pub fn telnet(stream: &mut TcpStream, host: &str, port: &u16) -> Result<String, std::io::Error> {
let mut telnet = Telnet::connect((host, *port), 4096)?;
let mut stream_buffer = [0; 4096];
stream.set_read_timeout(Some(Duration::new(1, 0)))?;
'main: loop {
match stream.read(&mut stream_buffer) {
Ok(read) => {
match String::from_utf8_lossy(&stream_buffer[0..read])
.as_ref()
.trim()
{
"quit" => break 'main,
_ => telnet.write(&stream_buffer[0..read]).unwrap(),
};
}
Err(_) => {}
};
match telnet.read_timeout(Duration::new(1, 0)) {
Ok(event) => {
if let Event::Data(buffer) = event {
stream.write(&buffer).unwrap();
}
}
Err(e) => {
println!("telnet read error: {:?}\n", e);
break 'main;
}
}
}
Ok(String::from("Telnet connection closed"))
}

View file

@ -1,18 +1,9 @@
use std::env; use std::io::{Read, Write};
use std::error::Error; use std::net::{TcpListener, TcpStream};
use std::fs;
use std::fs::DirEntry;
use std::io::Read;
use std::io::Write;
use std::net::TcpListener;
use std::net::TcpStream;
use std::path::Path;
use std::process::Command;
use std::time::Duration;
use gsm_agent::ThreadPool; use gsm_agent::ThreadPool;
use telnet::Event;
use telnet::Telnet; mod command;
fn main() { fn main() {
let listener = TcpListener::bind("0.0.0.0:4762").unwrap(); let listener = TcpListener::bind("0.0.0.0:4762").unwrap();
@ -27,82 +18,6 @@ fn main() {
} }
} }
fn cat(path: &str) -> String {
let path = Path::new(path);
String::from_utf8(fs::read(path).unwrap()).unwrap()
}
fn cd(destination: &str) -> Result<String, std::io::Error> {
// TODO: Change path per-worker instead of globally
let path = Path::new(destination);
match env::set_current_dir(path) {
Ok(_) => Ok(pwd()),
Err(e) => Err(e),
}
}
fn get(url: &str, destination: &str) -> Result<(), Box<dyn Error>> {
let response = minreq::get(url).send()?;
let path = Path::new(destination);
match fs::write(path, response.as_bytes()) {
Ok(_) => Ok(()),
Err(e) => Err(Box::new(e)),
}
}
fn ls(path: &str) -> String {
let mut files: Vec<String> = Vec::new();
let mut paths: Vec<DirEntry> = fs::read_dir(path).unwrap().map(|r| r.unwrap()).collect();
paths.sort_by(|a, b| a.file_name().cmp(&b.file_name()));
for file in paths {
files.push(file.file_name().into_string().unwrap());
}
files.join("\n")
}
fn pwd() -> String {
env::current_dir().unwrap().to_string_lossy().to_string()
}
fn shell<'a, T>(command: &str, args: T) -> String
where
T: Iterator<Item = &'a str>,
{
String::from_utf8(Command::new(command).args(args).output().unwrap().stdout).unwrap()
}
fn telnet(stream: &mut TcpStream, host: &str, port: &u16) -> Result<String, std::io::Error> {
let mut telnet = Telnet::connect((host, *port), 4096)?;
let mut stream_buffer = [0; 4096];
stream.set_read_timeout(Some(Duration::new(1, 0)))?;
'main: loop {
match stream.read(&mut stream_buffer) {
Ok(read) => {
match String::from_utf8_lossy(&stream_buffer[0..read])
.as_ref()
.trim()
{
"quit" => break 'main,
_ => telnet.write(&stream_buffer[0..read]).unwrap(),
};
}
Err(_) => {}
};
match telnet.read_timeout(Duration::new(1, 0)) {
Ok(event) => {
if let Event::Data(buffer) = event {
stream.write(&buffer).unwrap();
}
}
Err(e) => {
println!("telnet read error: {:?}\n", e);
break 'main;
}
}
}
Ok(String::from("Telnet connection closed"))
}
fn handle_connection(mut stream: TcpStream) { fn handle_connection(mut stream: TcpStream) {
loop { loop {
let mut buffer = [0; 1024]; let mut buffer = [0; 1024];
@ -118,13 +33,16 @@ fn handle_connection(mut stream: TcpStream) {
match command_iter.next() { match command_iter.next() {
Some(command) => { Some(command) => {
let response: String = match command { let response: String = match command {
"cat" => cat(command_iter.next().unwrap()), "cat" => command::cat(command_iter.next().unwrap()),
"cd" => match cd(&request[3..]) { "cd" => match command::cd(&request[3..]) {
Ok(s) => s.to_string(), Ok(s) => s.to_string(),
Err(e) => e.to_string(), Err(e) => e.to_string(),
}, },
"get" => { "get" => {
match get(command_iter.next().unwrap(), command_iter.next().unwrap()) { match command::get(
command_iter.next().unwrap(),
command_iter.next().unwrap(),
) {
Ok(_) => String::new(), Ok(_) => String::new(),
Err(e) => (*e).to_string(), Err(e) => (*e).to_string(),
} }
@ -132,17 +50,17 @@ fn handle_connection(mut stream: TcpStream) {
"ls" => { "ls" => {
let path = match command_iter.next() { let path = match command_iter.next() {
Some(s) => s.to_string(), Some(s) => s.to_string(),
None => pwd(), None => command::pwd(),
}; };
ls(path.as_ref()) command::ls(path.as_ref())
} }
"ping" => String::from("pong"), "ping" => String::from("pong"),
"pwd" => pwd(), "pwd" => command::pwd(),
"shell" => shell(command_iter.next().unwrap(), command_iter), "shell" => command::shell(command_iter.next().unwrap(), command_iter),
"telnet" => { "telnet" => {
let host = command_iter.next().unwrap(); let host = command_iter.next().unwrap();
let port = u16::from_str_radix(command_iter.next().unwrap(), 10).unwrap(); let port = u16::from_str_radix(command_iter.next().unwrap(), 10).unwrap();
match telnet(&mut stream, host, &port) { match command::telnet(&mut stream, host, &port) {
Ok(s) => s, Ok(s) => s,
Err(e) => e.to_string(), Err(e) => e.to_string(),
} }