build_package.py

import subprocess
import sys
import os
import argparse
import shutil
from pathlib import Path
import logging
import logging.handlers
import traceback

def install_pyinstaller():
    """Install PyInstaller if not already installed"""
    try:
        import PyInstaller
        print("PyInstaller is already installed")
    except ImportError:
        print("Installing PyInstaller...")
        subprocess.check_call([sys.executable, "-m", "pip", "install", "pyinstaller"])
        print("PyInstaller installed successfully")

def build_executable():
    """Build standalone executable using PyInstaller"""
    try:
        print("Building standalone executable...")
        
        # Configure logging with file rotation to prevent permission issues
        class SafeRotatingFileHandler(logging.handlers.RotatingFileHandler):
            def __init__(self, filename, mode='a', maxBytes=1024*1024, backupCount=3, encoding=None, delay=False):
                # Ensure the log directory exists and has proper permissions
                import time
                from pathlib import Path
                log_path = Path(filename)
                if log_path.parent and not log_path.parent.exists():
                    try:
                        log_path.parent.mkdir(parents=True, exist_ok=True)
                    except Exception as e:
                        print(f"Warning: Could not create log directory {log_path.parent}: {e}")
                
                # Try to handle permission issues gracefully
                try:
                    super().__init__(filename, mode, maxBytes, backupCount, encoding, delay)
                except (PermissionError, OSError) as e:
                    # If we can't write to the log file, create a unique log file name
                    try:
                        unique_filename = f"{log_path.stem}_{int(time.time())}{log_path.suffix}"
                        unique_path = log_path.parent / unique_filename if log_path.parent else Path(unique_filename)
                        super().__init__(str(unique_path), mode, maxBytes, backupCount, encoding, delay)
                        print(f"Warning: Could not write to {filename}, using {unique_path} instead")
                    except Exception as e2:
                        print(f"Error creating log file: {e2}")
                        # Fallback to console only logging
                        pass
        
        # Configure logging
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s - %(levelname)s - %(message)s',
            handlers=[
                SafeRotatingFileHandler('build_package.log', maxBytes=1024*1024, backupCount=3),  # 1MB files, keep 3 backups
                logging.StreamHandler()
            ]
        )
        logger = logging.getLogger(__name__)
        
        # Create PyInstaller command
        cmd = [
            "pyinstaller",
            "--onefile",  # Single executable file
            "--windowed",  # No console window (hidden mode)
            "--name", "InfraAgent",
            "--add-data", "config.json;.",  # Include config file
            "--add-data", "agentlogo.ico;.",  # Include icon file
            "--icon", "agentlogo.ico",  # Use the agent logo as icon
            "--hidden-import", "win32timezone",
            "--hidden-import", "win32api",
            "--hidden-import", "win32service",
            "--hidden-import", "win32serviceutil",
            "--hidden-import", "servicemanager",
            "--hidden-import", "win32event",
            "--hidden-import", "pystray",
            "--hidden-import", "PIL",
            "--hidden-import", "PIL.Image",
            "--hidden-import", "PIL.ImageDraw",
            "--hidden-import", "ctypes",
            "--hidden-import", "ctypes.wintypes",
            "main.py"
        ]
        
        # Run PyInstaller
        result = subprocess.run(cmd, capture_output=True, text=True)
        
        if result.returncode == 0:
            print("Executable built successfully!")
            print("Output location: dist/InfraAgent.exe")
            return True
        else:
            print("Error building executable:")
            print(result.stderr)
            return False
            
    except Exception as e:
        print(f"Error during build process: {e}")
        print(f"Traceback: {traceback.format_exc()}")
        return False

