"""
Linux on Hyper-V and Azure Test Code, ver. 1.0.0
Copyright (c) Microsoft Corporation

All rights reserved
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

See the Apache Version 2.0 License for specific language governing
permissions and limitations under the License.
"""
import os
import time
import paramiko
import socket


class SSHClient(object):
    """
    This class creates a paramiko.SSHClient() object that represents
    a session with an SSH server. You can use the SSHClient object to send
    commands to the remote host and manipulate files on the remote host.

    :param server: A server hostname or ip.
    :param host_key_file: The path to the user's .ssh key files.
    :param user: The username for the SSH connection. Default = 'root'.
    :param timeout: The optional timeout variable for the TCP connection.
    :param ssh_pwd: An optional password to use for authentication or for
                    unlocking the private key.
    :param ssh_key_file: SSH key pem data
    """

    def __init__(self, server, host_key_file='~/.ssh/known_hosts', user='root', timeout=None,
                 ssh_pwd=None, ssh_key_file=None):
        self.server = server
        self.host_key_file = host_key_file
        self.user = user
        self._timeout = timeout
        self._pkey = paramiko.RSAKey.from_private_key_file(ssh_key_file, password=ssh_pwd)
        self._ssh_client = paramiko.SSHClient()
        self._ssh_client.load_system_host_keys()
        self._ssh_client.load_host_keys(os.path.expanduser(host_key_file))
        self._ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        self.connect()

    def connect(self, num_retries=5):
        """
        Connect to an SSH server and authenticate with it.
        :type num_retries: int
        :param num_retries: The maximum number of connection attempts.
        """
        retry = 0
        while retry < num_retries:
            try:
                self._ssh_client.connect(self.server,
                                         username=self.user,
                                         pkey=self._pkey,
                                         timeout=self._timeout)
                return
            except socket.error as xxx_todo_changeme:
                (value, message) = xxx_todo_changeme.args
                if value in (51, 61, 111):
                    print('SSH Connection refused, will retry in 5 seconds')
                    time.sleep(5)
                    retry += 1
                else:
                    raise
            except paramiko.BadHostKeyException:
                print("{} has an entry in ~/.ssh/known_hosts and it doesn't match".format(
                        self.server))
                print('Edit that file to remove the entry and then hit return to try again')
                raw_input('Hit Enter when ready')
                retry += 1
            except EOFError: