setup.py

#!/usr/bin/env python3
"""
setup.py - ProjectSentry Initial Setup Wizard
Creates initial configuration and sets up the encryption system
"""

import os
import sys
import configparser
import getpass
import hashlib
from pathlib import Path
import json


def print_header():
    """Print setup wizard header"""
    print("=" * 60)
    print("🔐 PROJECT SENTRY - INITIAL SETUP WIZARD")
    print("=" * 60)
    print("Welcome to ProjectSentry! Let's set up your secure file vault.")
    print("This setup will only run once to configure your system.\n")


def get_master_password():
    """Get and verify master password from user"""
    print("🔑 STEP 1: SET YOUR MASTER PASSWORD")
    print("-" * 40)
    print("Your master password will be used to encrypt all files.")
    print("⚠️  IMPORTANT: If you lose this password, your files cannot be recovered!")
    print("💡 Recommendation: Use a strong password with at least 12 characters\n")
    
    while True:
        password1 = getpass.getpass("Enter your master password: ").strip()
        
        if len(password1) == 0:
            print("❌ Password cannot be empty. Please try again.\n")
            continue
            
        if len(password1) < 8:
            print("⚠️  Warning: Password is shorter than 8 characters.")
            confirm = input("Continue anyway? (y/n): ").lower()
            if confirm != 'y':
                continue
        
        password2 = getpass.getpass("Confirm your master password: ").strip()
        
        if password1 != password2:
            print("❌ Passwords don't match. Please try again.\n")
            continue
        
        # Create password hash for verification (not storage - we don't store passwords)
        password_strength = analyze_password_strength(password1)
        print(f"✅ Password set successfully!")
        print(f"💪 Password strength: {password_strength}")
        
        return password1


def analyze_password_strength(password):
    """Analyze password strength"""
    score = 0
    feedback = []
    
    if len(password) >= 8:
        score += 1
    if len(password) >= 12:
        score += 1
    if any(c.islower() for c in password):
        score += 1
    if any(c.isupper() for c in password):
        score += 1
    if any(c.isdigit() for c in password):
        score += 1
    if any(c in "!@#$%^&*()_+-=[]{}|;:,.<>?" for c in password):
        score += 1
    
    if score <= 2:
        return "Weak"
    elif score <= 4:
        return "Medium"
    else:
        return "Strong"


def setup_vault_directory():
    """Set up the vault directory for monitoring"""
    print("\n📁 STEP 2: CONFIGURE YOUR VAULT DIRECTORY")
    print("-" * 40)
    print("Choose a directory that ProjectSentry will monitor for new files.")
    print("Any file placed in this directory will be automatically encrypted.\n")
    
    while True:
        print("Options:")
        print("1. Use default vault directory (recommended)")
        print("2. Choose custom directory")
        print("3. Create new directory")
        
        choice = input("\nEnter your choice (1-3): ").strip()
        
        if choice == '1':
            # Default directory
            home_dir = Path.home()
            vault_dir = home_dir / "ProjectSentry_Vault"
            break
            
        elif choice == '2':
            # Custom directory
            custom_path = input("Enter full path to your vault directory: ").strip()
            vault_dir = Path(custom_path)
            
            if not vault_dir.exists():
                print(f"❌ Directory doesn't exist: {vault_dir}")
                create = input("Create this directory? (y/n): ").lower()
                if create != 'y':
                    continue
            break
            
        elif choice == '3':
            # Create new directory
            dir_name = input("Enter directory name: ").strip()
            if not dir_name:
                print("❌ Directory name cannot be empty.")
                continue
                
            location = input("Enter parent directory path (or press Enter for Desktop): ").strip()
            if not location:
                location = Path.home() / "Desktop"
            else:
                location = Path(location)
                
            vault_dir = location / dir_name
            break
        else:
            print("❌ Invalid choice. Please enter 1, 2, or 3.")
    
    # Create the directory if it doesn't exist
    try:
        vault_dir.mkdir(parents=True, exist_ok=True)
        print(f"✅ Vault directory ready: {vault_dir}")
        return str(vault_dir.absolute())
    except Exception as e:
        print(f"❌ Error creating directory: {e}")
        return None


