souleyez 2.16.0__py3-none-any.whl → 2.26.0__py3-none-any.whl

souleyez/ui/setup_wizard.py

@@ -7,7 +7,9 @@ import time
 import click
 import getpass
 import shutil
+import subprocess
 from pathlib import Path
+from typing import List, Dict, Optional
 from souleyez.ui.design_system import DesignSystem
 
 
@@ -26,16 +28,336 @@ def mark_wizard_completed():
     WIZARD_STATE_FILE.touch()
 
 
+def _check_disk_space(min_mb: int = 500) -> bool:
+    """Check if there's enough disk space for installs."""
+    try:
+        usage = shutil.disk_usage('/')
+        free_mb = usage.free // (1024 * 1024)
+        return free_mb >= min_mb
+    except Exception:
+        return True  # Assume OK if we can't check
+
+
+def _configure_sudoers(tool_name: str, tool_path: str) -> bool:
+    """
+    Add NOPASSWD sudoers entry for privileged tool.
+
+    Args:
+        tool_name: Name of the tool (used for sudoers filename)
+        tool_path: Absolute path to tool binary
+
+    Returns:
+        True if configured successfully, False otherwise
+    """
+    # Skip if running as root
+    if os.geteuid() == 0:
+        return True
+
+    username = getpass.getuser()
+    sudoers_file = f"/etc/sudoers.d/{tool_name}"
+    sudoers_line = f"{username} ALL=(ALL) NOPASSWD: {tool_path}"
+    tmp_file = f"/tmp/sudoers_{tool_name}_{os.getpid()}"
+
+    try:
+        # Write to temp file first
+        with open(tmp_file, 'w') as f:
+            f.write(sudoers_line + '\n')
+
+        # Validate syntax with visudo before applying
+        result = subprocess.run(
+            ['sudo', 'visudo', '-c', '-f', tmp_file],
+            capture_output=True,
+            timeout=30
+        )
+
+        if result.returncode != 0:
+            click.echo(f"    {click.style('!', fg='yellow')} Invalid sudoers syntax for {tool_name}, skipping")
+            os.unlink(tmp_file)
+            return False
+
+        # Safe to move to sudoers.d
+        subprocess.run(['sudo', 'mv', tmp_file, sudoers_file], check=True, timeout=30)
+        subprocess.run(['sudo', 'chmod', '0440', sudoers_file], check=True, timeout=30)
+
+        return True
+
+    except subprocess.TimeoutExpired:
+        click.echo(f"    {click.style('!', fg='yellow')} Timeout configuring sudoers for {tool_name}")
+        return False
+    except Exception as e:
+        click.echo(f"    {click.style('!', fg='yellow')} Error configuring sudoers for {tool_name}: {e}")
+        # Clean up temp file if it exists
+        if os.path.exists(tmp_file):
+            try:
+                os.unlink(tmp_file)
+            except Exception:
+                pass
+        return False
+
+
+def _install_desktop_shortcut():
+    """
+    Install desktop shortcut for SoulEyez in Applications menu.
+
+    This runs silently during setup - any errors are ignored to not
+    disrupt the setup flow.
+    """
+    try:
+        applications_dir = Path.home() / '.local' / 'share' / 'applications'
+        icons_dir = Path.home() / '.local' / 'share' / 'icons'
+        desktop_file = applications_dir / 'souleyez.desktop'
+        icon_dest = icons_dir / 'souleyez.png'
+
+        # Skip if already installed
+        if desktop_file.exists():
+            return
+
+        # Create directories
+        applications_dir.mkdir(parents=True, exist_ok=True)
+        icons_dir.mkdir(parents=True, exist_ok=True)
+
+        # Find and copy icon
+        try:
+            # Try importlib.resources first (Python 3.9+)
+            try:
+                from importlib.resources import files
+                icon_source = files('souleyez.assets').joinpath('souleyez-icon.png')
+                with open(icon_source, 'rb') as src:
+                    icon_data = src.read()
+            except (ImportError, TypeError, FileNotFoundError):
+                # Fallback: find icon relative to this file
+                icon_source = Path(__file__).parent.parent / 'assets' / 'souleyez-icon.png'
+                with open(icon_source, 'rb') as src:
+                    icon_data = src.read()
+
+            with open(icon_dest, 'wb') as dst:
+                dst.write(icon_data)
+        except Exception:
+            icon_dest = "utilities-terminal"  # Fallback to system icon
+
+        # Create .desktop file
+        desktop_content = f"""[Desktop Entry]
+Name=SoulEyez
+Comment=AI-Powered Penetration Testing Platform
+Exec=souleyez interactive
+Icon={icon_dest}
+Terminal=true
+Type=Application
+Categories=Security;System;Network;
+Keywords=pentest;security;hacking;nmap;metasploit;
+"""
+
+        desktop_file.write_text(desktop_content)
+
+        # Update desktop database (optional)
+        try:
+            subprocess.run(['update-desktop-database', str(applications_dir)],
+                          capture_output=True, check=False, timeout=5)
+        except Exception:
+            pass
+
+        click.echo(f"  {click.style('✓', fg='green')} Desktop shortcut: Added to Applications menu")
+
+    except Exception:
+        # Silently ignore errors - desktop shortcut is nice-to-have
+        pass
+
+
+def _run_tool_installs(
+    missing_tools: List[Dict],
+    wrong_version_tools: List[Dict],
+    distro: str
+) -> bool:
+    """
+    Prompt user and run install/upgrade commands.
+
+    Args:
+        missing_tools: List of missing tool dicts with 'name', 'install', 'tool_info'
+        wrong_version_tools: List of wrong version tool dicts
+        distro: Detected distribution
+
+    Returns:
+        True if any installs were run
+    """
+    from souleyez.utils.tool_checker import EXTERNAL_TOOLS, get_upgrade_command
+
+    all_installs = []
+
+    # Add wrong version tools (upgrades)
+    for tool in wrong_version_tools:
+        # Try to get upgrade command, fall back to install command
+        tool_info = tool.get('tool_info', {})
+        upgrade_cmd = get_upgrade_command(tool_info, distro) if tool_info else None
+        install_cmd = upgrade_cmd or tool.get('install', '')
+
+        all_installs.append({
+            'name': tool['name'],
+            'cmd': install_cmd,
+            'action': 'upgrade',
+            'tool_info': tool_info
+        })
+
+    # Add missing tools (installs)
+    for tool in missing_tools:
+        all_installs.append({
+            'name': tool['name'],
+            'cmd': tool['install'],
+            'action': 'install',
+            'tool_info': tool.get('tool_info', {})
+        })
+
+    if not all_installs:
+        return False
+
+    # Check disk space
+    if not _check_disk_space(500):
+        click.echo()
+        click.echo(f"  {click.style('!', fg='yellow')} Low disk space (<500MB free). Installs may fail.")
+        if not click.confirm("  Continue anyway?", default=False):
+            return False
+
+    # Show what will be installed/upgraded
+    click.echo()
+    click.echo(f"  {len(all_installs)} tool(s) to install/upgrade:")
+    for item in all_installs:
+        action_color = 'cyan' if item['action'] == 'upgrade' else 'green'
+        click.echo(f"    - {item['name']} ({click.style(item['action'], fg=action_color)})")
+
+    click.echo()
+    if not click.confirm("  Install/upgrade now?", default=False):
+        click.echo("  Skipped. You can run 'souleyez setup' later to install tools.")
+        return False
+
+    # Request sudo upfront
+    click.echo()
+    click.echo("  Requesting sudo access...")
+    result = subprocess.run(['sudo', '-v'], check=False)
+    if result.returncode != 0:
+        click.echo(f"  {click.style('x', fg='red')} Could not obtain sudo. Aborting installs.")
+        return False
+
+    # Update apt cache first (for apt-based installs)
+    has_apt_installs = any('apt' in item['cmd'] for item in all_installs)
+    if has_apt_installs:
+        click.echo()
+        click.echo("  Updating package lists...")
+        result = subprocess.run(['sudo', 'apt', 'update'], capture_output=True)
+        if result.returncode != 0:
+            click.echo(f"  {click.style('!', fg='yellow')} apt update had issues, continuing anyway...")
+
+    # Track results for summary
+    results = []
+
+    # Run each install command
+    for item in all_installs:
+        click.echo()
+        click.echo(f"  {click.style(item['action'].capitalize() + 'ing', fg='cyan')} {item['name']}...")
+
+        try:
+            # Run install command
+            proc = subprocess.run(
+                item['cmd'],
+                shell=True,  # nosec B602 - commands from trusted EXTERNAL_TOOLS
+                timeout=600  # 10 minute timeout per tool
+            )
+
+            if proc.returncode == 0:
+                click.echo(f"  {click.style('✓', fg='green')} {item['name']} {item['action']} complete")
+                results.append({'name': item['name'], 'success': True, 'tool_info': item['tool_info']})
+            else:
+                click.echo(f"  {click.style('✗', fg='red')} {item['name']} failed (exit {proc.returncode})")
+                results.append({'name': item['name'], 'success': False, 'tool_info': item['tool_info']})
+
+        except subprocess.TimeoutExpired:
+            click.echo(f"  {click.style('✗', fg='red')} {item['name']} timed out")
+            results.append({'name': item['name'], 'success': False, 'tool_info': item['tool_info']})
+        except Exception as e:
+            click.echo(f"  {click.style('✗', fg='red')} {item['name']} error: {e}")
+            results.append({'name': item['name'], 'success': False, 'tool_info': item['tool_info']})
+
+    # Configure sudoers for privileged tools that installed successfully
+    click.echo()
+    click.echo("  Configuring permissions for privileged tools...")
+
+    sudoers_configured = False
+    for res in results:
+        if not res['success']:
+            continue
+
+        tool_info = res['tool_info']
+        if not tool_info.get('needs_sudo'):
+            continue
+
+        sudoers_configured = True
+        # Find the actual binary path
+        command = tool_info.get('command', res['name'])
+        tool_path = shutil.which(command)
+
+        # Check alt_commands if primary not found
+        if not tool_path and tool_info.get('alt_commands'):
+            for alt in tool_info['alt_commands']:
+                tool_path = shutil.which(alt)
+                if tool_path:
+                    break
+
+        if tool_path:
+            if _configure_sudoers(res['name'], tool_path):
+                click.echo(f"    {click.style('✓', fg='green')} {res['name']} configured for passwordless sudo")
+            else:
+                click.echo(f"    {click.style('!', fg='yellow')} {res['name']} sudoers config failed")
+        else:
+            click.echo(f"    {click.style('!', fg='yellow')} {res['name']} binary not found, skipping sudoers")
+
+    if not sudoers_configured:
+        click.echo("    No privileged tools needed configuration.")
+
+    # Re-verify installed tools
+    click.echo()
+    click.echo("  Verifying installations...")
+
+    success_count = 0
+    fail_count = 0
+
+    for res in results:
+        tool_info = res['tool_info']
+        command = tool_info.get('command', res['name'])
+        alt_commands = tool_info.get('alt_commands')
+
+        # Check if tool is now available
+        found = shutil.which(command) is not None
+        if not found and alt_commands:
+            for alt in alt_commands:
+                if shutil.which(alt):
+                    found = True
+                    break
+
+        if found:
+            click.echo(f"    {click.style('✓', fg='green')} {res['name']} verified")
+            success_count += 1
+        else:
+            click.echo(f"    {click.style('✗', fg='red')} {res['name']} not found after install")
+            fail_count += 1
+
+    # Summary
+    click.echo()
+    if fail_count == 0:
+        click.echo(f"  {click.style('✓', fg='green')} All {success_count} tool(s) installed successfully!")
+    else:
+        click.echo(f"  {click.style('!', fg='yellow')} {success_count} succeeded, {fail_count} failed")
+        click.echo("  Failed tools may need manual installation.")
+
+    return True
+
 
 def _show_wizard_banner():
     """Display the SoulEyez ASCII banner for wizard steps."""
     click.echo()
