Notify remotely after running local command

notify-after.py

#!/usr/bin/env python3

# prerequisites:
# - sshexec.py
# - notify-send package on remote mechine

import argparse
import sshexec
import getpass
import subprocess

def run_process(exe):
    p = subprocess.Popen(exe, shell=True, universal_newlines=True,
                           stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    while(True):
        retcode = p.poll()
        line = p.stdout.readline()
        yield line
        if retcode!=None:
            for line in p.stdout.readlines():
                yield line
            break;

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Execute command and notify remotely.')
    parser.add_argument('user_host', type=str, help='user@host')
    parser.add_argument('message', type=str, help='notification message')
    parser.add_argument('cmd', type=str, help='the command to be executed')
    args = parser.parse_args()
    
    # check user@host
    try:
        user, host = args.user_host.split('@')
    except:
        print('Invalid user@host:', args.user_host)
        exit()
    
    # get and check password
    password = getpass.getpass('Remote password: ')
    if(not sshexec.try_connect(host, user, password)):
        print('Authentication failed, please check your login information')
        exit()
    
    # run the local command
    print('\33[34m[notify-after]: work begin...\33[0m')
    for line in run_process(args.cmd):
        print(line, end='')
    print('\33[34m[notify-after]: work end.\33[0m')
    
    # notify
    sshexec.execute(host, user, password, 'notify-send -u critical '+args.message)