main.py

from wsl_manager import WSLManager
import sys

class bcolors:
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKCYAN = '\033[96m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'

def select_distribution(manager, prompt="Select distribution"):
    success, stdout, stderr = manager.run_command("wsl -l -q", show_output=False)
    if not success and "no installed distributions" not in stdout.lower():
        print(f"{bcolors.FAIL}Error listing distributions{bcolors.ENDC}")
        return None
    
    lines = stdout.strip().split('\n')
    distro_names = []
    
    for line in lines:
        cleaned_line = line.strip().replace('\x00', '').replace('\r', '')
        if cleaned_line and not cleaned_line.lower().startswith('windows subsystem'):
            distro_names.append(cleaned_line)
    
    if not distro_names:
        print(f"{bcolors.WARNING}No WSL distributions found{bcolors.ENDC}")
        return None
    
    print(f"\n{bcolors.OKCYAN}{prompt}:{bcolors.ENDC}")
    for i, distro in enumerate(distro_names, 1):
        print(f"{i}. {distro}")
    
    try:
        choice = int(input(f"\nSelect distribution (1-{len(distro_names)}): ")) - 1
        if 0 <= choice < len(distro_names):
            return distro_names[choice]
        else:
            print(f"{bcolors.FAIL}Invalid selection{bcolors.ENDC}")
            return None
    except:
        print(f"{bcolors.FAIL}Invalid input{bcolors.ENDC}")
        return None

def print_menu(manager):
    print(f"\n{bcolors.HEADER}=== WSL Pentest Environment Manager ==={bcolors.ENDC}")
    
    if manager.current_config:
        config_name = manager.current_config['template_name']
        distro = manager.current_config['distro']
        print(f"{bcolors.OKGREEN}Current Config: {config_name} ({distro}){bcolors.ENDC}")
        print("─" * 50)
        print("1. Manage configurations (create/select/delete)")
        print("2. Create template")
        print("3. Clone template")
        print("4. Setup & Install packages")
        print("5. Manage distributions (list/start/delete)")
        print("6. Current config info")
        print("7. Exit")
    else:
        print(f"{bcolors.WARNING}No configuration selected{bcolors.ENDC}")
        print("─" * 50)
        print("1. Manage configurations (create/select/delete)")
        print("2. Manage distributions (list/delete)")
        print("3. Exit")
        print(f"\n{bcolors.OKCYAN}Select a configuration first to access template features{bcolors.ENDC}")

def main():
    manager = WSLManager()
    
    while True:
        print_menu(manager)
        choice = input("\nSelect option: ").strip()
        
        if choice == '1':
            while True:
                print(f"\n{bcolors.HEADER}=== Manage Configurations ==={bcolors.ENDC}")
                print("1. Select existing configuration")
                print("2. Create new configuration")
                print("3. Delete configuration")
                print("4. Check for outdated configurations")
                print("5. Back to main menu")
                sub_choice = input("Choose option (1-5): ").strip()
                
                if sub_choice == '1':
                    manager.select_config()
                elif sub_choice == '2':
                    manager.create_new_template_config()
                elif sub_choice == '3':
                    manager.delete_config()
                elif sub_choice == '4':
                    manager.check_outdated_configs()
                elif sub_choice == '5':
                    break
                else:
                    print(f"{bcolors.FAIL}Invalid option{bcolors.ENDC}")
        
        elif not manager.current_config:
            if choice == '2':
                while True:
                    print(f"\n{bcolors.HEADER}=== Manage Distributions ==={bcolors.ENDC}")
                    print("1. List distributions")
                    print("2. Delete distribution")
                    print("3. Back to main menu")
                    sub_choice = input("Choose option (1-3): ").strip()
                    
                    if sub_choice == '1':
                        print(f"\n{bcolors.OKCYAN}WSL Distributions:{bcolors.ENDC}")
                        print(manager.list_distributions())
                    elif sub_choice == '2':
                        distro = select_distribution(manager, "Select distribution to delete")
                        if distro:
                            manager.delete_distribution(distro)
                    elif sub_choice == '3':
                        break
                    else:
                        print(f"{bcolors.FAIL}Invalid option{bcolors.ENDC}")
            elif choice == '3':
                print(f"{bcolors.OKGREEN}Goodbye!{bcolors.ENDC}")
                sys.exit(0)
            else:
                print(f"{bcolors.FAIL}Please select a configuration first for template features{bcolors.ENDC}")
                continue
        
        elif choice == '2':
            manager.create_template()
        
        elif choice == '3':
            source_distro = select_distribution(manager, "Select distribution to clone")
            if source_distro:
                name = input(f"Enter new instance name (or press Enter for '{source_distro}-clone'): ").strip()
                if not name:
                    name = f"{source_distro}-clone"
                    print(f"{bcolors.OKCYAN}Using automatic name: {name}{bcolors.ENDC}")
                manager.clone_distribution(source_distro, name)
        
        elif choice == '4':
            distro = select_distribution(manager, "Select distribution for setup & package installation")
            if distro:
                if manager.current_config:
                    print("\nPackage installation options:")
                    print("1. Install from current config (auto)")
                    print("2. Universal packages only")
                    print("3. Custom packages")
                    
                    pkg_choice = input("Select option (1-3): ").strip()
                    
                    if pkg_choice == '1':
                        manager.setup_and_install_packages(distro, "auto")
                    elif pkg_choice == '2':
                        manager.setup_and_install_packages(distro, "universal")
                    elif pkg_choice == '3':
                        manager.setup_and_install_packages(distro, "custom")
                    else:
                        print("Invalid option")
                else:
                    print("\nSelect package type:")
                    print("1. Universal packages")
                    print("2. Custom packages")
                    
                    pkg_choice = input("Select option (1-2): ").strip()
                    
                    if pkg_choice == '1':
                        manager.setup_and_install_packages(distro, "universal")
                    elif pkg_choice == '2':
                        manager.setup_and_install_packages(distro, "custom")
                    else:
                        print("Invalid option")
        
        elif choice == '5':
            while True:
                print(f"\n{bcolors.HEADER}=== Manage Distributions ==={bcolors.ENDC}")
                print("1. List distributions")
                print("2. Start distribution")
                print("3. Delete distribution")
                print("4. Back to main menu")
                sub_choice = input("Choose option (1-4): ").strip()
                
                if sub_choice == '1':
                    print(f"\n{bcolors.OKCYAN}WSL Distributions:{bcolors.ENDC}")
                    print(manager.list_distributions())
                elif sub_choice == '2':
                    distro = select_distribution(manager, "Select distribution to start")
                    if distro:
                        manager.start_distribution(distro)
                elif sub_choice == '3':
                    distro = select_distribution(manager, "Select distribution to delete")
                    if distro:
                        manager.delete_distribution(distro)
                elif sub_choice == '4':
                    break
                else:
                    print(f"{bcolors.FAIL}Invalid option{bcolors.ENDC}")
        
        elif choice == '6':
            info = manager.get_current_config_info()
            print(f"\n{bcolors.OKCYAN}Current Configuration:{bcolors.ENDC}")
            for key, value in info.items():
                print(f"{key.replace('_', ' ').title()}: {value}")
        
        elif choice == '7':
            print(f"{bcolors.OKGREEN}Goodbye!{bcolors.ENDC}")
            sys.exit(0)
        
        else:
            print(f"{bcolors.FAIL}Invalid option{bcolors.ENDC}")

if __name__ == "__main__":
    main()