-    click.echo(click.style("   ███████╗ ██████╗ ██╗   ██╗██╗     ███████╗██╗   ██╗███████╗███████╗", fg='bright_cyan', bold=True))
-    click.echo(click.style("   ██╔════╝██╔═══██╗██║   ██║██║     ██╔════╝╚██╗ ██╔╝██╔════╝╚══███╔╝", fg='bright_cyan', bold=True))
-    click.echo(click.style("   ███████╗██║   ██║██║   ██║██║     █████╗   ╚████╔╝ █████╗    ███╔╝ ", fg='bright_cyan', bold=True))
-    click.echo(click.style("   ╚════██║██║   ██║██║   ██║██║     ██╔══╝    ╚██╔╝  ██╔══╝   ███╔╝  ", fg='bright_cyan', bold=True))
-    click.echo(click.style("   ███████║╚██████╔╝╚██████╔╝███████╗███████╗   ██║   ███████╗███████╗", fg='bright_cyan', bold=True))
-    click.echo(click.style("   ╚══════╝ ╚═════╝  ╚═════╝ ╚══════╝╚══════╝   ╚═╝   ╚══════╝╚══════╝", fg='bright_cyan', bold=True))
+    click.echo(click.style("   ███████╗ ██████╗ ██╗   ██╗██╗     ███████╗██╗   ██╗███████╗███████╗", fg='bright_cyan', bold=True) + click.style("        ▄██▄", fg='bright_blue', bold=True))
+    click.echo(click.style("   ██╔════╝██╔═══██╗██║   ██║██║     ██╔════╝╚██╗ ██╔╝██╔════╝╚══███╔╝", fg='bright_cyan', bold=True) + click.style("      ▄█▀  ▀█▄", fg='bright_blue', bold=True))
+    click.echo(click.style("   ███████╗██║   ██║██║   ██║██║     █████╗   ╚████╔╝ █████╗    ███╔╝ ", fg='bright_cyan', bold=True) + click.style("     █   ◉   █", fg='bright_blue', bold=True))
+    click.echo(click.style("   ╚════██║██║   ██║██║   ██║██║     ██╔══╝    ╚██╔╝  ██╔══╝   ███╔╝  ", fg='bright_cyan', bold=True) + click.style("     █  ═══  █", fg='bright_blue', bold=True))
+    click.echo(click.style("   ███████║╚██████╔╝╚██████╔╝███████╗███████╗   ██║   ███████╗███████╗", fg='bright_cyan', bold=True) + click.style("      ▀█▄  ▄█▀", fg='bright_blue', bold=True))
+    click.echo(click.style("   ╚══════╝ ╚═════╝  ╚═════╝ ╚══════╝╚══════╝   ╚═╝   ╚══════╝╚══════╝", fg='bright_cyan', bold=True) + click.style("        ▀██▀", fg='bright_blue', bold=True))
     click.echo()
     click.echo(click.style("  Created by CyberSoul SecurITy", fg='bright_blue'))
     click.echo()
