#!/usr/bin/python3
# Agent-push tool for Dwarfguard
# Copyright (c) Jan Otte, Dwarf Technologies s.r.o. All rights reserved.

import sys
import getopt
import os
import signal
import subprocess
import time
import multiprocessing as mp
import random
import math
import re
import tempfile
from enum import IntEnum

class Devtype(IntEnum):
    ANYDEVICE = 0
    LINUXBOX = 1
    ADVROUTER = 2
    OPENWRT = 3
    TELTROUTER = 4
    MAXDEVICES = 5

class Loglevel(IntEnum):
    NONE = 0
    ERROR = 1
    FINISH = 2
    INFO = 3
    DEBUG = 4

class Retcode(IntEnum):
    NOERROR = 0
    INTERRUPT = 1
    PARSE = 2
    ROUTE = 3
    CONNECTION = 4
    DEVTYPE = 5
    REMOVAL = 6
    ARCHIVE = 7
    UNCOMPRESS = 8
    INSTALL = 9
    DEVID = 10
    NOOLDAGENT = 11
    NONEWAGENT = 12

version = "0.2.0"
errout = subprocess.DEVNULL
filename = None
teststring="test"
agent_dir="/opt/adwarfg"
agent_tar_ext="tgz"
agent_ipk_ext="ipk"
agent_fname="dwarfg_agent.sh"
agent_fname_prefix="adwarfg"
archive_path=""
pwds_from_file={}
manager = mp.Manager()
terminate = manager.Event()

def_pause = 10
def_cycles = 3
def_loglevel = Loglevel.FINISH
def_agpath="./agents"
def_maxproc=2
repush=0

maxproc = def_maxproc
loglevel=def_loglevel
base_agents_path=os.path.abspath(def_agpath)
pause=def_pause
cycles=def_cycles

def print_help():
    print('\n',
'\nAgent multi-push script for Dwarfguard remote monitoring software, version ' + version,
'\nThis script expects a list of devices in a file. Each line of the file represents one device:',
'\ndevice_address:auth_method:auth_token:removal_mode',
'\n    <device_address> ... either IP or DNS name',
'    <auth_method> ... one of ( password | passfile | keyfile )',
'        password ... auth_token contains password text',
'        passfile ... auth_token contains path to file containing (only a) password',
'        keyfile  ... auth_token contains path to ssh private key',
'    <auth_token> ... see above',
'    <removal_mode> ... one of ( remove | noremove | uninstall)',
'        remove ... if agent there already, remove it and install new',
'        noremove ... if agent there already, do nothing',
'        uninstall ... special: never install, only remove if agent there already',
'\n    Examples:',
'\n        172.17.3.122:keyfile:/home/user/.ssh/id_rsa:remove',
'        myrouter:passfile:./myfile:remove',
'        themachine:password:lk23\:gK@2-:noremove',
'\n    Once the script runs, it creates two files:',
'        <input_filename>.ok ... all installed devices and assigned DevIDs',
'        <input_filename>.err ... all failed and skipped devices (on Ctrl-C)',
'\n    Script supports following parameters:',
'        <filepath> ... path to the input file. Mandatory parameter + must come first.',
'            Other parameters are optional:',
'        --jobs <x> ... number of installations running at one time [' + str(def_maxproc) + ']',
'        --repush ... force "remove" parameter for all',
'        --showerr ... show errors from subcommands [do not show]',
'        --agent_dir <path> ... agent archives are located in <path> [' + def_agpath + ']',
'        --loglevel <x> ... where x is in 0 .. 4 (NONE/ERROR/FINISH/INFO/DEBUG) [' + str(def_loglevel) + ']',
'        --wait <x> ... x is in seconds, x is the time wait for checking if agent registered [' + str(def_pause) + ']',
'        --cycles <x> ... x is number of checking cycles. If all passes and not registered => error [' + str(def_cycles) + ']\n', sep="\n"
)

def ctrl_c_handle(signum, frame):
    global terminate
    mprint_err("\nTerminating... (waiting for pushes in late stage, can take some time)")
    terminate.set()

