import sys
import paramiko
'''
a wrapper for remote shell
usage:
python rsh.py user host password 'commands'
e.g.
python rsh.py root 192.168.192.206 123456 'cd home/quant_group;ls -a'
'''

if __name__ == '__main__':
    recv_nbytes = 1024

    args = sys.argv
    host, port, user, password = args[1], int(args[2]), args[3], args[4]
    command_begin = 4
    if '--tag=' in args[4]:
        tag = args[4][6:] + ":"
        command_begin += 1
    else:
        tag = ''
        remain_args = ''
    for arg in args[command_begin:]:
        if ' ' not in arg:
            remain_args += ' ' + arg
        else:
            remain_args += ' "' + arg + '" '
    print tag + remain_args

    def print_err(str_value):
        sys.stderr.write(tag + str_value)
        sys.stderr.flush()

    def print_out(str_value):
        sys.stdout.write(tag + str_value)
        sys.stdout.flush()

    client = paramiko.SSHClient()
    client.load_system_host_keys()
    client.get_transport()
    client.set_missing_host_key_policy(paramiko.WarningPolicy())
    client.connect(host, port, user, password)
    transport = client.get_transport()
    channel = transport.open_session(None)
    channel.settimeout(None)
    channel.exec_command(remain_args)
    while True:
        if channel.recv_ready():
            ret_str = channel.recv(recv_nbytes)
            if len(ret_str) == 0:
                break
            print_out(ret_str)
        if channel.recv_stderr_ready():
            ret_str = channel.recv_stderr(recv_nbytes)
            if len(ret_str) == 0:
                break
            print_err(ret_str)
        if channel.exit_status_ready():
            ret_str = channel.recv(recv_nbytes)
            if len(ret_str) != 0:
                print_out(ret_str)
            ret_str = channel.recv_stderr(recv_nbytes)
            if len(ret_str) != 0:
                print_err(ret_str)
            client.close()
            sys.exit(channel.recv_exit_status())
