import paramiko

class Ssh:
    def __init__(self, server, username="root"):
        self.ssh_connection = None
        self.server = server
        self.username = username

    def __enter__(self):
        self.ssh_connection = paramiko.SSHClient()
        self.ssh_connection.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        self.ssh_connection.connect(self.server, port=2201, username=self.username)
        return self

    def __exit__(self, tp, e, traceback):
        self.ssh_connection.close()

    def check_output(self, command):
        ssh_stdin, ssh_stdout, ssh_stderr = self.ssh_connection.exec_command(command)
        return ssh_stdout.read().decode()

    def check_return_code(self, command):
        ssh_stdin, ssh_stdout, ssh_stderr = self.ssh_connection.exec_command(command)
        return ssh_stdout.channel.recv_exit_status()