def pg_msg(level, output=sys.stdout, *args):
    if level <= loglevel:
        if level == Loglevel.FINISH:
            print("\t==> ", *args, file=output, sep='')
        else:
            print(*args, file=output, sep='')

def mprint(level, *args):
    pg_msg(level, sys.stdout, *args)

def mprint_err(*args):
    pg_msg(Loglevel.ERROR, sys.stderr, *args)

def mprint_finerr(*args):
    pg_msg(Loglevel.FINISH, sys.stderr, "ERR: ", *args)

def mprint_finok(*args):
    pg_msg(Loglevel.FINISH, sys.stdout, "OK : ", *args)

def mprint_finun(*args):
    pg_msg(Loglevel.FINISH, sys.stdout, "PURGED : ", *args)

def target_cleanup(conncmd, agent_archive_fname, dirname):
    global errout
    acmd = conncmd + [ "[ -d " + agent_archive_fname + " ] && rm -r " + agent_archive_fname + " 2>/dev/null" ]
    if len(dirname) > 0:
        acmd[-1] = acmd[-1] + "; [ -d " + dirname + " ] && rm -r " + dirname + " 2>/dev/null"
    mprint(Loglevel.DEBUG, '\tCleanup: ', ' '.join(str(i) for i in acmd))
    res = subprocess.run(acmd, stdout=subprocess.PIPE, stderr=errout, text=True)