def configure_file_types():
    """Configure which file types to encrypt"""
    print("\n📄 STEP 3: FILE TYPE CONFIGURATION")
    print("-" * 40)
    print("Choose which types of files ProjectSentry should encrypt:\n")
    
    # Predefined file type categories
    file_categories = {
        'documents': {
            'name': 'Documents',
            'extensions': ['.txt', '.doc', '.docx', '.pdf', '.rtf', '.odt'],
            'description': 'Text documents, Word files, PDFs'
        },
        'images': {
            'name': 'Images',
            'extensions': ['.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff', '.svg'],
            'description': 'Photos and image files'
        },
        'videos': {
            'name': 'Videos',
            'extensions': ['.mp4', '.avi', '.mkv', '.mov', '.wmv', '.flv', '.webm'],
            'description': 'Video files'
        },
        'audio': {
            'name': 'Audio',
            'extensions': ['.mp3', '.wav', '.flac', '.aac', '.ogg', '.wma'],
            'description': 'Music and audio files'
        },
        'archives': {
            'name': 'Archives',
            'extensions': ['.zip', '.rar', '.7z', '.tar', '.gz'],
            'description': 'Compressed archive files'
        },
        'code': {
            'name': 'Code/Development',
            'extensions': ['.py', '.js', '.html', '.css', '.cpp', '.java', '.c', '.h'],
            'description': 'Programming source code'
        }
    }
    
    print("Available file type categories:")
    for i, (key, category) in enumerate(file_categories.items(), 1):
        print(f"{i}. {category['name']}: {category['description']}")
    
    print(f"{len(file_categories) + 1}. All files (encrypt everything)")
    print(f"{len(file_categories) + 2}. Custom selection")
    
    while True:
        choice = input(f"\nEnter your choice (1-{len(file_categories) + 2}): ").strip()
        
        try:
            choice_num = int(choice)
            
            if 1 <= choice_num <= len(file_categories):
                # Single category selected
                selected_key = list(file_categories.keys())[choice_num - 1]
                selected_extensions = file_categories[selected_key]['extensions']
                print(f"✅ Selected: {file_categories[selected_key]['name']}")
                return selected_extensions, False
                
            elif choice_num == len(file_categories) + 1:
                # All files
                print("✅ Selected: All files will be encrypted")
                return [], True  # Empty list with encrypt_all=True
                
            elif choice_num == len(file_categories) + 2:
                # Custom selection
                return configure_custom_file_types(file_categories)
            else:
                print(f"❌ Invalid choice. Please enter 1-{len(file_categories) + 2}")
                
        except ValueError:
            print("❌ Please enter a valid number.")


def configure_custom_file_types(file_categories):
    """Allow custom file type selection"""
    print("\n🔧 CUSTOM FILE TYPE SELECTION")
    print("Select multiple categories (enter numbers separated by commas):")
    
    for i, (key, category) in enumerate(file_categories.items(), 1):
        print(f"{i}. {category['name']}")
    
    while True:
        selection = input("Enter category numbers (e.g., 1,3,5): ").strip()
        
        try:
            selected_nums = [int(x.strip()) for x in selection.split(',')]
            
            # Validate all numbers
            if all(1 <= num <= len(file_categories) for num in selected_nums):
                selected_extensions = []
                selected_names = []
                
                for num in selected_nums:
                    key = list(file_categories.keys())[num - 1]
                    selected_extensions.extend(file_categories[key]['extensions'])
                    selected_names.append(file_categories[key]['name'])
                
                print(f"✅ Selected: {', '.join(selected_names)}")
                return selected_extensions, False
            else:
                print(f"❌ Invalid numbers. Please use 1-{len(file_categories)}")
                
        except ValueError:
            print("❌ Please enter numbers separated by commas (e.g., 1,3,5)")