@@ -44,7 +366,7 @@ def _show_wizard_banner():
 def run_setup_wizard() -> bool:
     """
     Run the setup wizard for new users.
-    
+
     Returns:
         bool: True if wizard completed, False if skipped/cancelled
     """
@@ -53,20 +375,20 @@ def run_setup_wizard() -> bool:
         from souleyez.auth import get_current_user, Tier
         user = get_current_user()
         is_pro = user and user.tier == Tier.PRO
-        
+
         # Step 1: Welcome
         if not _wizard_welcome(is_pro):
             mark_wizard_completed()  # Mark as completed even if skipped
             return False
-        
+
         # Step 2: Encryption Setup
         encryption_enabled = _wizard_encryption_setup()
-        
+
         # Step 3: Create First Engagement
         engagement_info = _wizard_create_engagement()
         if not engagement_info:
             return False
-        
+
         # Step 4: Tool Check
         tool_status = _wizard_tool_check()
 
@@ -78,7 +400,7 @@ def run_setup_wizard() -> bool:
             automation_prefs = _wizard_automation_prefs()
         else:
             automation_prefs = {'enabled': False, 'mode': None}
-        
+
         # Step 7: Deliverable Templates (Pro only)
         if is_pro:
             templates = _wizard_deliverables(engagement_info.get('type'))