def proc_agent_push(addr, pwd, keyfpath, install_new = True, remove_old = False, pause = def_pause, cycles = def_cycles):
    global archive_path, errout
    device_type = Devtype.ANYDEVICE
    mprint(Loglevel.INFO, "Pushing agent to ", addr)
    cmds = [ 'nc', '-z', addr, '22' ]
    res = subprocess.run(cmds, stdout=subprocess.DEVNULL, stderr=errout, text=True)
    if 0 != res.returncode:
        return Retcode.ROUTE, ""
    if pwd:
        os.environ["SSHPASS"] = pwd
        cmds = [ 'sshpass', '-e', 'ssh', '-o', 'StrictHostKeyChecking=no' ]
        mprint(Loglevel.DEBUG, '\tSetting SSHPASS to ', pwd, '\n')
    else:
        cmds = [ 'ssh' , '-q', '-o', 'StrictHostKeyChecking=no', '-o', 'BatchMode=yes', '-i', keyfpath ]
    conncmd = cmds + [ 'root@' + addr ]
    acmd = conncmd + [ 'exit' ]
    mprint(Loglevel.DEBUG, '\tRunning: ', ' '.join(str(i) for i in acmd))
    res = subprocess.run(acmd, stdout=subprocess.PIPE, stderr=errout, text=True)
    if 0 != res.returncode:
        return Retcode.CONNECTION, ""
    acmd = conncmd + [ "if status lan >/dev/null 2>&1; then echo " + teststring + "; else echo ; fi" ]
    mprint(Loglevel.DEBUG, '\tRunning: ', ' '.join(str(i) for i in acmd))
    res = subprocess.run(acmd, stdout=subprocess.PIPE, stderr=errout, text=True)
    if 0 != res.returncode:
        return Retcode.CONNECTION, ""
    if res.stdout.strip() == teststring:
        device_type = Devtype.ADVROUTER
        archive_path = base_agents_path + "/agent_advantech/"+ agent_fname_prefix + "_advantech.tgz"
        archive_dirstring = agent_fname_prefix
        agent_archive_fname = agent_fname_prefix + '.' + agent_tar_ext
    else:
        acmd = conncmd + [ "if [ -e /etc/issue ] ; then echo " + teststring + "; else echo ; fi" ]
        mprint(Loglevel.DEBUG, '\tRunning: ', ' '.join(str(i) for i in acmd))
        res = subprocess.run(acmd, stdout=subprocess.PIPE, stderr=errout, text=True)
        if 0 != res.returncode:
            return Retcode.CONNECTION, ""
        if res.stdout.strip() == teststring:
            device_type = Devtype.LINUXBOX
            archive_path = base_agents_path + "/agent_linux/" + agent_fname_prefix + "_linux.tgz"
            archive_dirstring = agent_fname_prefix + "_linux"
            agent_archive_fname = agent_fname_prefix + '.' + agent_tar_ext
        else:
            archive_dirstring = ""
            acmd = conncmd + [ "if [ -e /etc/banner ] && grep -i openwrt /etc/banner >/dev/null ; then echo " + teststring + "; else echo ; fi" ]
            mprint(Loglevel.DEBUG, '\tRunning: ', ' '.join(str(i) for i in acmd))
            res = subprocess.run(acmd, stdout=subprocess.PIPE, stderr=errout, text=True)
            if 0 != res.returncode:
                return Retcode.CONNECTION, ""
            if res.stdout.strip() == teststring:
                device_type = Devtype.OPENWRT
                archive_path = base_agents_path + "/agent_owrt/" + agent_fname_prefix + "_owrt.ipk"
                agent_archive_fname = agent_fname_prefix + '.' + agent_ipk_ext
            else:
                acmd = conncmd + [ "if [ -e /etc/banner ] && grep -i teltonika /etc/banner >/dev/null ; then echo " + teststring + "; else echo ; fi" ]
                mprint(Loglevel.DEBUG, '\tRunning: ', ' '.join(str(i) for i in acmd))
                res = subprocess.run(acmd, stdout=subprocess.PIPE, stderr=errout, text=True)
                if 0 != res.returncode:
                    return Retcode.CONNECTION, ""
                if res.stdout.strip() == teststring:
                    device_type = Devtype.TELTROUTER
                    archive_path = base_agents_path + "/agent_teltonika/" + agent_fname_prefix + "_teltonika.ipk"
                    agent_archive_fname = agent_fname_prefix + '.' + agent_ipk_ext
                else:
                    print("Returned: " + res.stdout)
                    return Retcode.DEVTYPE, ""
    mprint(Loglevel.INFO, "\t", str(addr), ": Device type is ", str(device_type))
    acmd = conncmd + [ "if [ -e \"" + agent_dir + "\" ] ; then echo " + teststring + "; else echo ; fi " ]
    mprint(Loglevel.DEBUG, '\tRunning: ', ' '.join(str(i) for i in acmd))
    res = subprocess.run(acmd, stdout=subprocess.PIPE, stderr=errout, text=True)
    if 0 != res.returncode:
        return Retcode.CONNECTION, ""
    if res.stdout.strip() == teststring:
        mprint(Loglevel.INFO, "\t", str(addr), ": Existing agent detected")
        target_cleanup(conncmd, agent_archive_fname, archive_dirstring)
        if remove_old or repush:
            if device_type == Devtype.ADVROUTER:
                acmd = conncmd + [ "if umupdate -d " + agent_fname_prefix + " >/dev/null 2>&1; then echo " +teststring+ "; else echo ; fi " ]
                mprint(Loglevel.DEBUG, '\tRunning: ', ' '.join(str(i) for i in acmd))
                res = subprocess.run(acmd, stdout=subprocess.PIPE, stderr=errout, text=True)
                if 0 != res.returncode:
                    return Retcode.CONNECTION, ""
                if res.stdout.strip() != teststring:
                    return Retcode.REMOVAL, ""
            elif device_type == Devtype.OPENWRT or device_type == Devtype.TELTROUTER:
                acmd = conncmd + [ "opkg remove adwarfg >/dev/null 2>&1" ]
                mprint(Loglevel.DEBUG, '\tRunning: ', ' '.join(str(i) for i in acmd))
                res = subprocess.run(acmd, stdout=subprocess.PIPE, stderr=errout, text=True)
                if 0 != res.returncode:
                    return Retcode.CONNECTION, ""
                # retest - agent could have been installed outside of opkg
                acmd = conncmd + [ "if [ -e \"" + agent_dir + "\" ] ; then echo " + teststring + "; else echo ; fi " ]
                mprint(Loglevel.DEBUG, '\tRunning: ', ' '.join(str(i) for i in acmd))
                res = subprocess.run(acmd, stdout=subprocess.PIPE, stderr=errout, text=True)
                if 0 != res.returncode:
                    return Retcode.CONNECTION, ""
                if res.stdout.strip() == teststring:
                    mprint(Loglevel.INFO, "\t", str(addr), ": Agent installed outside of opkg, calling manual uninstall...")
                    acmd = conncmd + [ "if " +agent_dir+ "/uninstall.sh >/dev/null 2>&1; then echo " +teststring+ "; else echo ; fi " ]
                    mprint(Loglevel.DEBUG, '\tRunning: ', ' '.join(str(i) for i in acmd))
                    res = subprocess.run(acmd, stdout=subprocess.PIPE, stderr=errout, text=True)
                    if 0 != res.returncode:
                        return Retcode.CONNECTION, ""
                    if res.stdout.strip() != teststring:
                        return Retcode.REMOVAL, ""
            else:
                acmd = conncmd + [ "if " +agent_dir+ "/uninstall.sh >/dev/null 2>&1; then echo " +teststring+ "; else echo ; fi " ]
                mprint(Loglevel.DEBUG, '\tRunning: ', ' '.join(str(i) for i in acmd))
                res = subprocess.run(acmd, stdout=subprocess.PIPE, stderr=errout, text=True)
                if 0 != res.returncode:
                    return Retcode.CONNECTION, ""
                if res.stdout.strip() != teststring:
                    return Retcode.REMOVAL, ""
            mprint(Loglevel.INFO, "\t", str(addr), ": Existing agent was removed")
        else:
            return Retcode.REMOVAL, ""
    else:
        if not install_new:
            return Retcode.NOOLDAGENT, ""
    if not install_new:
        return Retcode.NONEWAGENT, ""
    if not os.access(archive_path, os.R_OK):
        mprint(Loglevel.ERROR, "Archive path ", archive_path, " is not a readable file.")
        return Retcode.ARCHIVE, ""
    acmd = conncmd + [ "cat >/root/" + agent_archive_fname ]
    archive = open(archive_path, "rb")
    mprint(Loglevel.DEBUG, '\tUsing agent archive from ', archive_path)
    mprint(Loglevel.DEBUG, '\tRunning: ', ' '.join(str(i) for i in acmd))
    res = subprocess.run(acmd, stdin=archive, stdout=subprocess.PIPE, stderr=errout, text=True)
    archive.close()
    if 0 != res.returncode:
        return Retcode.ARCHIVE, ""
    if device_type != Devtype.OPENWRT and device_type != Devtype.TELTROUTER:
        acmd = conncmd + [ "if tar xzf " + agent_archive_fname + " >/dev/null 2>&1; then echo "+teststring+"; else echo ; fi" ]
        mprint(Loglevel.DEBUG, '\tRunning: ', ' '.join(str(i) for i in acmd))
        res = subprocess.run(acmd, stdout=subprocess.PIPE, stderr=errout, text=True)
        if 0 != res.returncode:
            return Retcode.CONNECTION, ""
        if res.stdout.strip() != teststring:
            target_cleanup(conncmd, agent_archive_fname, archive_dirstring)
            return Retcode.UNCOMPRESS, ""
    if device_type == Devtype.ADVROUTER:
        acmd = conncmd + [ "if umupdate -a /root/" + agent_archive_fname + " >/dev/null 2>&1; then echo " + teststring + "; else echo ; fi" ]
    elif device_type == Devtype.OPENWRT or device_type == Devtype.TELTROUTER:
        acmd = conncmd + [ "if opkg install " + agent_archive_fname + " >/dev/null 2>&1; then echo " + teststring + "; else echo ; fi" ]
    else:
        acmd = conncmd + [ "if /root/" +archive_dirstring+ "/" + agent_fname + " >/dev/null 2>&1; then echo "+teststring+"; else echo ; fi" ]
    mprint(Loglevel.DEBUG, '\tRunning: ', ' '.join(str(i) for i in acmd))
    res = subprocess.run(acmd, stdout=subprocess.PIPE, stderr=errout, text=True)
    if 0 != res.returncode:
        return Retcode.CONNECTION, ""
    if res.stdout.strip() != teststring:
        target_cleanup(conncmd, agent_archive_fname, archive_dirstring)
        return Retcode.INSTALL, ""
    target_cleanup(conncmd, agent_archive_fname, archive_dirstring)
    mprint(Loglevel.INFO, "\t", str(addr), ": Agent was installed")
    acmd = conncmd + [ "[ -f "+agent_dir+"/settings.ini ] && grep \"^g_devid=\" "+agent_dir+"/settings.ini" ]
    for x in range(0, max(1,cycles)):
        time.sleep(pause)
        mprint(Loglevel.DEBUG, '\tRunning: ', ' '.join(str(i) for i in acmd))
        res = subprocess.run(acmd, stdout=subprocess.PIPE, stderr=errout, text=True)
        devid = res.stdout[8:].strip().strip('\"')
        if 0 == res.returncode and len(devid) > 0:
            break
        devid = ""
    if len(devid) == 0:
        return Retcode.DEVID, ""
    return 0, devid

