SaturnArch-Discover/saturnDiscover

#!/usr/bin/env python

import os
import sys
import json

def tryReadFile(fname, defval):
    try:
        with open(fname, 'r') as file:
            return file.read().replace('\n', '')
    except:
        pass

    return defval

def restartService():
    os.system("systemctl restart saturnDiscover")

def loadHosts():
    HOSTS = {}

    try:
        with open('/etc/slurm-llnl/hosts.json', 'r') as file:
            HOSTS = json.load(file)
    except:
        HOSTS = {}

    return HOSTS

def updateFile(filename, value):
    with open(filename, "w") as outfile:
        outfile.write(value)

def help():
    print("Saturn discover setting tool")
    print()
    print("Using: saturnDiscover <command> [ops]")
    print("supported commands:")
    print("    add         Add value (to list)")
    print("    set         Set value")
    print("    list        List values (from list)")
    print("    get         Get value")
    print("    reload      Reload cluster config")
    print("    lookup      Lookup for value in list")
    print("")
    print("Every command have sub help, when you write help in ops")

def helpAdd():
    print("Saturn discover setting tool")
    print("command ADD - add value to list")
    print()
    print("Using: saturnDiscover add <target> <value>")
    print("supported tragets:")
    print("    res         Node resources list")

def helpSet():
    print("Saturn discover setting tool")
    print("command SET - set value")
    print()
    print("Using: saturnDiscover set <target> <value>")
    print("supported tragets:")
    print("    mac         Node mac address")

def helpList():
    print("Saturn discover setting tool")
    print("command LIST - print values in list")
    print()
    print("Using: saturnDiscover list <target>")
    print("supported tragets:")
    print("    res         Node resources list")
    print("    nodes       Discovered nodes list")

def helpGet():
    print("Saturn discover setting tool")
    print("command GET - get value")
    print()
    print("Using: saturnDiscover get <target>")
    print("supported tragets:")
    print("    mac         Node mac address")

def helpLookup():
    print("Saturn discover setting tool")
    print("command LOOKUP - get value from list")
    print()
    print("Using: saturnDiscover lookup <targetList> <tragetRow>")
    print("supported tragetsList:")
    print("    hostmac         Node mac address lookup by hostname or ip address")

def helpReload():
    print("Saturn discover setting tool")
    print("command RELOAD - reload cluster configuration on all nodes")
    print()
    print("Using: saturnDiscover reload")

argc = len(sys.argv)

if (argc <= 1):
    help()
    exit(0)

## prase command
if (sys.argv[1] == "add"):
    if (argc != 4):
        helpAdd()
        exit(1)
    
    if (sys.argv[2] == "res"):
        res = json.loads(tryReadFile("/etc/slurm-llnl/localRes", "[]"))
        res.append(sys.argv[3])
        updateFile("/etc/slurm-llnl/localRes", json.dumps(res))

    restartService()

elif (sys.argv[1] == "set"):
    if (argc != 4):
        helpSet()
        exit(1)

    if (sys.argv[2] == "mac"):
        updateFile("/etc/slurm-llnl/localMac", sys.argv[3])

    restartService()

elif (sys.argv[1] == "list"):
    if (argc != 3):
        helpList()
        exit(1)

    if (sys.argv[2] == "res"):
        res = json.loads(tryReadFile("/etc/slurm-llnl/localRes", "[]"))
        print(",".join(res["res"]))

    if (sys.argv[2] == "nodes"):
        nodes = loadHosts()

        print("NAME      MASTER    IP                MAC                CPUs    RAM[GB]    VERSION    GRES")
        print("SELF")

        for node in nodes.values():
            isMaster = ""
            if node["type"] == "master":
                isMaster = "YES"

            print(f'{node["name"]: <10}{isMaster: <10}{node["ip"]: <18}{node["mac"]: <19}{node["cpus"]: <8}{int(node["rams"]  / (1024.**2) / 1000): <11}{node["version"]: <11}{",".join(node["res"])}')

elif (sys.argv[1] == "get"):
    if (argc != 3):
        helpGet()
        exit(1)

    if (sys.argv[2] == "mac"):
        print(tryReadFile("/etc/slurm-llnl/localMac", "00:00:00:00:00:00"))

elif (sys.argv[1] == "lookup"):
    if (argc != 4):
        helpLookup()
        exit(1)

    if (sys.argv[2] == "hostmac"):
        nodes = loadHosts()

        for node in nodes.values():
            if node["ip"] == sys.argv[3] or node["name"] == sys.argv[3]:
                print(node["mac"])
                exit(0)


        print("00:00:00:00:00:00")

elif (sys.argv[1] == "reload"):
    if (argc != 2):
        helpReload()
        exit(1)

    ver = int(tryReadFile("/etc/slurm-llnl/confVer", 0))
    updateFile("/etc/slurm-llnl/confVer", str(ver + 1))

    restartService()

else:
    help()
    exit(1)