pyro/pentest_tool
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update main.rs

Browse Source 
re-worked the add a project and remove project functions to be better.

# import_project
added logic to automatically set customer and project names based on moved folder names, falls back to user input, but the project and notes folders should be properly set for imported projects anyways.

# remove_project
added logic to save the list of projects after the distrobox removal happens, this way your projects.conf file should be updated after removing the project.
This commit is contained in:
Pyro57000
2024-12-03 13:35:32 -06:00
committed by GitHub
parent 08921e0515
commit 2ad293af7c
1 changed files with 71 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

84
pentest_tool/src/main.rs
View File

@@ -2,6 +2,7 @@ use std::collections::HashMap;
use std::env;
use std::fs;
use std::fs::read_to_string;
use std::io::stdin;
use std::io::Write;
use std::path::PathBuf;
use std::process;
@@ -234,21 +235,19 @@ fn new_project(projects: &mut Vec<Project>, project_dir: &PathBuf, notes_dir: &P
let mut existing_folders = String::new();
let mut customer_name = String::new();
let mut project_name = String::new();
println!("customer name?");
std::io::stdin().read_line(&mut customer_name).unwrap();
println!("project name?");
std::io::stdin().read_line(&mut project_name).unwrap();
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 customer_name = customer_name.trim_end().to_owned();
let project_name = project_name.trim_end().to_owned();
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);
@@ -257,12 +256,12 @@ fn new_project(projects: &mut Vec<Project>, project_dir: &PathBuf, notes_dir: &P
println!("notes destination: {}", new_note_dir.display());
let folder_move_success = process::Command::new("mv")
.arg("-i")
.arg(files_to_copy)
.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(notes_to_copy)
.arg(&notes_to_copy)
.arg(new_note_dir.display().to_string())
.status().expect("unable to call the system mv command");
if folder_move_success.success(){
@@ -277,13 +276,71 @@ fn new_project(projects: &mut Vec<Project>, project_dir: &PathBuf, notes_dir: &P
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);
new_project_dir.push(&project_name);
new_note_dir.push(&project_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);
@@ -355,7 +412,7 @@ fn new_project(projects: &mut Vec<Project>, project_dir: &PathBuf, notes_dir: &P
}
fn remove_project(projects: &mut Vec<Project>){
fn remove_project(projects: &mut Vec<Project>, config_path: &PathBuf){
for project in projects.clone(){
println!("{} {} {}", project.id, project.customer, project.project_name);
}
@@ -392,6 +449,7 @@ fn remove_project(projects: &mut Vec<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")
@@ -716,7 +774,7 @@ Current Project: {} {}
"4" => start_pentest(),
"5" => save_projects(&projects, &config_path),
"6" => {new_id = new_id + 1; new_project(&mut projects, &base_files, &base_notes, &tools_dir, &boxtemplate, &config_path, new_id)},
"7" => remove_project(&mut projects),
"7" => remove_project(&mut projects, &config_path),
"8" => project_standalone_terminal(active_project.clone(), terminal.clone()),
"9" => project_inline_terminal(active_project.clone()),
"10" => open_in_dolphin("files", active_project.clone()),