import socket
try:
    from io import BytesIO as StringIO
except ImportError:
    from StringIo import StringIO

from paramiko.client import SSHClient, AutoAddPolicy

from cisco_olt_client.exceptions import ConnectionError


class OltClient:

    def __init__(
        self,
        hostname,
        username,
        password,
        auto_connect=True,
        recv_timeout=1,
        recv_chunk_size=4096,
        missing_host_key_policy_cls=AutoAddPolicy,
        ssh_connection_options=None
    ):
        self.hostname = hostname
        self.username = username
        self.password = password
        self.auto_connect = auto_connect
        self.recv_timeout = recv_timeout
        self.recv_chunk_size = recv_chunk_size
        self.missing_host_key_policy_cls = missing_host_key_policy_cls
        self.shell = None
        self.ssh_client = None
        self.ssh_connection_options = dict(
            look_for_keys=False,
            allow_agent=False,
            timeout=5
        ) if ssh_connection_options is None else ssh_connection_options

    def get_ssh_client(self):
        '''
        Prepare default ssh client and set missing host key policy
        '''
        ssh_client = SSHClient()
        ssh_client.set_missing_host_key_policy(
            self.missing_host_key_policy_cls())
        return ssh_client

    def get_olt_shell(self, ssh_client, recv_msg=True):
        '''
        Initiate OLT interactive cli interface

        :param bool recv_msg: Indicates that the login message (everything that
            is printed by OLT on login) should be read so it will not appear in
            output of subsequent command.

        :return: OLT interactive cli interface in form of paramiko's channel
        :rtype: :class:`paramiko.Channel`
        :raises paramiko.AuthenticationException:
        '''
        ssh_client.connect(
            self.hostname,
            username=self.username,
            password=self.password,
            **self.ssh_connection_options
        )
        shell = ssh_client.invoke_shell()
        shell.settimeout(self.recv_timeout)
        if recv_msg:
            exec_command(shell, '')
        return shell

    def connect(self):
        # the ssh_client reference is needed otherwise it will get garbage
        # collected and causing the shell to close immediately
        self.ssh_client = self.get_ssh_client()
        self.shell = self.get_olt_shell(self.ssh_client)

    def disconnect(self):
        self.shell.close()
        self.ssh_client.close()

    def raw_exec_command(self, cmd, **exec_options):
        if self.shell is None:
            if self.auto_connect:
                self.connect()
            else:
                raise ConnectionError(
                    'Connectio not ready and auto_connect flag is False')
        return exec_command(self.shell, cmd, **exec_options)

    def __enter__(self):
        self.connect()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.disconnect()


def exec_command(
    shell,
    cmd,
    buffer=None,
    recv_chunk_size=4096,
    output_as_file=False,
    new_line=True,
    new_line_char='\n'
):
    buffer = StringIO() if buffer is None else buffer
    if not cmd.endswith(new_line_char):
        cmd += new_line_char
    shell.send(cmd)
    while True:
        try:
            buffer.write(shell.recv(recv_chunk_size))
        except socket.timeout:
            if shell.closed:
                raise ConnectionError('Connection to OTL failed')
            break
    return buffer if output_as_file else buffer.getvalue()

@@ -160,16 +160,61 @@
 Policy%0A%0A
+from cisco_olt_client.command import Command%0A
 from cis
@@ -2575,16 +2575,761 @@
 lose()%0A%0A
+    def execute(self, cmd, args=None):%0A        '''%0A        Use the client to execute given command%0A%0A        If given %60cmd%60 is already instance of%0A        :class:%60~cisco_olt_client.command.Command%60 it'll be executed using the%0A        clien, otherwise %60cmd%60 and %60args%60 will be used to create an instance%0A        for you.%0A%0A        :param cmd: Command to execute%0A        :type cmd: :class:%60~cisco_olt_client.command.Command%60 or str%0A%0A        :param args: Additional command arguments. In python2, if you need to%0A            preserve order of arguments use ordered dict or list of 2-tuples%0A        :type args: list or dict%0A        '''%0A        if not instance(cmd, Command):%0A            cmd = Command(cmd, args=args)%0A        return cmd.execute(self)%0A%0A%0A
     def 

Add execute function on client