def process_line(line, lineno, terminate):
    if terminate.is_set():
        return (Retcode.INTERRUPT, line, None)
    entry = [ token.replace('\:', ':') for token in re.split(r'(?<!\\):', line) ]
    if len(entry) < 4:
        mprint_finerr("line ", str(lineno), " [Syntax (skipping)]")
        return (Retcode.PARSE, line, None)
    if entry[1] == "password":
        passwd = entry[2]
        keypassfile = None
    elif entry[1] == "passfile":
        keypassfile = None
        if pwds_from_file.get(entry[2]) is not None:
            passwd = pwds_from_file.get(entry[2])
        else:
            try:
                with open(entry[2]) as f:
                    pwds_from_file[entry[2]] = f.readline().strip('\n')
            except:
                mprint_finerr(entry[0], " [Password file, line ", str(lineno), " \"", entry[2], "\"]")
                return (Retcode.PARSE, line, None)
            passwd = pwds_from_file[entry[2]]
    elif entry[1] == "keyfile":
        if not os.path.isfile(entry[2]):
            mprint_finerr(entry[0], " [Keyfile, line ", str(lineno), " \"", entry[2], "\"]")
            return (Retcode.PARSE, line, None)
        passwd = None
        keypassfile = entry[2]
    else:
        mprint_finerr(entry[0], " [Bad auth_metod, line ", str(lineno), " \"", entry[1], "\"]")
        return (Retcode.PARSE, line, None)
    removal = True
    installation = True
    if entry[3] == "remove":
        pass
    elif entry[3] == "noremove":
        removal = False
    elif entry[3] == "uninstall":
        installation = False
    else:
        mprint_finerr(entry[0], " [Bad removal mode, line ", str(lineno), " \"", entry[3], "\"]")
        return (Retcode.PARSE, line, None)
    auth_token_print = re.sub(r'[:]', r'\\:', entry[2])
    reslinefrag = str(entry[0]+":"+entry[1]+":"+auth_token_print+":"+entry[3])
    if (ret := proc_agent_push(entry[0], passwd, keypassfile, installation, removal, pause, cycles))[0]:
        if Retcode.ROUTE == ret[0]:
            mprint_finerr(entry[0], " [Unreachable]");
        elif Retcode.CONNECTION == ret[0]:
            mprint_finerr(entry[0], " [Communication/Login]");
        elif Retcode.DEVTYPE == ret[0]:
            mprint_finerr(entry[0], " [Device-type detection]");
        elif Retcode.REMOVAL == ret[0]:
            mprint_finerr(entry[0], " [Unable to remove agent]");
        elif Retcode.ARCHIVE == ret[0]:
            mprint_finerr(entry[0], " [Agent archive not found]");
        elif Retcode.UNCOMPRESS == ret[0]:
            mprint_finerr(entry[0], " [Uncompressing agent]");
        elif Retcode.INSTALL == ret[0]:
            mprint_finerr(entry[0], " [Agent install]");
        elif Retcode.DEVID == ret[0]:
            mprint_finerr(entry[0], " [Device registration]");
        elif Retcode.NOOLDAGENT == ret[0]:
            mprint_finun(entry[0], " [N/A]");
        elif Retcode.NONEWAGENT == ret[0]:
            mprint_finun(entry[0], " [Uninstalled]");
        else:
            mprint_finerr(entry[0], " [Code: ", str(ret[0]), "]")
        return (ret[0], line, None)
    else:
        # ok (retval is 0, we have devid in second list element
        mprint_finok(entry[0], " [", ret[1], "]")
        return (ret[0], str(reslinefrag+":"+ret[1]), ret[1])

