#!/usr/bin/python
import atexit
import argparse
import getpass
import ssl
import sys

from pyVmomi import vim, vmodl
from pyVim.connect import SmartConnect, Disconnect

def get_args():
    parser = argparse.ArgumentParser(description='Configure ESXi for non-root user and Lockdown Mode')
    parser.add_argument('-s', '--host', required=True, help='IP of the ESXi host')
    parser.add_argument('-u', '--user', required=True, help='Root user to connect initially')
    parser.add_argument('-n', '--new_user', required=True, help='Name of the new user to create (e.g., svc_vcenter)')
    parser.add_argument('-p', '--new_pass', required=True, help='Password for the new user')
    
    args = parser.parse_args()
    # Securely enter root password
    args.root_pass = getpass.getpass(prompt='Enter root password: ')
    return args

def create_user_with_minimal_privileges(content, username, password):
    """
    Create user, new role (minimal privileges) and assign permissions.
    """
    auth_manager = content.authorizationManager
    user_manager = content.user_manager

    # 1. Create user
    print(f"[*] Creating user: {username}...")
    try:
        user_manager.CreateUser(username, password, None)
        print(f"[+] User {username} has been created.")
    except vmodl.fault.AlreadyExists:
        print(f"[!] User {username} already exists. Skipping user creation.")
    except Exception as e:
        print(f"[!] Error creating user: {e}")
        sys.exit(1)

    # 2. Create new Role with minimum privileges (Minimal Privileges)
    # vCenter requires this permission to create the vpxuser and manage the host.
    role_name = "vCenterAdminMinimal"
    
    # List of required privileges (Important: Host.Local.ManageUserGroups)
    required_privileges = [
        'Host.Config.*',           # Configure host
        'Host.Cim.*',              # Hardware management (CIM)
        'Host.Inventory.*',        # Manage inventory
        'Host.Local.ManageUserGroups', # REQUIRED: To create vpxuser
        'System.Read',
        'System.Write'
    ]
    
    # Check if role already exists
    role_id = None
    for role in auth_manager.role_list:
        if role.name == role_name:
            role_id = role.roleId
            print(f"[!] Role '{role_name}' already exists. Using existing role.")
            break
    
    if role_id is None:
        print(f"[*] Creating new Role: {role_name}...")
        try:
            # Map string privilege to object privilege id
            priv_ids = []
            avail_privs = auth_manager.privilege_list
            priv_dict = {p.name: p.privId for p in avail_privs}
            
            for p in required_privileges:
                if p in priv_dict:
                    priv_ids.append(priv_dict[p])
                elif '*' in p:
                    # Handle wildcard (e.g., Host.Config.*)
                    base = p.replace('*', '')
                    for key in priv_dict:
                        if key.startswith(base):
                            priv_ids.append(priv_dict[key])

            auth_manager.AddRole(role_name, priv_ids)
            role_id = next(r.roleId for r in auth_manager.role_list if r.name == role_name)
            print(f"[+] Role '{role_name}' has been created with {len(priv_ids)} privileges.")
        except Exception as e:
            print(f"[!] Error creating role: {e}")
            sys.exit(1)

    # 3. Assign Role to User on Root Folder
    print(f"[*] Assigning permissions to user {username}...")
    perm = vim.Permission()
    perm.principal = username
    perm.group = False
    perm.roleId = role_id
    perm.propagate = True
    
    try:
        auth_manager.SetEntityPermissions(content.rootFolder, [perm])
        print(f"[+] Permissions assigned successfully.")
    except Exception as e:
        print(f"[!] Error assigning permissions: {e}")
        sys.exit(1)

def configure_lockdown_mode(host_system, username):
    """
    Configure Lockdown Mode and add user to Exception List.
    """
    print("[*] Configuring Lockdown Mode...")
    
    # Get Access Manager
    host_config_manager = host_system.configManager
    access_manager = host_config_manager.hostAccessManager
    
    # Get current exception list
    # Note: You might need to change user to domain\user if using domain
    current_exceptions = list(access_manager.exceptionUsers) if access_manager.exceptionUsers else []
    
    if username not in current_exceptions:
        current_exceptions.append(username)
        print(f"[*] Adding user {username} to the exception list...")
    
    # Update Exception List before enabling Lockdown (If not, the user might be blocked immediately)
    try:
        access_manager.UpdateLockdownExceptions(current_exceptions)
        print("[+] Updated Exception List.")
    except Exception as e:
        print(f"[!] Warning: Could not update exception list: {e}")
        print("[!] If Lockdown is enabled, the user might be blocked. Please check manually.")

    # Enable Lockdown Mode (lockdownNormal)
    try:
        # State: lockdownNormal (Normal), lockdownStrict (Strict)
        access_manager.ChangeLockdownMode("lockdownNormal")
        print("[+] Enabled Lockdown Mode (lockdownNormal).")
        print("[+] Root account has been locked directly.")
    except Exception as e:
        print(f"[!] Error enabling Lockdown Mode: {e}")

def main():
    args = get_args()
    
    # SSL context ignoring verify (for Lab environment)
    context = ssl.SSLContext(ssl.PROTOCOL_TLSv1_2)
    context.verify_mode = ssl.CERT_NONE

    # Connect to ESXi using ROOT
    try:
        si = SmartConnect(host=args.host, user=args.user, pwd=args.root_pass, sslContext=context)
        atexit.register(Disconnect, si)
        content = si.RetrieveContent()
        
        # Get Host System object (Get first host in the list)
        host_system = content.rootFolder.childEntity[0].hostFolder.childEntity[0].hostFolder.childEntity[0].host[0]
        
        # 1. Create User & Role
        create_user_with_minimal_privileges(content, args.new_user, args.new_pass)
        
        # 2. Configure Lockdown Mode
        configure_lockdown_mode(host_system, args.new_user)
        
        print("\n" + "="*50)
        print("SUCCESS! System has been configured securely.")
        print("Please use the information below to add host to vCenter:")
        print(f"Host IP: {args.host}")
        print(f"Username: {args.new_user}")
        print(f"Password: {args.new_pass}")
        print("="*50)

    except vmodl.MethodFault as e:
        print(f"[!] vAPI/VMware Error: {e.msg}")
        sys.exit(1)
    except Exception as e:
        print(f"[!] Undefined Error: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()