#!/usr/bin/env python3


from os import path, getpid
from subprocess import Popen, PIPE
from sys import argv, platform


class Proc:
    """
    Find a running process from Python
    """
    @classmethod
    def _list_windows(cls):
        """
        Find all running process with wmi
        :return: list of dictionaries with descriptions of found processes
        """
        execlist = []
        separate = b'\r\r\n'
        out, err = Popen(['wmic', 'process', 'get', 'CommandLine,ExecutablePath,Name,ProcessId', '/format:list'],
                         stdout=PIPE, stderr=PIPE).communicate()
        for line in out.split(separate + separate):
            execpid, exename, exepath, cmdline = None, None, None, None
            for subline in line.split(separate):
                if b'ProcessId=' in subline:
                    execpid = subline.split(b'=')[1].decode('utf-8')
                if b'Name=' in subline:
                    exename = subline.split(b'=')[1].decode('utf-8')
                if b'ExecutablePath=' in subline:
                    exepath = subline.split(b'=')[1].decode('utf-8')
                if b'CommandLine=' in subline:
                    cmdline = subline.split(b'=')[1].decode('utf-8')
            if execpid and exename:
                execlist.append({'execpid': execpid, 'exename': exename, 'exepath': exepath, 'cmdline': cmdline})
        return execlist

    @classmethod
    def _list_linux(cls):
        """
        Find all running process with ps
        :return: list of dictionaries with descriptions of found processes
        """
        execlist = []
        out, err = Popen(['/bin/ps', '-eo', 'pid,args'], stdout=PIPE, stderr=PIPE).communicate()
        for line in out.splitlines():
            execpid = line.split()[0].decode('utf-8')
            exepath = line.split()[1].decode('utf-8')
            exename = path.basename(exepath)
            cmdline = line.split(None, 1)[1].decode('utf-8')
            if execpid and exename:
                execlist.append({'execpid': execpid, 'exename': exename, 'exepath': exepath, 'cmdline': cmdline})
        return execlist

    @classmethod
    def list(cls):
        """
        Find all running process
        :return: list of dictionaries with descriptions of found processes
        """
        if platform.startswith('linux') or platform.startswith('darwin'):
            return cls._list_linux()
        elif platform.startswith('win32'):
            return cls._list_windows()
        else:
            return None

    @classmethod
    def search(cls, find: str, exclude: str = None):
        """
        Find specified processes
        :param find: find process pid, name or arguments
        :param exclude: exclude process pid, name or arguments
        :return: list of dictionaries with descriptions of found processes
        """
        proc_found = []
        try:
            for proc in cls.list():
                if exclude and (exclude in proc['execpid'] or exclude in proc['exename'] or
                                exclude in proc['exepath'] or exclude in proc['cmdline']):
                    pass
                elif find in proc['execpid'] or find in proc['exename'] or \
                        find in proc['exepath'] or find in proc['cmdline']:
                    proc_found.append(proc)
        except TypeError as ex:
            print('ON', platform, 'PLATFORM', 'search ERROR:', ex)
        finally:
            if len(proc_found) == 0:
                return None
            else:
                return proc_found

    @classmethod
    def kill(cls, pid: int):
        """
        Kill the process by means of the OS
        :param pid: process ID
        :return: None
        """
        if platform.startswith('linux') or platform.startswith('darwin'):
            Popen(['kill', '-s', 'SIGKILL', str(pid)])
        elif platform.startswith('win32'):
            Popen(['taskkill', '/PID', str(pid), '/F'])


if __name__ == "__main__":
    from argparse import ArgumentParser

    args = ArgumentParser(
        prog='Proc',
        description='Find a running process from Python',
        epilog='Dependencies: Python 3 (tested version 3.9.5)'
    )
    args.add_argument('--find', type=str, required=False,
                      help='find process pid, name or arguments')
    args.add_argument('--exclude', type=str, default=None, required=False,
                      help='exclude process pid, name or arguments')
    args.add_argument('--self', action='store_true', required=False,
                      help='find a clones of self')
    args.add_argument('--kill', type=int, required=False,
                      help='kill the process with pid')
    args = vars(args.parse_args())
    processes = None
    if args['kill']:
        Proc.kill(args['kill'])
    elif args['find']:
        processes = Proc.search(args['find'], args['exclude'])
    elif args['self']:
        processes = Proc.search(' '.join(argv), str(getpid()))
    else:
        processes = Proc.list()
    if processes:
        for process in processes:
            print(process)