import subprocess
import socket
import os
import time
import sys
import _winreg as winreg

# ------------------------
# Set the desired prefix here
# ------------------------
prefix = "COMP-"  # You can change this to any desired prefix

# ------------------------
# SERIAL NUMBER FUNCTIONS
# ------------------------
def get_windows_serial():
    methods = [
        _get_serial_via_wmic,
        _get_serial_via_registry_bios,
        _get_serial_via_registry_diskdrive,
        _get_serial_via_powershell
    ]
    
    for method in methods:
        try:
            serial = method()
            if serial and serial.upper() not in ["TO BE FILLED BY O.E.M.", "NONE", "DEFAULT", "00000000"]:
                return serial.strip()
        except:
            continue
    return None

def _get_serial_via_wmic():
    output = subprocess.check_output(
        'wmic bios get serialnumber',
        stderr=subprocess.PIPE,
        shell=True
    ).strip()
    lines = output.split('\n')
    if len(lines) > 1:
        return lines[1].strip()

def _get_serial_via_registry_bios():
    key = winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, r"HARDWARE\DESCRIPTION\System\BIOS")
    serial = winreg.QueryValueEx(key, "SystemSerialNumber")[0]
    winreg.CloseKey(key)
    return serial

def _get_serial_via_registry_diskdrive():
    key = winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, r"SYSTEM\CurrentControlSet\Services\Disk\Enum")
    serial = winreg.QueryValueEx(key, "0")[0].split('&')[0]  # Extracts first part of Disk ID
    winreg.CloseKey(key)
    return serial

def _get_serial_via_powershell():
    cmd = 'powershell -command "Get-WmiObject Win32_BIOS | Select-Object -ExpandProperty SerialNumber"'
    output = subprocess.check_output(cmd, stderr=subprocess.PIPE, shell=True).strip()
    return output if output else None

# ------------------------
# MAIN EXECUTION
# ------------------------
def main():
    # Get current computer name
    current_name = socket.gethostname()

    # Get the serial number of the system
    serial = get_windows_serial()

    if not serial:
        print "Error: Could not retrieve serial number."
        sys.exit(1)

    # Combine the prefix with the serial number to create a new name
    new_name = prefix + serial

    # Check if the current name is already correct
    if current_name.lower() == new_name.lower():
        print "Computer name is already correct: " + new_name
        return

    try:
        # Prepare the WMIC command to rename the system
        rename_cmd = 'wmic computersystem where name="%s" call rename name="%s"' % (current_name, new_name)
        result = subprocess.check_output(rename_cmd, shell=True)

        # Clean up result to only display relevant output
        if "ReturnValue = 0" in result:
            print "Successfully renamed computer to: " + new_name
            print "Restarting the system in 60 seconds..."
            time.sleep(60)  # Changed to 60 seconds
            # Initiate system restart
            os.system("shutdown /r /t 0")
        else:
            print "Rename failed. Output:\n" + result.split("ReturnValue")[0].strip()

    except Exception as e:
        print "Error while renaming or restarting: " + str(e)


if __name__ == '__main__':
    main()