from unittest import TestCase
from unittest.mock import mock_open, patch

from paramiko import ProxyCommand

from pyinfra.connectors.sshuserclient import SSHClient
from pyinfra.connectors.sshuserclient.client import AskPolicy, get_ssh_config

SSH_CONFIG_DATA = '''
# Comment
Host 127.0.0.1
    IdentityFile /id_rsa
    IdentityFile /id_rsa2
    User testuser
    Port 33
    ProxyCommand ssh thing

Include other_file
'''

SSH_CONFIG_OTHER_FILE = '''
Host 192.168.1.1
    User "otheruser"
    ProxyCommand None
    ForwardAgent yes
    UserKnownHostsFile ~/.ssh/test3
'''

SSH_CONFIG_OTHER_FILE_PROXYJUMP = '''
Host 192.168.1.2
    User "otheruser"
    ProxyJump nottestuser@127.0.0.1
    ForwardAgent yes
'''

BAD_SSH_CONFIG_DATA = '''
&
'''

LOOPING_SSH_CONFIG_DATA = '''
Include other_file
'''


class TestSSHUserConfigMissing(TestCase):
    def setUp(self):
        get_ssh_config.cache = {}

    @patch(
        'pyinfra.connectors.sshuserclient.client.path.exists',
        lambda path: False,
    )
    def test_load_ssh_config_no_exist(self):
        client = SSHClient()

        _, config, forward_agent, missing_host_key_policy, host_keys_file \
            = client.parse_config('127.0.0.1')

        assert config.get('port') == 22


@patch(
    'pyinfra.connectors.sshuserclient.client.path.exists',
    lambda path: True,
)
@patch(
    'pyinfra.connectors.sshuserclient.config.glob.iglob',
    lambda path: ['other_file'],
)
@patch(
    'pyinfra.connectors.sshuserclient.config.path.isfile',
    lambda path: True,
)
@patch(
    'pyinfra.connectors.sshuserclient.config.path.expanduser',
    lambda path: path,
)
@patch(
    'pyinfra.connectors.sshuserclient.config.path.isabs',
    lambda path: True,
)
@patch(
    'paramiko.config.LazyFqdn.__str__',
    lambda self: '',
)
class TestSSHUserConfig(TestCase):
    def setUp(self):
        get_ssh_config.cache = {}

    @patch(
        'pyinfra.connectors.sshuserclient.client.open',
        mock_open(read_data=SSH_CONFIG_DATA),
        create=True,
    )
    @patch(
        'pyinfra.connectors.sshuserclient.config.open',
        mock_open(read_data=SSH_CONFIG_OTHER_FILE),
        create=True,
    )
    def test_load_ssh_config(self):
        client = SSHClient()

        _, config, forward_agent, missing_host_key_policy, host_keys_file \
            = client.parse_config('127.0.0.1')

        assert config.get('key_filename') == ['/id_rsa', '/id_rsa2']
        assert config.get('username') == 'testuser'
        assert config.get('port') == 33
        assert isinstance(config.get('sock'), ProxyCommand)
        assert forward_agent is False
        assert isinstance(missing_host_key_policy, AskPolicy)
        assert host_keys_file == '~/.ssh/known_hosts'  # OpenSSH default

        _, other_config, forward_agent, missing_host_key_policy, host_keys_file \
            = client.parse_config('192.168.1.1')

        assert other_config.get('username') == 'otheruser'
        assert forward_agent is True
        assert isinstance(missing_host_key_policy, AskPolicy)
        assert host_keys_file == '~/.ssh/test3'

    @patch(
        'pyinfra.connectors.sshuserclient.client.open',
        mock_open(read_data=BAD_SSH_CONFIG_DATA),
        create=True,
    )
    def test_invalid_ssh_config(self):
        client = SSHClient()

        with self.assertRaises(Exception) as context:
            client.parse_config('127.0.0.1')

        assert context.exception.args[0] == 'Unparsable line &'

    @patch(
        'pyinfra.connectors.sshuserclient.client.open',
        mock_open(read_data=LOOPING_SSH_CONFIG_DATA),
        create=True,
    )
    @patch(
        'pyinfra.connectors.sshuserclient.config.open',
        mock_open(read_data=LOOPING_SSH_CONFIG_DATA),
        create=True,
    )
    def test_include_loop_ssh_config(self):
        client = SSHClient()

        with self.assertRaises(Exception) as context:
            client.parse_config('127.0.0.1')

        assert context.exception.args[0] == 'Include loop detected in ssh config file: other_file'

    @patch(
        'pyinfra.connectors.sshuserclient.client.open',
        mock_open(read_data=SSH_CONFIG_DATA),
        create=True,
    )
    @patch(
        'pyinfra.connectors.sshuserclient.config.open',
        mock_open(read_data=SSH_CONFIG_OTHER_FILE_PROXYJUMP),
        create=True,
    )
    @patch('pyinfra.connectors.sshuserclient.SSHClient.connect')
    @patch('pyinfra.connectors.sshuserclient.SSHClient.gateway')
    def test_load_ssh_config_proxyjump(self, fake_gateway, fake_ssh_connect):
        client = SSHClient()

        # Load the SSH config with ProxyJump configured
        _, config, forward_agent, _, _ = client.parse_config(
            '192.168.1.2',
            {'port': 1022},
            ssh_config_file='other_file',
        )

        fake_ssh_connect.assert_called_once_with(
            '127.0.0.1',
            _pyinfra_ssh_config_file='other_file',
            port='33',
            sock=None,
            username='nottestuser',
        )
        fake_gateway.assert_called_once_with('192.168.1.2', 1022, '192.168.1.2', 1022)