@@ -95,10 +417,10 @@ def run_setup_wizard() -> bool:
             templates,
             is_pro
         )
-        
+
         mark_wizard_completed()
         return True
-        
+
     except (KeyboardInterrupt, click.Abort):
         click.echo(click.style("\n\n  Setup wizard cancelled.", fg='yellow'))
         mark_wizard_completed()  # Don't show wizard again
@@ -110,13 +432,13 @@ def _wizard_welcome(is_pro: bool = False) -> bool:
     """Show welcome screen with ASCII banner."""
     DesignSystem.clear_screen()
     width = DesignSystem.get_terminal_width()
-    
+
     # Header
     click.echo("\n┌" + "─" * (width - 2) + "┐")
     click.echo("│" + click.style(" WELCOME TO SOULEYEZ - SETUP WIZARD ".center(width - 2), bold=True, fg='cyan') + "│")
     click.echo("└" + "─" * (width - 2) + "┘")
     click.echo()
-    
+
     # ASCII Art Banner - SOULEYEZ with all-seeing eye on the right
     click.echo(click.style("   ███████╗ ██████╗ ██╗   ██╗██╗     ███████╗██╗   ██╗███████╗███████╗", fg='bright_cyan', bold=True) + click.style("        ▄██▄", fg='bright_blue', bold=True))
     click.echo(click.style("   ██╔════╝██╔═══██╗██║   ██║██║     ██╔════╝╚██╗ ██╔╝██╔════╝╚══███╔╝", fg='bright_cyan', bold=True) + click.style("      ▄█▀  ▀█▄", fg='bright_blue', bold=True))
