install_agent_admin.py

#!/usr/bin/env python3
"""
Infrastructure Agent Installation Script for Administrator Installation
This script installs the agent to run for all users while correctly detecting the logged-in user.
"""

import os
import sys
import subprocess
import shutil
import argparse
import logging
from pathlib import Path
import ctypes
import winreg

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler('install_agent_admin.log')
    ]
)
logger = logging.getLogger(__name__)

def is_admin():
    """Check if the script is running with administrator privileges"""
    try:
        return ctypes.windll.shell32.IsUserAnAdmin()
    except:
        return False

def install_to_all_users_startup(src_exe_path):
    """Install the agent to the All Users startup folder"""
    try:
        # Get the All Users startup folder path
        all_users_startup = os.path.join(
            os.environ.get('PROGRAMDATA', 'C:\\ProgramData'),
            'Microsoft\\Windows\\Start Menu\\Programs\\StartUp'
        )
        
        # Create the directory if it doesn't exist
        os.makedirs(all_users_startup, exist_ok=True)
        
        # Copy the executable to the startup folder
        dest_path = os.path.join(all_users_startup, 'InfraAgent.exe')
        shutil.copy2(src_exe_path, dest_path)
        logger.info(f"Copied agent to: {dest_path}")
        
        # Create a batch file to run the agent with the correct parameters
        bat_content = '''@echo off
REM Infrastructure Agent Startup Script (All Users)
REM This script runs the agent with the correct parameters for all users

cd /d "%~dp0"
if exist "InfraAgent.exe" (
    start "" /MIN "InfraAgent.exe" --run-continuous --hidden
) else (
    echo ERROR: InfraAgent.exe not found
)
'''
        bat_path = os.path.join(all_users_startup, 'StartInfraAgent.bat')
        with open(bat_path, 'w') as f:
            f.write(bat_content)
        logger.info(f"Created startup script: {bat_path}")
        
        return True
    except Exception as e:
        logger.error(f"Error installing to all users startup: {e}")
        return False

def install_as_windows_service(src_exe_path):
    """Install the agent as a Windows service (more robust approach)"""
    try:
        # Create a service installation script
        service_script_content = f'''@echo off
REM Infrastructure Agent Service Installation Script

REM Install the agent as a Windows service
"{src_exe_path}" --install-service

REM Start the service
net start InfraAgent

echo Service installation completed.
pause
'''
        
        # Save the service script
        script_path = os.path.join(os.path.dirname(src_exe_path), 'install_service.bat')
        with open(script_path, 'w') as f:
            f.write(service_script_content)
        
        logger.info(f"Created service installation script: {script_path}")
        logger.info("To install as a service, run the script as administrator")
        
        return True
    except Exception as e:
        logger.error(f"Error creating service installation script: {e}")
        return False

def create_registry_run_key(src_exe_path):
    """Create a registry entry to run the agent at startup for all users"""
    try:
        # Registry path for all users run key
        key_path = r"SOFTWARE\Microsoft\Windows\CurrentVersion\Run"
        
        # Open the registry key with write access
        with winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, key_path, 0, winreg.KEY_WRITE) as key:
            # Set the value to run the agent
            winreg.SetValueEx(
                key, 
                "InfraAgent", 
                0, 
                winreg.REG_SZ, 
                f'"{src_exe_path}" --run-continuous --hidden'
            )
        
        logger.info("Added agent to Windows Run registry for all users")
        return True
    except Exception as e:
        logger.error(f"Error adding to registry Run key: {e}")
        return False

def main():
    parser = argparse.ArgumentParser(description='Install Infrastructure Agent for all users')
    parser.add_argument('--method', choices=['startup', 'service', 'registry'], 
                       default='startup', help='Installation method')
    parser.add_argument('--source', help='Path to InfraAgent.exe (default: current directory)')
    
    args = parser.parse_args()
    
    # Check for admin privileges
    if not is_admin():
        logger.error("This script must be run as administrator!")
        logger.info("Please run this script with administrator privileges.")
        return False
    
    # Determine source path
    if args.source:
        src_exe_path = args.source
    else:
        # Look for InfraAgent.exe in current directory or dist subdirectory
        possible_paths = [
            'InfraAgent.exe',
            'dist/InfraAgent.exe',
            'package/InfraAgent.exe'
        ]
        
        src_exe_path = None
        for path in possible_paths:
            if os.path.exists(path):
                src_exe_path = os.path.abspath(path)
                break
        
        if not src_exe_path:
            logger.error("Could not find InfraAgent.exe. Please specify --source parameter.")
            return False
    
    logger.info(f"Installing agent from: {src_exe_path}")
    
    # Perform installation based on method
    if args.method == 'startup':
        success = install_to_all_users_startup(src_exe_path)
        if success:
            logger.info("Agent installed to All Users startup folder successfully!")
            logger.info("The agent will now run for all users and correctly detect the logged-in user.")
    elif args.method == 'service':
        success = install_as_windows_service(src_exe_path)
        if success:
            logger.info("Service installation script created successfully!")
            logger.info("Run the generated install_service.bat as administrator to complete service installation.")
    elif args.method == 'registry':
        success = create_registry_run_key(src_exe_path)
        if success:
            logger.info("Agent added to registry Run key successfully!")
            logger.info("The agent will run for all users at startup.")
    
    return success

if __name__ == "__main__":
    try:
        success = main()
        sys.exit(0 if success else 1)
    except Exception as e:
        logger.error(f"Unexpected error: {e}")
        sys.exit(1)