def initworker():
    signal.signal(signal.SIGINT, signal.SIG_IGN)
    signal.signal(signal.SIGTERM, signal.SIG_IGN)
    signal.signal(signal.SIGHUP, signal.SIG_IGN)

def main():
    global terminate, loglevel, maxproc, errout
    for sig in ('TERM', 'HUP', 'INT'):
        signal.signal(getattr(signal, 'SIG'+sig), ctrl_c_handle)
    try:
        with open(filename) as f:
            lines = [ line.strip() for line in f ]
    except OSError:
        mprint_err("Failed to read file \"", filename, "\"")
    ctr=0
    results = []
    todolist = []
    for line in lines:
        ctr += 1
        todolist.append((line, ctr, terminate))
    pool = mp.Pool(processes = maxproc, initializer=initworker)
    results = []
    results = pool.starmap(process_line, todolist)
    print("Devices:", str(len(lines)))
    out_e = open(filename+".err", "w")
    out_ok = open(filename+".ok", "w")
    cnt_ok = 0
    cnt_err = 0
    cnt_skip = 0
    for i in results:
        if i[0]:
            out_e.write(i[1]+"\n")
            if i[0] == Retcode.INTERRUPT:
                cnt_skip += 1
            else:
                cnt_err += 1
        else:
            out_ok.write(i[1]+"\n")
            cnt_ok += 1
    print("\tSuccess:", str(cnt_ok), "\n\tFailures:", str(cnt_err), "\n\tSkipped:", str(cnt_skip))
    out_e.close()
    out_ok.close()
    pool.close()
    pool.join()