@@ -166,15 +488,15 @@ def _wizard_encryption_setup() -> bool:
     click.echo("│" + click.style(" STEP 2: ENCRYPTION SETUP ".center(width - 2), bold=True, fg='cyan') + "│")
     click.echo("└" + "─" * (width - 2) + "┘")
     click.echo()
-    
+
     crypto = CryptoManager()
-    
+
     if crypto.is_encryption_enabled():
         click.echo("  " + click.style("✓ Encryption already configured", fg='green'))
         click.echo()
         click.pause("  Press any key to continue...")
         return True
-    
+
     click.echo("  SoulEyez encrypts all credentials (passwords, API keys, tokens)")
     click.echo("  with a master password. This is " + click.style("required", fg='green', bold=True) + " for security.")
     click.echo()
@@ -204,46 +526,46 @@ def _wizard_encryption_setup() -> bool:
     click.echo("    • At least one number")
     click.echo("    • At least one special character (!@#$%^&*)")
     click.echo()
-    click.echo("  " + click.style("⚠️  If you lose this password, encrypted credentials cannot be recovered!", 
+    click.echo("  " + click.style("⚠️  If you lose this password, encrypted credentials cannot be recovered!",
                                   fg='yellow', bold=True))
     click.echo()
-    
+
     import re
     password_set = False
     while not password_set:
         password = getpass.getpass("  Enter master password: ")
-        
+
         # Validate password strength
         if len(password) < 12:
             click.echo(click.style("  ✗ Password must be at least 12 characters.", fg='red'))
             continue
-        
+
         if not re.search(r'[a-z]', password):
             click.echo(click.style("  ✗ Password must contain at least one lowercase letter.", fg='red'))
             continue
-        
+
         if not re.search(r'[A-Z]', password):
             click.echo(click.style("  ✗ Password must contain at least one uppercase letter.", fg='red'))
             continue
-        
+
         if not re.search(r'\d', password):
             click.echo(click.style("  ✗ Password must contain at least one number.", fg='red'))
             continue
-        
+
         if not re.search(r'[!@#$%^&*()_+\-=\[\]{};:\'",.<>?/\\|`~]', password):
             click.echo(click.style("  ✗ Password must contain at least one special character.", fg='red'))
             continue
-        
+
         password_confirm = getpass.getpass("  Confirm master password: ")
         if password != password_confirm:
             click.echo(click.style("  ✗ Passwords don't match!", fg='red'))
             continue
-        
+
         password_set = True
-    
+
     click.echo()
     click.echo("  Enabling encryption...")
-    
+
     try:
         if crypto.enable_encryption(password):
             click.echo("  " + click.style("✓ Encryption enabled!", fg='green'))
@@ -255,7 +577,7 @@ def _wizard_encryption_setup() -> bool:
         click.echo("  " + click.style(f"✗ Error: {e}", fg='red'))
         click.pause("  Press any key to continue...")
         return False
-    
+
     click.pause("  Press any key to continue...")
     return True
 
@@ -293,9 +615,9 @@ def _wizard_create_engagement() -> dict:
     click.echo()
     click.echo("  " + click.style("NOTE:", fg='yellow') + " Type affects default automation and scan aggressiveness")
     click.echo()
-    
+
     type_choice = click.prompt("  Select option", type=click.IntRange(1, 5), default=1, show_default=False)
-    
+
     engagement_types = {
         1: 'penetration_test',
         2: 'bug_bounty',
@@ -303,7 +625,7 @@ def _wizard_create_engagement() -> dict:
         4: 'red_team',
         5: 'custom'
     }
-    
+
     engagement_type = engagement_types[type_choice]
 
     # Create engagement and set it as active
@@ -338,7 +660,8 @@ def _wizard_create_engagement() -> dict:
 def _wizard_tool_check() -> dict:
     """Check installed tools using the centralized tool_checker module."""
     from souleyez.utils.tool_checker import (
-        check_tool_version, EXTERNAL_TOOLS, get_install_command, detect_distro
+        check_tool_version, EXTERNAL_TOOLS, get_install_command, detect_distro,
+        get_tool_version, get_upgrade_command
     )
 
     DesignSystem.clear_screen()
@@ -373,6 +696,7 @@ def _wizard_tool_check() -> dict:
     total = len(wizard_tools)
     tool_status = {}
     wrong_version_tools = []
+    missing_tools = []
 
     for display_name, (category, tool_key) in wizard_tools.items():
         tool_info = EXTERNAL_TOOLS[category][tool_key]
@@ -390,7 +714,6 @@ def _wizard_tool_check() -> dict:
             version_str = version_status['version']
             if not version_str:
                 # Try to get version using the actual found command
-                from souleyez.utils.tool_checker import get_tool_version
                 version_str = get_tool_version(actual_cmd)
 
             # Check if version is OK
@@ -410,8 +733,9 @@ def _wizard_tool_check() -> dict:
                     'name': display_name,
                     'installed': ver_display,
                     'required': min_ver,
-                    'install': get_install_command(tool_info, distro),
-                    'note': version_status.get('version_note')
+                    'install': get_upgrade_command(tool_info, distro) or get_install_command(tool_info, distro),
+                    'note': version_status.get('version_note'),
+                    'tool_info': tool_info
                 })
                 found += 1  # Still counts as found, just needs upgrade
             else:
