nixos_config/util/wol_man.py

from time import time
import subprocess
from _socket import socket
import argparse
import time
import socket
import subprocess

TARGET_LINK = "f0:2f:74:30:0c:d0"
TARGET_HOSTNAME = "tower.home.the.malice.zone"
TARGET_IP = socket.gethostbyname(TARGET_HOSTNAME)

def remote_cmd(command):
    print(f"Executing: ssh tower '{command}'")
    response = subprocess.run(f"ssh tower '{command}'", shell=True, capture_output=True, text=True)
    return (response.returncode, response.stdout, response.stderr)

def weboot():
    print("Rebooting to windows")
    code, stdout, stderr = remote_cmd("sudo efibootmgr -n 0000")
    if code == 0:
        code, stdout, stderr = remote_cmd("sudo reboot")
        time.sleep(60)
    if code != 0:
        raise RuntimeError(f"weboot failed {stderr}")
    return code, stdout, stderr

def send_wol():
    print(f"Broadcasting WOL packet to {TARGET_LINK}")
    response = subprocess.run(f"wakeonlan {TARGET_LINK}", shell=True, capture_output=True, text=True)
    return (response.returncode, response.stdout, response.stderr)

def get_os():
    print("Determining target OS")
    code, stdout, stderr = remote_cmd("whoami")
    if code == 255 and "No route to host" in stderr:
        print("Target is offline")
        return "offline"
    elif code != 0:
        raise RuntimeError(f"win_or_linux failed {stderr}")
    if "windows" in stdout.lower():
        print("Target is running Windows")
        return "windows"
    else:
        print("Target is running Linux")
        return "linux"

def swap_os():
    match get_os():
        case "windows":
            code, stdout, stderr = remote_cmd("shutdown /r")
        case "linux":
            code, stdout, stderr = weboot()
        case "offline":
            return False
    return True

def boot_linux():
    print("Booting Linux")
    match get_os():
        case "windows":
            print("Tower is busy")
            return False
        case "linux":
            print("Linux is already running")
            return True
        case "offline":
            print("Booting via WOL")
            send_wol()
            time.sleep(60)
            return True

def boot_windows():
    print("Booting Windows")
    match get_os():
        case "windows":
            print("Windows is already running")
            return True
        case "linux":
            print("Tower is busy")
            return False
        case "offline":
            print("Booting via WOL")
            boot_linux()
            weboot()
            if get_os() == "windows":
                return True
            else:
                return False

def reboot():
    print("Rebooting target")
    orig_os = get_os()
    match orig_os:
        case "windows":
            code, stdout, stderr = remote_cmd("shutdown /r")
            time.sleep(60)
            weboot()
            return True
        case "linux":
            code, stdout, stderr = remote_cmd("sudo reboot")
            return True
        case "offline":
            print("Target is offline, cannot reboot")
            return False
        
def shutdown():
    print("Shutting down target")
    match get_os():
        case "windows":
            code, stdout, stderr = remote_cmd("shutdown /s")
            time.sleep(30)
            return True
        case "linux":
            code, stdout, stderr = remote_cmd("sudo shutdown -h now")
            time.sleep(30)
            return True
        case "offline":
            print("Target is already offline")
            return False

def action_do(action):
    match action:
        case "get_os":
            os = get_os()
            print(f"Current OS: {os}")
        case "boot_linux":
            if boot_linux():
                print("Linux is now running")
            else:
                print("Failed to boot Linux")
        case "boot_windows":
            if boot_windows():
                print("Windows is now running")
            else:
                print("Failed to boot Windows")
        case "swap_os":
            if swap_os():
                print("OS swapped successfully")
            else:
                print("Failed to swap OS")
        case "reboot":
            if reboot():
                print("Reboot command sent successfully")
            else:
                print("Failed to reboot")
        case "shutdown":
            if shutdown():
                print("Shutdown command sent successfully")
            else:
                print("Failed to shutdown")
        case "winupdate":
            if get_os() != "offline":
                print("Target is in use, belaying update 1h")
                time.sleep(3600)
                action_do("winupdate")
            if boot_windows():
                print("Initiating Windows update")
            else:
                print("Failed to boot Windows for update")
                return
            code, stdout, stderr = remote_cmd("powershell -Command \"& {Install-WindowsUpdate -AcceptAll -AutoReboot}\"")
            if code == 0:
                print("Windows update initiated successfully")
            else:
                print(f"Failed to initiate Windows update: {stderr}")
            #first reboot to windows
            time.sleep(300)
            weboot()
            #second reboot after update
            time.sleep(600)
            reboot()
            #shutdown
            time.sleep(300)
            shutdown()

def main():
    parser = argparse.ArgumentParser(description="WOL and reboot utility")
    parser.add_argument("action", choices=["get_os", "boot_linux", "boot_windows", "swap_os", "reboot", "shutdown", "winupdate"], help="Action to perform")
    action_do(parser.parse_args().action)
    
if __name__ == "__main__":
    main()