Refactored the code to be split
into different modules instead of one giant .rs file this should make it more organized
This commit is contained in:
pyro57000
339
pentest_tool/src/project_controls.rs
Normal file
339
pentest_tool/src/project_controls.rs
Normal file
@@ -0,0 +1,339 @@
|
||||
use std::env;
|
||||
use std::fs;
|
||||
use std::io::stdin;
|
||||
use std::io::Write;
|
||||
use std::path::PathBuf;
|
||||
use std::process;
|
||||
use std::thread;
|
||||
use std::time::Duration;
|
||||
use std::str::FromStr;
|
||||
use crate::Project;
|
||||
|
||||
pub fn switch_project(projects: &mut Vec<Project>){
|
||||
for project in projects.clone(){
|
||||
if project.active == false{
|
||||
println!("{} {}|{}", project.id, project.customer, project.project_name);
|
||||
}
|
||||
}
|
||||
println!("\nnew project selection?\n");
|
||||
let mut response = String::new();
|
||||
std::io::stdin().read_line(&mut response).unwrap();
|
||||
if response.len() > 1{
|
||||
let new_id:i32 = response.trim_end().parse().expect("error converting to i32");
|
||||
for project in projects{
|
||||
if project.id == new_id{
|
||||
project.active = true;
|
||||
println!("project found switching to {} {}", project.customer, project.project_name);
|
||||
env::set_var("CURRENT_PROJECT_BOX", project.boxname.clone());
|
||||
}
|
||||
else if project.id != new_id{
|
||||
project.active = false;
|
||||
}
|
||||
else{
|
||||
println!("error unknown project id")
|
||||
}
|
||||
}
|
||||
}
|
||||
else{
|
||||
println!("error we need user input here dummy!");
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
pub fn save_projects(projects: &Vec<Project>, config_path: &PathBuf){
|
||||
let mut save_file_path = config_path.clone();
|
||||
let mut active_set = false;
|
||||
save_file_path.pop();
|
||||
save_file_path.push("projects.conf");
|
||||
let mut save_file = fs::File::create(save_file_path).expect("error creating save_file");
|
||||
save_file.write_all(b"customer:name:notes:files:active:time:box_name\n").expect("error writing first line to file");
|
||||
for project in projects{
|
||||
let default = format!{"{}:{}:{}:{}:", project.customer, project.project_name, project.notes_folder.display(), project.files_folder.display()};
|
||||
let mut _outline = String::new();
|
||||
if project.active{
|
||||
if active_set == false{
|
||||
_outline = format!("{}yes:{}\n", default, project.boxname);
|
||||
active_set = true;
|
||||
}
|
||||
|
||||
}
|
||||
else{
|
||||
_outline = format!("{}no:{}\n", default, project.boxname);
|
||||
}
|
||||
save_file.write_all(_outline.as_bytes()).expect("error writing outline");
|
||||
}
|
||||
}
|
||||
|
||||
pub fn new_project(projects: &mut Vec<Project>, project_dir: &PathBuf, notes_dir: &PathBuf, tools_dir: &PathBuf, boxtemplate: &String, config_path: &PathBuf, new_id: i32){
|
||||
let mut new_project_dir = project_dir.clone();
|
||||
let mut new_note_dir = notes_dir.clone();
|
||||
let mut existing_folders = String::new();
|
||||
let mut customer_name = String::new();
|
||||
let mut project_name = String::new();
|
||||
println!("do you have an existing notes and folder structure to copy over?\ny/n");
|
||||
std::io::stdin().read_line(&mut existing_folders).unwrap();
|
||||
let mut customer_name = customer_name.trim_end().to_owned();
|
||||
let mut project_name = project_name.trim_end().to_owned();
|
||||
if existing_folders.contains("y") || existing_folders.contains("Y"){
|
||||
println!("NOTE THIS WILL OVERWRITE CUSTOMER NAME AND PROJECT NAME WITH THE 2nd TO LAST AND LAST FOLDER NAMES IN THE PATH YOU'RE COPYING RESPECTIVELY");
|
||||
let mut files_to_copy = String::new();
|
||||
let mut notes_to_copy = String::new();
|
||||
println!("path to project folder folder to copy:");
|
||||
std::io::stdin().read_line(&mut files_to_copy).unwrap();
|
||||
println!("path to notes folder to copy:");
|
||||
std::io::stdin().read_line(&mut notes_to_copy).unwrap();
|
||||
//to get rid of the new line chars
|
||||
files_to_copy.pop();
|
||||
notes_to_copy.pop();
|
||||
println!("files to copy: {}", files_to_copy);
|
||||
println!("notes to copy: {}", notes_to_copy);
|
||||
println!("files destination: {}", new_project_dir.display());
|
||||
println!("notes destination: {}", new_note_dir.display());
|
||||
let folder_move_success = process::Command::new("mv")
|
||||
.arg("-i")
|
||||
.arg(&files_to_copy)
|
||||
.arg(new_project_dir.display().to_string())
|
||||
.status().expect("unable to call the system mv command");
|
||||
let note_move_success = process::Command::new("mv")
|
||||
.arg("-i")
|
||||
.arg(¬es_to_copy)
|
||||
.arg(new_note_dir.display().to_string())
|
||||
.status().expect("unable to call the system mv command");
|
||||
if folder_move_success.success(){
|
||||
println!("we copied the project folder correctly!!");
|
||||
}
|
||||
else{
|
||||
println!("failed to copy the project folder, try to move it manually!");
|
||||
}
|
||||
if note_move_success.success(){
|
||||
println!("we copied the notes folder correctly!!");
|
||||
}
|
||||
else{
|
||||
println!("failed to copy the notes folder, try to move it manually!");
|
||||
}
|
||||
// this lovely set of code takes the folder you gave to copy and tries to find the customername and project name based on directory names
|
||||
// this solves a case where the entered customer name or project name does not match the files copied
|
||||
let copied_files = PathBuf::from(&files_to_copy);
|
||||
customer_name = copied_files.file_name().unwrap().to_str().unwrap().to_owned();
|
||||
new_project_dir.push(&customer_name);
|
||||
new_note_dir.push(&customer_name);
|
||||
match fs::read_dir(&new_project_dir){
|
||||
Ok(entries) => {
|
||||
for entry in entries{
|
||||
match entry {
|
||||
Ok(entry) => {
|
||||
let entry_path = entry.path();
|
||||
if entry_path.is_dir(){
|
||||
if let Some(dir_name) = entry_path.file_name(){
|
||||
if let Some(dir_name_str) = dir_name.to_str(){
|
||||
if dir_name_str.contains("pentest"){
|
||||
project_name = dir_name_str.to_owned();
|
||||
println!("pentest folder found! assuming projectname...");
|
||||
new_project_dir.push(&project_name);
|
||||
new_note_dir.push(&project_name);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => eprintln!("error reading entry name: {}", e)
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => eprintln!("Error reading directory entries: {}", e)
|
||||
}
|
||||
let mut customer_name_response = String::new();
|
||||
let mut project_name_response = String::new();
|
||||
println!("Customer_name: {}", &customer_name);
|
||||
println!("Is this correct?");
|
||||
stdin().read_line(&mut customer_name_response).unwrap();
|
||||
println!("Project_name: {}", &project_name);
|
||||
println!("Is this corrrect?");
|
||||
stdin().read_line(&mut project_name_response).unwrap();
|
||||
if customer_name_response.contains(|c: char| c == 'n' || c=='N') || project_name_response.contains(|c: char| c == 'n' || c=='N'){
|
||||
println!("oops sorry about that, tried to guess based on the second to last and last folder names");
|
||||
if customer_name_response.contains(|c: char| c == 'n' || c =='N'){
|
||||
let mut name_response = String::new();
|
||||
println!("what is the correct customer name?");
|
||||
stdin().read_line(&mut name_response).unwrap();
|
||||
name_response.pop();
|
||||
customer_name = name_response.to_owned();
|
||||
}
|
||||
if project_name_response.contains(|c: char| c == 'n' || c == 'N'){
|
||||
let mut project_response = String::new();
|
||||
println!("what is the correct project name?");
|
||||
stdin().read_line(&mut project_response).unwrap();
|
||||
project_response.pop();
|
||||
project_name = project_response.to_owned();
|
||||
}
|
||||
}
|
||||
}
|
||||
else{
|
||||
println!("customer name?");
|
||||
std::io::stdin().read_line(&mut customer_name).unwrap();
|
||||
println!("project name?");
|
||||
std::io::stdin().read_line(&mut project_name).unwrap();
|
||||
// to remove newline characters
|
||||
customer_name.pop();
|
||||
project_name.pop();
|
||||
new_project_dir.push(&customer_name);
|
||||
new_note_dir.push(&customer_name);
|
||||
new_project_dir.push(&project_name);
|
||||
new_note_dir.push(&project_name);
|
||||
fs::create_dir_all(&new_project_dir).expect("error creating new files folder");
|
||||
fs::create_dir_all(&new_note_dir).expect("error creating new notes folder");
|
||||
}
|
||||
thread::sleep(Duration::from_secs(2));
|
||||
let box_name = format!("atarchbox_{}", customer_name);
|
||||
let mut box_name_path = new_project_dir.clone();
|
||||
box_name_path.push("boxname");
|
||||
let mut box_name_file = fs::File::create(box_name_path).expect("Error creating box name file");
|
||||
box_name_file.write_all(&box_name.as_bytes()).expect("error writing boxname to box file");
|
||||
let pentest_volume = format!("{}:/pentest:rw", new_project_dir.display());
|
||||
let toold_volume = format!("{}:/tools:rw", tools_dir.display());
|
||||
println!("distrobox create --root --clone {} --volume {} --volume {} --name {}", boxtemplate, toold_volume, pentest_volume, box_name);
|
||||
let distrobox_result = process::Command::new("distrobox")
|
||||
.arg("create")
|
||||
.arg("--root")
|
||||
.arg("--clone")
|
||||
.arg(boxtemplate)
|
||||
.arg("--volume")
|
||||
.arg(&toold_volume)
|
||||
.arg("--volume")
|
||||
.arg(&pentest_volume)
|
||||
.arg("--name")
|
||||
.arg(&box_name)
|
||||
.status()
|
||||
.expect("error getting distrobox status");
|
||||
if distrobox_result.success(){
|
||||
println!("we made a distrobox oh boy!");
|
||||
let distrobox_start_result = process::Command::new("distrobox")
|
||||
.arg("enter")
|
||||
.arg("--root")
|
||||
.arg(&box_name)
|
||||
.arg("--")
|
||||
.arg("sudo")
|
||||
.arg("-s")
|
||||
.arg("ln")
|
||||
.arg("-sf")
|
||||
.arg("/pentest/boxname")
|
||||
.arg("/etc/boxname")
|
||||
.status()
|
||||
.expect("error getting response from distrobox start");
|
||||
if distrobox_start_result.success(){
|
||||
println!("distrobox was started as well!!!! good job me!");
|
||||
}
|
||||
else{
|
||||
println!("ooof did not start successfully try entering it yoruself");
|
||||
println!("distrobox enter --rrot {} -- sudo -s ln -sf /pentest/boxname /etc/boxname", &box_name);
|
||||
}
|
||||
}
|
||||
else{
|
||||
println!("ooof distrobox did not work.... try creating it yourself");
|
||||
println!("distrobox create --root --clone {} --volume {} --volume {} --name {}", boxtemplate, &toold_volume, &pentest_volume, &box_name);
|
||||
println!("distrobox enter --rrot {} -- sudo -s ln -sf /pentest/boxname /etc/boxname", &box_name);
|
||||
}
|
||||
let new_project = Project{customer: customer_name.trim_end().to_owned(),
|
||||
project_name: project_name.trim_end().to_owned(),
|
||||
notes_folder: new_note_dir,
|
||||
files_folder:new_project_dir,
|
||||
active: false,
|
||||
id: new_id,
|
||||
boxname: box_name,
|
||||
};
|
||||
projects.push(new_project);
|
||||
save_projects(projects, config_path);
|
||||
|
||||
}
|
||||
|
||||
|
||||
pub fn remove_project(projects: &mut Vec<Project>, config_path: &PathBuf){
|
||||
for project in projects.clone(){
|
||||
println!("{} {} {}", project.id, project.customer, project.project_name);
|
||||
}
|
||||
let mut project_to_remove = String::new();
|
||||
println!("project to remove?");
|
||||
std::io::stdin().read_line(&mut project_to_remove).unwrap();
|
||||
if project_to_remove.len() > 1{
|
||||
let mut project_to_keep = Vec::new();
|
||||
if project_to_remove.len() > 0{
|
||||
let remove_id: i32 = project_to_remove.trim_end().parse().unwrap();
|
||||
let mut project_set = false;
|
||||
for project in projects.clone(){
|
||||
if project.id == remove_id{
|
||||
println!("will remove {} {}", project.customer, project.project_name);
|
||||
project_set = true;
|
||||
let _distrobox_stop_status = process::Command::new("distrobox").arg("stop").arg("--root").arg(&project.boxname).status().expect("error stopping distrobox");
|
||||
let distrobox_rm_status = process::Command::new("distrobox-rm")
|
||||
.arg("--root")
|
||||
.arg("-f")
|
||||
.arg(&project.boxname)
|
||||
.status().expect("error calling distrobox");
|
||||
if distrobox_rm_status.success(){
|
||||
println!("Distrobox Removal Successful!!!");
|
||||
}
|
||||
else{
|
||||
println!("Distrobox Removal Failed, manual removal required!");
|
||||
}
|
||||
}
|
||||
else {
|
||||
println!("{} {} will be kept", project.customer, project.project_name);
|
||||
project_to_keep.push(project);
|
||||
}
|
||||
}
|
||||
if project_set{
|
||||
projects.clear();
|
||||
projects.append(&mut project_to_keep);
|
||||
save_projects(&projects, config_path);
|
||||
}
|
||||
else{
|
||||
println!("error no prjects found to remove")
|
||||
}
|
||||
|
||||
}
|
||||
else{
|
||||
println!("we need user in put here dummy!!");
|
||||
}
|
||||
}
|
||||
else{
|
||||
println!("we need input here dummy!");
|
||||
}
|
||||
}
|
||||
|
||||
pub fn get_projects(config_path: &PathBuf) -> Vec<Project>{
|
||||
let mut mut_config_path = config_path.clone();
|
||||
mut_config_path.pop();
|
||||
mut_config_path.push("projects.conf");
|
||||
let mut projects = Vec::new();
|
||||
let projects_string = fs::read_to_string(mut_config_path).expect("error reading projects file");
|
||||
let project_lines:Vec<&str> = projects_string.split("\n").collect();
|
||||
let mut first = 0;
|
||||
let mut already_active = false;
|
||||
for line in project_lines{
|
||||
first = first + 1;
|
||||
if first != 1{
|
||||
if line.len() > 1{
|
||||
let settings: Vec<&str> = line.split(":").collect();
|
||||
let customer = settings[0].to_owned();
|
||||
let project = settings[1].to_owned();
|
||||
let notes_string = settings[2].to_owned();
|
||||
let folder_string = settings[3].to_owned();
|
||||
let notes_folder = PathBuf::from_str(¬es_string.trim_end()).expect("error reading notes string");
|
||||
let project_folder = PathBuf::from_str(&folder_string.trim_end()).expect("error reading folding sering");
|
||||
let mut active = false;
|
||||
let boxname = settings[5].to_owned();
|
||||
if settings[4] == "yes"{
|
||||
if already_active == false{
|
||||
env::set_var("CURRENT_PROJECT_BOX", boxname.clone());
|
||||
already_active = true;
|
||||
active = true;
|
||||
}
|
||||
}
|
||||
let new_project = Project{customer: customer, project_name: project, files_folder: project_folder, notes_folder: notes_folder, active: active, id: first, boxname: boxname};
|
||||
println!("{} {} LOADED!", &new_project.customer, &new_project.project_name);
|
||||
projects.push(new_project);
|
||||
}
|
||||
}
|
||||
}
|
||||
return projects
|
||||
}
|
||||
Reference in New Issue
Block a user