@@ -423,6 +747,11 @@ def _wizard_tool_check() -> dict:
         else:
             click.echo(f"  {click.style('✗', fg='red')} {display_name:<15} NOT FOUND")
             tool_status[display_name] = {'found': False, 'path': None}
+            missing_tools.append({
+                'name': display_name,
+                'install': get_install_command(tool_info, distro),
+                'tool_info': tool_info
+            })
 
     click.echo()
     click.echo(f"  Found {found}/{total} recommended tools")
@@ -435,17 +764,34 @@ def _wizard_tool_check() -> dict:
             click.echo(f"    - {tool['name']}: installed v{tool['installed']}, needs v{tool['required']}+")
             if tool.get('note'):
                 click.echo(f"      {click.style(tool['note'], fg='bright_black')}")
+
+    # Offer to install/upgrade tools
+    if missing_tools or wrong_version_tools:
         click.echo()
-        click.echo("  " + click.style("To upgrade:", fg='cyan'))
-        for tool in wrong_version_tools:
-            click.echo(f"    {tool['install']}")
+        # Run the install flow (will prompt user)
+        if _run_tool_installs(missing_tools, wrong_version_tools, distro):
+            # Re-check tools after install
+            click.echo()
+            click.echo("  Re-checking tool availability...")
+            click.echo()
 