def create_config_file(vault_directory, file_extensions, encrypt_all):
    """Create the config.ini file"""
    print("\n⚙️ STEP 4: CREATING CONFIGURATION FILE")
    print("-" * 40)
    
    config = configparser.ConfigParser()
    
    # Settings section
    config['Settings'] = {
        'watch_directory': vault_directory,
        'delete_original_after_encryption': 'false',
        'encrypted_file_extension': '.enc',
        'encrypt_all_files': str(encrypt_all).lower()
    }
    
    # File types section
    if not encrypt_all:
        config['FileTypes'] = {
            'allowed_extensions': ','.join(file_extensions)
        }
    
    # Hardware section
    config['Hardware'] = {
        'dongle_id': 'ProjectSentryKey_Alpha_9182',
        'hardware_required': 'true'
    }
    
    # System section
    config['System'] = {
        'setup_completed': 'true',
        'version': '1.0'
    }
    
    try:
        with open('config.ini', 'w') as f:
            config.write(f)
        print("✅ Configuration file created successfully!")
        return True
    except Exception as e:
        print(f"❌ Error creating config file: {e}")
        return False


def create_system_info_file(vault_directory):
    """Create a system info file for reference"""
    info = {
        'setup_date': str(datetime.now()),
        'vault_directory': vault_directory,
        'system_version': '1.0',
        'setup_completed': True
    }
    
    try:
        with open('system_info.json', 'w') as f:
            json.dump(info, f, indent=2)
    except:
        pass  # Not critical if this fails


def print_setup_complete(vault_directory):
    """Print setup completion message"""
    print("\n" + "=" * 60)
    print("🎉 SETUP COMPLETED SUCCESSFULLY!")
    print("=" * 60)
    print("Your ProjectSentry system is now configured and ready to use.\n")
    
    print("📋 SETUP SUMMARY:")
    print(f"   🔐 Master password: Set and verified")
    print(f"   📁 Vault directory: {vault_directory}")
    print(f"   ⚙️  Configuration: Saved to config.ini")
    print(f"   🔑 Hardware: Mock dongle (for testing)")
    
    print("\n🚀 NEXT STEPS:")
    print("1. Run 'python monitor.py' to start the file monitor")
    print("2. Add files to your vault directory to test encryption")
    print("3. The system will automatically encrypt new files")
    
    print("\n⚠️  IMPORTANT REMINDERS:")
    print("• Remember your master password - it cannot be recovered!")
    print("• Keep your ProjectSentry dongle connected when using the system")
    print("• Encrypted files will have '.enc' extension")
    
    print("\n" + "=" * 60)


def main():
    """Main setup function"""
    from datetime import datetime
    
    # Check if setup already completed
    if os.path.exists('config.ini'):
        config = configparser.ConfigParser()
        config.read('config.ini')
        
        if config.has_option('System', 'setup_completed') and config.getboolean('System', 'setup_completed'):
            print("⚠️  ProjectSentry is already set up!")
            print("Configuration file 'config.ini' already exists.")
            
            choice = input("Do you want to run setup again? This will overwrite existing settings (y/n): ")
            if choice.lower() != 'y':
                print("Setup cancelled. Use 'python monitor.py' to start monitoring.")
                return
    
    print_header()
    
    try:
        # Step 1: Master Password
        master_password = get_master_password()
        
        # Step 2: Vault Directory
        vault_directory = setup_vault_directory()
        if not vault_directory:
            print("❌ Setup failed: Could not configure vault directory")
            return
        
        # Step 3: File Types
        file_extensions, encrypt_all = configure_file_types()
        
        # Step 4: Create Configuration
        if not create_config_file(vault_directory, file_extensions, encrypt_all):
            print("❌ Setup failed: Could not create configuration file")
            return
        
        # Create system info file
        create_system_info_file(vault_directory)
        
        # Show completion message
        print_setup_complete(vault_directory)
        
        print("\nSetup completed! You can now run 'python monitor.py' to start monitoring.")
        
    except KeyboardInterrupt:
        print("\n\n⚠️  Setup cancelled by user.")
        print("Run 'python setup.py' again to complete setup.")
    except Exception as e:
        print(f"\n❌ Setup failed with error: {e}")
        print("Please try running setup again.")


if __name__ == "__main__":
    main()