def create_package():
    """Create a complete package for distribution"""
    try:
        print("Creating distribution package...")
        
        # Configure logging with file rotation to prevent permission issues
        class SafeRotatingFileHandler(logging.handlers.RotatingFileHandler):
            def __init__(self, filename, mode='a', maxBytes=1024*1024, backupCount=3, encoding=None, delay=False):
                # Ensure the log directory exists and has proper permissions
                import time
                from pathlib import Path
                log_path = Path(filename)
                if log_path.parent and not log_path.parent.exists():
                    try:
                        log_path.parent.mkdir(parents=True, exist_ok=True)
                    except Exception as e:
                        print(f"Warning: Could not create log directory {log_path.parent}: {e}")
                
                # Try to handle permission issues gracefully
                try:
                    super().__init__(filename, mode, maxBytes, backupCount, encoding, delay)
                except (PermissionError, OSError) as e:
                    # If we can't write to the log file, create a unique log file name
                    try:
                        unique_filename = f"{log_path.stem}_{int(time.time())}{log_path.suffix}"
                        unique_path = log_path.parent / unique_filename if log_path.parent else Path(unique_filename)
                        super().__init__(str(unique_path), mode, maxBytes, backupCount, encoding, delay)
                        print(f"Warning: Could not write to {filename}, using {unique_path} instead")
                    except Exception as e2:
                        print(f"Error creating log file: {e2}")
                        # Fallback to console only logging
                        pass
        
        # Configure logging
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s - %(levelname)s - %(message)s',
            handlers=[
                SafeRotatingFileHandler('build_package.log', maxBytes=1024*1024, backupCount=3),  # 1MB files, keep 3 backups
                logging.StreamHandler()
            ]
        )
        logger = logging.getLogger(__name__)
        
        # Create package directory
        package_dir = Path("package")
        if package_dir.exists():
            try:
                print(f"Removing existing package directory: {package_dir}")
                shutil.rmtree(package_dir)
                print("Existing package directory removed successfully")
            except Exception as e:
                print(f"Error removing existing package directory: {e}")
                print("\nPossible solutions:")
                print("1. Close any applications that might be using files in the package folder")
                print("2. Run this script as administrator")
                print("3. Restart your computer and try again")
                print("4. Use the force_delete.ps1 script: powershell -ExecutionPolicy Bypass -File force_delete.ps1")
                return False
        package_dir.mkdir()
        
        # Copy executable
        exe_path = Path("dist/InfraAgent.exe")
        if exe_path.exists():
            shutil.copy(exe_path, package_dir / "InfraAgent.exe")
        else:
            print("Error: Executable not found. Please build it first.")
            return False
        
        # Copy config file
        config_path = Path("config.json")
        if config_path.exists():
            shutil.copy(config_path, package_dir / "config.json")
        
        # Copy icon file
        icon_file = Path("agentlogo.ico")
        if icon_file.exists():
            shutil.copy(icon_file, package_dir / "agentlogo.ico")
            print("Copied agentlogo.ico to package")
        else:
            print("Warning: agentlogo.ico not found")
        
        # Create run script
        run_script = """@echo off
start "" "InfraAgent.exe" --run-continuous
exit
"""
        with open(package_dir / "run_agent.bat", "w") as f:
            f.write(run_script)
        
        # Create README file
        readme_content = """Infrastructure Agent Package
========================

This package contains the standalone Infrastructure Agent executable that can run 
without requiring Python or additional libraries to be installed.

Files:
- InfraAgent.exe: The standalone agent executable with custom icon
- config.json: Configuration file (can be modified before installation)
- run_agent.bat: Script to manually run the agent

To run:
1. Double-click run_agent.bat to start the agent with a visible taskbar icon
   OR
2. Run InfraAgent.exe directly

The agent will:
- Collect PC information (OS, CPU, RAM, disk, network, etc.)
- Send data to the configured server
- Display an icon in the taskbar
"""
        with open(package_dir / "README.txt", "w") as f:
            f.write(readme_content)
        
        print(f"Package created successfully in {package_dir}/")
        return True
        
    except Exception as e:
        print(f"Error creating package: {e}")
        print(f"Traceback: {traceback.format_exc()}")
        return False

def main():
    # Check if we're running in windowed mode (frozen executable without console)
    is_windowed = getattr(sys, 'frozen', False) and sys.stdout is None
    
    # Initialize args with default values
    import argparse
    args = argparse.Namespace(
        build=False,
        package=False
    )
    
    # Only try to parse arguments if we're not in windowed mode
    if not is_windowed:
        try:
            parser = argparse.ArgumentParser(description='Build Infrastructure Agent package')
            parser.add_argument('--build', action='store_true', help='Build executable')
            parser.add_argument('--package', action='store_true', help='Create distribution package')
            
            # Try to parse arguments only if we have arguments to parse
            if len(sys.argv) > 1:
                args = parser.parse_args()
        except Exception as e:
            # Silently ignore argument parsing errors to prevent crashes in windowed mode
            print(f"Silently ignoring argument parsing error in build package: {e}")
    
    # If no arguments provided, do both (default behavior)
    if not args.build and not args.package:
        args.build = True
        args.package = True
    
    # Install PyInstaller if needed
    install_pyinstaller()
    
    # Build executable if requested
    if args.build:
        if not build_executable():
            print("Failed to build executable")
            return False
    
    # Create package if requested
    if args.package:
        if not create_package():
            print("Failed to create package")
            return False
    
    print("Build process completed successfully!")
    return True

if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)