import paramiko
import re


class ShellHandler:

    def __init__(self, host, user, psw):
        self.ssh = paramiko.SSHClient()
        self.ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        self.ssh.connect(host, username=user, password=psw, port=22)

        channel = self.ssh.invoke_shell()
        self.stdin = channel.makefile('wb')
        self.stdout = channel.makefile('r')

    def __del__(self):
        self.ssh.close()

    def execute(self, cmd):
        """

        :param cmd: the command to be executed on the remote computer
        :examples:  execute('ls')
                    execute('finger')
                    execute('cd folder_name')
        """
        cmd = cmd.strip('\n')
        self.stdin.write(cmd + '\n')
        finish = 'end of stdOUT buffer. finished with exit status'
        echo_cmd = 'echo {} $?'.format(finish)
        self.stdin.write(echo_cmd + '\n')
        shin = self.stdin
        self.stdin.flush()

        shout = []
        sherr = []
        exit_status = 0
        for line in self.stdout:
            if str(line).startswith(cmd) or str(line).startswith(echo_cmd):
                # up for now filled with shell junk from stdin
                shout = []
            elif str(line).startswith(finish):
                # our finish command ends with the exit status
                exit_status = int(str(line).rsplit(maxsplit=1)[1])
                if exit_status:
                    # stderr is combined with stdout.
                    # thus, swap sherr with shout in a case of failure.
                    sherr = shout
                    shout = []
                break
            else:
                # get rid of 'coloring and formatting' special characters
                shout.append(re.compile(r'(\x9B|\x1B\[)[0-?]*[ -/]*[@-~]').sub('', line).
                             replace('\b', '').replace('\r', ''))


s = ShellHandler('linux-lab-055.ece.uw.edu', 'hyu3', ':Eric200002182919')
s.execute("vncserver -list")


# ssh = paramiko.SSHClient()
# ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
# # linux-lab-055.ece.uw.edu
# ssh.connect("linux-lab-055.ece.uw.edu", username="hyu3", password=":Eric200002182919")
# ssh_stdin, ssh_stdout, ssh_stderr = ssh.exec_command("vncserver -list")
#
# xdisplay = ssh_stdout.readlines()[-1]
#
# if "DISPLAY" in xdisplay:
#     ssh_stdin, ssh_stdout, ssh_stderr = ssh.exec_command('vncserver -interface 0.0.0.0 && vncserver -list')
#     xdisplay = ssh_stdout.readlines()[-1]
#
# print(xdisplay.split()[0][1:])