-    click.echo()
+            found = 0
+            wrong_version_tools = []
+            for display_name, (category, tool_key) in wizard_tools.items():
+                tool_info = EXTERNAL_TOOLS[category][tool_key]
+                version_status = check_tool_version(tool_info)
+
+                if version_status['installed'] and not version_status['needs_upgrade']:
+                    tool_status[display_name] = {'found': True, 'path': shutil.which(tool_info['command'])}
+                    found += 1
+                elif version_status['installed'] and version_status['needs_upgrade']:
+                    tool_status[display_name] = {'found': True, 'needs_upgrade': True}
+                    wrong_version_tools.append({'name': display_name})
+                    found += 1
+                else:
+                    tool_status[display_name] = {'found': False, 'path': None}
 
-    if found < total:
-        click.echo("  " + click.style("TIP:", fg='yellow', bold=True) +
-                  " You can install missing tools later. SoulEyez will work")
-        click.echo("  with the tools you have available.")
+            click.echo(f"  Now have {found}/{total} tools available")
 
     click.echo()
     click.pause("  Press any key to continue...")
@@ -556,14 +902,16 @@ def _install_ollama() -> bool:
 
     click.echo()
     click.echo("  Installing Ollama...")
-    click.echo("  " + click.style("(This may take a minute)", fg='bright_black'))
+    click.echo("  " + click.style("(This may take a minute - downloading ~100MB)", fg='bright_black'))
     click.echo()
 
     try:
-        # Run the official Ollama install script - let it output directly to terminal
+        # Run the official Ollama install script with timeout
+        # Add curl timeouts to fail faster on network issues
         result = subprocess.run(
-            ['bash', '-c', 'curl -fsSL https://ollama.ai/install.sh | sh'],
-            check=False
+            ['bash', '-c', 'curl --connect-timeout 30 --max-time 300 -fsSL https://ollama.ai/install.sh | sh'],
+            check=False,
+            timeout=360  # 6 minute total timeout
         )
 
         if result.returncode == 0:
@@ -575,10 +923,25 @@ def _install_ollama() -> bool:
         else:
             click.echo()
             click.echo("  " + click.style("✗ Installation failed", fg='red'))
-            click.echo("  You can install manually from https://ollama.ai")
+            click.echo("  " + click.style("This is usually a network issue (slow connection or timeout).", fg='yellow'))
+            click.echo()
+            click.echo("  To install manually:")
+            click.echo("    curl -fsSL https://ollama.ai/install.sh | sh")
+            click.echo()
+            click.echo("  Or visit: https://ollama.ai")
             click.pause("  Press any key to continue...")
             return False
 
