From e0e2082341028d1f87acb8d723c2162afe6fb745 Mon Sep 17 00:00:00 2001 From: William Brawner Date: Wed, 27 Jul 2022 21:42:08 -0600 Subject: [PATCH] Move commands to separate file --- src/command.rs | 79 ++++++++++++++++++++++++++++++++++ src/main.rs | 112 +++++++------------------------------------------ 2 files changed, 94 insertions(+), 97 deletions(-) create mode 100644 src/command.rs diff --git a/src/command.rs b/src/command.rs new file mode 100644 index 0000000..8046574 --- /dev/null +++ b/src/command.rs @@ -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 { + // 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> { + 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 = Vec::new(); + let mut paths: Vec = 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, +{ + String::from_utf8(Command::new(command).args(args).output().unwrap().stdout).unwrap() +} + +pub fn telnet(stream: &mut TcpStream, host: &str, port: &u16) -> Result { + 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")) +} \ No newline at end of file diff --git a/src/main.rs b/src/main.rs index 094a53c..ac48cfb 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,18 +1,9 @@ -use std::env; -use std::error::Error; -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 std::io::{Read, Write}; +use std::net::{TcpListener, TcpStream}; use gsm_agent::ThreadPool; -use telnet::Event; -use telnet::Telnet; + +mod command; fn main() { 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 { - // 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> { - 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 = Vec::new(); - let mut paths: Vec = 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, -{ - String::from_utf8(Command::new(command).args(args).output().unwrap().stdout).unwrap() -} - -fn telnet(stream: &mut TcpStream, host: &str, port: &u16) -> Result { - 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) { loop { let mut buffer = [0; 1024]; @@ -118,13 +33,16 @@ fn handle_connection(mut stream: TcpStream) { match command_iter.next() { Some(command) => { let response: String = match command { - "cat" => cat(command_iter.next().unwrap()), - "cd" => match cd(&request[3..]) { + "cat" => command::cat(command_iter.next().unwrap()), + "cd" => match command::cd(&request[3..]) { Ok(s) => s.to_string(), Err(e) => e.to_string(), }, "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(), Err(e) => (*e).to_string(), } @@ -132,17 +50,17 @@ fn handle_connection(mut stream: TcpStream) { "ls" => { let path = match command_iter.next() { Some(s) => s.to_string(), - None => pwd(), + None => command::pwd(), }; - ls(path.as_ref()) + command::ls(path.as_ref()) } "ping" => String::from("pong"), - "pwd" => pwd(), - "shell" => shell(command_iter.next().unwrap(), command_iter), + "pwd" => command::pwd(), + "shell" => command::shell(command_iter.next().unwrap(), command_iter), "telnet" => { let host = command_iter.next().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, Err(e) => e.to_string(), }