pyro/pentest_tool
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
9fa18003375721a3c18218aa176d6b2a559bfbe1
pentest_tool/pentest_tool/src/project_controls.rs
Pyro57000 6d92032d06 added some functions to switch between managing 
personal projects and managing work projects
as well as the ability to separate your current
project list into work and personal.
2025-06-02 12:18:59 -05:00

631 lines
28 KiB
Rust
Raw Blame History

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;
use std::process::Command;
use std::thread;
use std::time::Duration;
use std::str::FromStr;
use colored::Colorize;
use crate::get_user_input;
use crate::open_overwrite;
use crate::tableize;
use crate::Project;
use crate::box_controls::make_box;
use crate::print_success;
use crate::print_error;
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;
print_success(format!("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{
print_error("error selecting new project", String::from("unknown project id selected."));
}
}
}
else{
print_error("error selecting new project", String::from("we need 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 save_file_res = fs::File::create(save_file_path);
if save_file_res.is_err(){
let error = save_file_res.err().unwrap().to_string();
print_error("error creating save file.", error);
return;
}
let mut save_file = save_file_res.unwrap();
save_file.write_all(b"Current Loaded Projects Config File\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, project.stage);
active_set = true;
}
}
else{
_outline = format!("{}no:{}:{}\n", default, project.boxname, project.stage);
}
let write_res = save_file.write_all(_outline.as_bytes());
if write_res.is_err(){
let error = write_res.err().unwrap().to_string();
print_error("error saving projects.conf file!", error);
return;
}
write_res.unwrap();
}
}
#[allow(unused)]
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, upcoming_files: &PathBuf, upcoming_notes: &PathBuf, fingerprint: bool){
let mut new_project_dir = PathBuf::new();
let mut new_note_dir = PathBuf::new();
let mut existing_folders = String::new();
let mut customer_name = String::new();
let mut project_name = String::new();
let mut project_stage = String::new();
loop{
let mut stage_response = String::new();
println!("what stage is this project in?");
print!("
1.) current
2.) upcoming
");
let stage_result = stdin().read_line(&mut stage_response);
if stage_result.is_err(){
println!("we need input here dummy, try again...");
}
else{
match &stage_response.trim_end(){
&"1" => {project_stage = "current".to_owned(); break;},
&"2" => {project_stage = "upcoming".to_owned(); break;},
_ => println!("unknown option, try again...")
}
}
}
if project_stage.contains("current"){
new_project_dir = project_dir.clone();
new_note_dir = notes_dir.clone();
}
else if project_stage.contains("upcoming"){
new_project_dir = upcoming_files.clone();
new_note_dir = upcoming_notes.clone();
}
else{
println!("unknown stage!!")
}
println!("{}", new_project_dir.display());
println!("do you have an existing notes and folder structure to copy over?\ny/n");
std::io::stdin().read_line(&mut existing_folders).unwrap();
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(&notes_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 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,
stage: project_stage.to_owned()
};
if project_stage.contains("current"){
make_box(&new_project, &tools_dir, &boxtemplate, true, fingerprint);
}
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!("no project selected! returning...");
}
}
else{
println!("no project selected! returning...");
}
}
pub fn get_projects(config_path: &PathBuf, show: bool) -> Option<Vec<Project>>{
let mut mut_config_path = config_path.clone();
mut_config_path.pop();
mut_config_path.push("projects.conf");
let mut bkup_config_path = mut_config_path.clone();
bkup_config_path.pop();
bkup_config_path.push("projects.conf.bkup");
let bkup_result = fs::copy(&mut_config_path, &bkup_config_path);
if bkup_result.is_err(){
let error = bkup_result.err().unwrap();
println!("{}","error backing up the projects.conf file!".red());
println!("{}", error.to_string().red());
}
let mut projects = Vec::new();
let projects_string_res = fs::read_to_string(mut_config_path);
if projects_string_res.is_err(){
let error = projects_string_res.err().unwrap();
println!("{}", "Error Loading Project!".red());
println!("{}", error.to_string().red());
return None;
}
let projects_string = projects_string_res.unwrap();
let project_lines:Vec<&str> = projects_string.split("\n").collect();
let mut first = 0;
let mut already_active = false;
for line in project_lines{
//println!("{}", line);
first = first + 1;
if first != 1{
if line.len() > 1{
let settings: Vec<&str> = line.split(":").collect();
if settings.len() > 5{
// debug config file...
/*let mut count = 0;
for settin in &settings{
println!("{}: {}", count, settin);
count = count + 1;
}*/
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(&notes_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 project_stage = settings[6].to_owned();
let new_project = Project{customer: customer, project_name: project, files_folder: project_folder, notes_folder: notes_folder, active: active, id: first, boxname: boxname, stage: project_stage};
if show{
println!("{} {} {}", &new_project.customer, &new_project.project_name, "LOADED!".green());
}
projects.push(new_project);
}
}
}
}
return Some(projects)
}
pub fn print_upcoming_projects(projects: &Vec<Project>){
let mut lines = vec![String::from("CUSTOMER||PROJECT")];
for project in projects{
if project.stage.contains("upcoming"){
let new_line = format!("{}||{}", project.customer, project.project_name);
lines.push(new_line);
}
}
tableize(lines);
}
pub fn promote_project(projects: &mut Vec<Project>, config_path: &PathBuf, project_dir: &PathBuf, notes_dir: &PathBuf, tools_dir: &PathBuf, boxtemplate: &String, fingerprint: bool){
let working_projects = projects.clone();
for project in &working_projects{
if project.stage.contains("upcoming"){
println!("{}.) {}:{}", project.id, project.customer, project.project_name);
}
}
println!("which project to promote?");
let mut selection = String::new();
let result = stdin().read_line(&mut selection);
if result.is_err(){
println!("we need input here dummy try again....");
}
result.unwrap();
println!("{}", project_dir.display());
let promote_id: i32 = selection.trim_end().parse().unwrap();
let mut projects_to_save = Vec::new();
for project in &working_projects{
if project.id == promote_id{
let mut promoted_project = project.clone();
let mut new_files_dir = project_dir.clone();
let mut new_notes_dir = notes_dir.clone();
new_files_dir.push(&promoted_project.customer);
new_notes_dir.push(&promoted_project.customer);
fs::create_dir_all(&new_files_dir).unwrap();
fs::create_dir_all(&new_notes_dir).unwrap();
let folder_move_success = process::Command::new("mv")
.arg("-i")
.arg(&project.files_folder)
.arg(&new_files_dir.display().to_string())
.status().expect("unable to call the system mv command");
let note_move_success = process::Command::new("mv")
.arg("-i")
.arg(&project.notes_folder)
.arg(&new_notes_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!!");
let mut remove_folder = PathBuf::new();
remove_folder.push(&project.files_folder);
remove_folder.pop();
let remove_files_res = fs::remove_dir_all(remove_folder);
if remove_files_res.is_err(){
println!("error removing the original files folder form the upcomming folder, manual cleanup required");
}
else{
remove_files_res.unwrap();
println!("upcoming files folder cleanup successful!");
}
}
else{
println!("failed to copy the project folder, try to move it manually!");
}
if note_move_success.success(){
print_success("we copied the notes folder correctly!!");
let mut remove_folder = PathBuf::new();
remove_folder.push(&project.files_folder);
remove_folder.pop();
let remove_notes_res = fs::remove_dir_all(remove_folder);
if remove_notes_res.is_err(){
println!("error removing the original notes folder form the upcomming folder, manual cleanup required");
}
else{
remove_notes_res.unwrap();
println!("upcoming notes folder cleanup successful!");
}
}
else{
println!("failed to copy the notes folder, try to move it manually!");
}
new_files_dir.push(&promoted_project.project_name);
new_notes_dir.push(&promoted_project.project_name);
promoted_project.files_folder = new_files_dir;
promoted_project.notes_folder = new_notes_dir;
promoted_project.stage = "current".to_owned();
let cs_response = get_user_input("will you need to be using cobalt strike for this project?");
if cs_response.to_lowercase().contains("y"){
let cs_path = get_user_input("path to your current cobalt strike directory?");
let copy_result = Command::new("cp").arg("-R").arg(&cs_path).arg(&promoted_project.files_folder).status();
if copy_result.is_err(){
println!("oof we had an error copying... you'll have to copy this manually");
}
else{
let copy_exit = copy_result.unwrap();
if copy_exit.success() == false{
println!("oof we had an error copying... you'll have to copy this manually");
println!("run cp -R {} {}", &cs_path, &promoted_project.files_folder.display());
}
}
}
thread::sleep(Duration::from_secs(3));
make_box(&promoted_project, tools_dir, boxtemplate, true, fingerprint);
projects_to_save.push(promoted_project);
}
else{
projects_to_save.push(project.clone());
}
}
projects.clear();
projects.append(&mut projects_to_save);
save_projects(&projects_to_save, config_path);
print_success("project promoted successfully!");
}
pub fn list_projects(projects: &Vec<Project>){
let mut lines = vec!["customer||project||stage".to_owned()];
for project in projects{
let new_line = format!("{}||{}||{}", project.customer, project.project_name, project.stage);
lines.push(new_line);
}
tableize(lines);
}
pub fn separate_personal_work_projects(config_path: &PathBuf){
let mut projects_conf_path = config_path.clone();
projects_conf_path.pop();
projects_conf_path.push("projects.conf");
let mut working_conf_path = config_path.clone();
working_conf_path.pop();
working_conf_path.push("projects.work");
let mut personal_conf_path = config_path.clone();
personal_conf_path.pop();
personal_conf_path.push("projects.personal");
println!("{}", projects_conf_path.display());
let project_read_res = read_to_string(&projects_conf_path);
if project_read_res.is_err(){
print_error("error reading current projects config file!", project_read_res.err().unwrap().to_string());
return;
}
let project_string = project_read_res.unwrap();
let project_lines: Vec<&str> = project_string.split("\n").collect();
let mut personal_projects = Vec::new();
let mut work_projects = Vec::new();
let mut default = String::new();
for line in project_lines{
let words: Vec<&str> = line.split(":").collect();
if words.len() > 3{
if words[0].contains("default"){
default = line.to_owned();
}
else{
println!("{} {}", words[0], words[1]);
if get_user_input("should this project be added to your personal projects config files?").to_lowercase().contains("y"){
personal_projects.push(line.to_owned());
print_success(format!("{}:{} {}",words[0], words[1], "successfully added to personal config!"));
if get_user_input("Do you also want to keep it in your work projects config?").to_lowercase().contains("y"){
work_projects.push(line.to_owned());
}
}
else{
work_projects.push(line.to_owned());
}
}
}
}
let work_config_open_res = fs::OpenOptions::new().create(true).write(true).open(working_conf_path);
if work_config_open_res.is_err(){
print_error("error opening work projects config file!", work_config_open_res.err().unwrap().to_string());
return;
}
let person_config_open_res = fs::OpenOptions::new().create(true).write(true).open(personal_conf_path);
if person_config_open_res.is_err(){
print_error("error opening personal config file!", person_config_open_res.err().unwrap().to_string());
return;
}
let mut work_config_file = work_config_open_res.unwrap();
let mut personal_config_file = person_config_open_res.unwrap();
let work_write_success = write!(work_config_file, "Work Config File\n");
let personal_write_success = write!(personal_config_file, "Personal Config File\n");
if work_write_success.is_err(){
print_error("error writing to work config file!", work_write_success.err().unwrap().to_string());
return;
}
else{
work_write_success.unwrap();
}
if personal_write_success.is_err(){
print_error("error writing personal config file!", personal_write_success.err().unwrap().to_string());
return;
}
else{
personal_write_success.unwrap();
}
write!(personal_config_file, "{}\n", default).unwrap();
write!(work_config_file, "{}\n", default).unwrap();
for project in work_projects{
write!(work_config_file, "{}\n", project).unwrap();
}
for project in personal_projects{
write!(personal_config_file, "{}\n", project).unwrap();
}
print_success("projects separated successfully!");
}
pub fn swith_to_personal(config: &PathBuf) -> Option<Vec<Project>>{
let mut projects_path = config.clone();
projects_path.pop();
let mut personal_projects = config.clone();
personal_projects.pop();
projects_path.push("projects.conf");
personal_projects.push("projects.personal");
let personal_projects_read_res = read_to_string(&personal_projects);
if personal_projects_read_res.is_err(){
print_error("error reading personal projects!", personal_projects_read_res.err().unwrap().to_string());
return None;
}
let person_projects_string = personal_projects_read_res.unwrap();
let open_res = open_overwrite(&projects_path);
if open_res.is_none(){
print_error("error opening projects.conf file!", "".to_string());
return None;
}
let mut project_conf = open_res.unwrap();
write!(project_conf, "{}", person_projects_string).unwrap();
let new_projects = get_projects(config, true);
return new_projects;
}
pub fn swith_to_work(config: &PathBuf) -> Option<Vec<Project>>{
let mut projects_path = config.clone();
projects_path.pop();
let mut work_projects = config.clone();
work_projects.pop();
projects_path.push("projects.conf");
work_projects.push("projects.work");
let work_projects_read_res = read_to_string(work_projects);
if work_projects_read_res.is_err(){
print_error("error reading personal projects!", work_projects_read_res.err().unwrap().to_string());
return None;
}
let work_projects_string = work_projects_read_res.unwrap();
let open_res = open_overwrite(&projects_path);
if open_res.is_none(){
print_error("", "error opening projects.conf file!".to_string());
return None;
}
let mut project_conf = open_res.unwrap();
write!(project_conf, "{}", work_projects_string).unwrap();
let new_projects = get_projects(config, true);
return new_projects;
}
Reference in New Issue View Git Blame Copy Permalink