+    except subprocess.TimeoutExpired:
+        click.echo()
+        click.echo("  " + click.style("✗ Installation timed out", fg='red'))
+        click.echo("  " + click.style("The download took too long. Check your internet connection.", fg='yellow'))
+        click.echo()
+        click.echo("  To install manually:")
+        click.echo("    curl -fsSL https://ollama.ai/install.sh | sh")
+        click.pause("  Press any key to continue...")
+        return False
+
     except Exception as e:
         click.echo()
         click.echo(click.style(f"  ✗ Error: {e}", fg='red'))
@@ -722,7 +1085,7 @@ def _wizard_deliverables(engagement_type: str) -> list:
     DesignSystem.clear_screen()
     _show_wizard_banner()
     width = 60
-    
+
     click.echo("┌" + "─" * (width - 2) + "┐")
     click.echo("│" + click.style(" STEP 7: REPORT TEMPLATES ".center(width - 2), bold=True, fg='cyan') + "│")
     click.echo("└" + "─" * (width - 2) + "┘")
@@ -732,7 +1095,7 @@ def _wizard_deliverables(engagement_type: str) -> list:
     click.echo()
     click.echo("  " + click.style("NOTE:", fg='yellow') + " Templates help generate professional reports from your findings")
     click.echo()
-    
+
     # Default templates based on engagement type
     type_templates = {
         'penetration_test': ['Executive Summary', 'Technical Findings Report', 'Vulnerability Details'],
@@ -741,9 +1104,9 @@ def _wizard_deliverables(engagement_type: str) -> list:
         'red_team': ['Attack Narrative', 'Remediation Roadmap'],
         'custom': ['Technical Findings Report']
     }
-    
+
     recommended = type_templates.get(engagement_type, ['Technical Findings Report'])
-    
+
     click.echo(f"  Recommended for {engagement_type.replace('_', ' ').title()}:")
     for template in recommended:
         click.echo(f"    {click.style('✓', fg='green')} {template}")
@@ -785,7 +1148,7 @@ def _wizard_summary(encryption_enabled, engagement_info, tool_status, ai_enabled
         click.echo(f"  {click.style('✓', fg='green')} AI Features: Enabled (Ollama)")
     else:
         click.echo(f"  {click.style('○', fg='bright_black')} AI Features: Not configured")
-    
+
     # Show tier status
     if is_pro:
         click.echo(f"  {click.style('💎', fg='magenta')} License: PRO")
@@ -794,7 +1157,7 @@ def _wizard_summary(encryption_enabled, engagement_info, tool_status, ai_enabled
             click.echo(f"  {click.style('✓', fg='green')} Auto-Chain: ON ({auto_mode} mode)")
         else:
             click.echo(f"  {click.style('✓', fg='green')} Auto-Chain: OFF")
-        
+
         if templates:
             click.echo(f"  {click.style('✓', fg='green')} Templates: {len(templates)} selected")
     else:
@@ -806,11 +1169,14 @@ def _wizard_summary(encryption_enabled, engagement_info, tool_status, ai_enabled
         click.echo("    • MSF Integration - Advanced attack chains")
         click.echo("    • Reports - Professional deliverables")
         click.echo(f"    {click.style('→ cybersoulsecurity.com/upgrade', fg='cyan')}")
-    
+
     click.echo()
     click.echo("  " + click.style("You're ready to start!", bold=True, fg='cyan'))
     click.echo()
 
+    # Install desktop shortcut automatically
+    _install_desktop_shortcut()
+
     # Prompt for interactive tutorial
     click.echo("  ┌" + "─" * 56 + "┐")
     click.echo("  │" + click.style(" Would you like to run the interactive tutorial?", fg='cyan').center(65) + "│")