if __name__ == "__main__":
    if len(sys.argv) < 2:
        mprint_err("Expecting at least one argument - path to file, see help below.")
        print_help()
        sys.exit(1)
    if not os.access(sys.argv[1], os.R_OK):
        mprint_err("The provided file name (", sys.argv[1], ") is not available. See help below.")
        print_help()
        sys.exit(1)
    filename = sys.argv[1]
    try:
        opts, args = getopt.gnu_getopt(sys.argv[2:], 'h', [ 'help', 'jobs=', 'agent_dir=', 'loglevel=', 'wait=', 'cycles=', 'showerrs' ])
        for opt, arg in opts:
            if opt == "-h" or opt == "--help":
                print_help()
                sys.exit(0)
            if opt == "--jobs":
                maxproc = int(arg)
                if maxproc > os.cpu_count() * 4:
                    maxproc = os.cpu_count()* 4 
                    mprint_err("Limiting number of processes to ", maxproc)
            if opt == "--repush":
                repush = 1
            if opt == "--agent_dir":
                base_agents_path = os.path.abspath(arg)
            if opt == "--loglevel":
                if int(arg) >= Loglevel.NONE and int(arg) <= Loglevel.DEBUG:
                    loglevel = Loglevel(int(arg))
                else:
                    mprint_err("Invalid loglevel set (", arg, "). Must be between ", Loglevel.NONE, ' and ', Loglevel.DEBUG)
            if opt == "--wait":
                pause = int(arg)
            if opt == "--cycles":
                cycles = int(arg)
            if opt == "--showerrs":
                errout=sys.stderr
    except getopt.GetoptError:
        print("Error when parsing cmdline arguments from: ", ' '.join(sys.argv))
        print_help()
        sys.exit(1)
    main()