from ftplib import error_perm
import os
from unittest.mock import MagicMock
from unittest.mock import mock_open
from unittest.mock import patch

from boto.s3.key import Key
from garcon import activity  # noqa
from garcon.contrib import ftp  # noqa

import paramiko

import pytest


US_SFTP_CREDS = dict(
    username='user',
    host='host',
    port=22)

CANADA_FTP_CREDS = dict(
    username='user',
    password='password',
    host='host',
    port=21)

US_FILE_NAMES_LIST = [
    'TRACK_ORCH.txt.gz', 'PRODUCT_TRACK_ORCH.txt.gz', 'PRODUCT_ORCH.txt.gz']

CANADA_FILE_NAMES_LIST = [
    'ORCH_Canada_Track.txt.gz',
    'ORCH_Canada_Product.txt.gz']

US_PRIVATE_KEY_PATH = './id_rsa_soundscan_us'


def test_upload_from_local_to_ftp(monkeypatch):
    # Mock sftp client object
    sftp_obj = MagicMock()
    monkeypatch.setattr(sftp_obj, 'chdir', MagicMock(return_value=''))
    monkeypatch.setattr(sftp_obj, 'mkdir', MagicMock(return_value=''))
    monkeypatch.setattr(sftp_obj, 'put', MagicMock(return_value=''))
    monkeypatch.setattr(sftp_obj, 'close', MagicMock(return_value=''))
    monkeypatch.setattr(
        paramiko.SFTPClient, 'from_transport',
        MagicMock(return_value=sftp_obj))

    # Mock transport object
    transport_obj = MagicMock()
    monkeypatch.setattr(transport_obj, 'connect', MagicMock(return_value=''))
    monkeypatch.setattr(transport_obj, 'close', MagicMock(return_value=''))
    monkeypatch.setattr(
        paramiko, 'Transport', MagicMock(return_value=transport_obj))

    # Mock paramiko.RSAKey object
    monkeypatch.setattr(
        paramiko.RSAKey, 'from_private_key_file',
        MagicMock(return_value=''))

    ftp._upload_from_local_to_ftp(
        US_SFTP_CREDS, 'test', 'test_remote_dir', 'test_local_dir',
        'test_key_path')

    assert not sftp_obj.mkdir.called
    assert sftp_obj.chdir.called
    assert sftp_obj.put.called
    assert sftp_obj.close.called

    # Test if IOError exception thrown
    monkeypatch.setattr(
        sftp_obj, 'chdir',
        MagicMock(return_value='', side_effect=IOError('Boom!')))

    with pytest.raises(IOError):
        ftp._upload_from_local_to_ftp(
            US_SFTP_CREDS, 'test', 'test_remote_dir', 'test_local_dir',
            'test_key_path')
        assert sftp_obj.mkdir.called

    # Mock regular ftp client object
    ftp_obj = MagicMock()
    monkeypatch.setattr(ftp, 'FTP', MagicMock(return_value=ftp_obj))

    m = mock_open()
    with patch('builtins.open', m, create=True):
        ftp._upload_from_local_to_ftp(
            CANADA_FTP_CREDS, 'test', 'test_remote_dir', 'test_local_dir')

        assert ftp_obj.mkd.called
        assert ftp_obj.close.called
        assert ftp_obj.storbinary.called


@patch('os.stat')
def test_download_from_sftp(mock_stat, monkeypatch):
    mock_stat().st_size = 5

    copyfile_obj = MagicMock()
    monkeypatch.setattr(ftp.shutil, 'copyfileobj', copyfile_obj)

    # Mock sftp client object
    sftp_obj = MagicMock()
    sftp_obj.open = mock_open()
    sftp_obj.stat.return_value.st_size = 5
    monkeypatch.setattr(
        paramiko.SFTPClient, 'from_transport',
        MagicMock(return_value=sftp_obj))

    # Mock transport object
    transport_obj = MagicMock()
    monkeypatch.setattr(
        paramiko, 'Transport', MagicMock(return_value=transport_obj))

    # Mock paramiko.RSAKey object
    monkeypatch.setattr(
        paramiko.RSAKey, 'from_private_key_file',
        MagicMock(return_value='test_rsa_key'))

    open_local_obj = mock_open()

    with patch('garcon.contrib.ftp.open', open_local_obj, create=True):
        # assert sftp call with a RSA key & password
        ftp._download_from_sftp(
            US_SFTP_CREDS, 'test', 'test_remote_dir', 'test_local_dir',
            'test_key_path')

        transport_obj.connect.assert_called_with(
            username='user',
            password=None,
            pkey='test_rsa_key')

        sftp_obj.chdir.assert_called_with('test_remote_dir')
        open_local_obj.assert_called_with('test_local_dir/test', 'wb')
        sftp_obj.open.assert_called_with('test')
        copyfile_obj.assert_called_with(sftp_obj.open(), open_local_obj())
        sftp_obj.close.assert_called_with()

        # assert sftp call with just a password (no RSA key)
        transport_obj.reset_mock()
        sftp_creds = dict(
            username='user',
            password='password',
            host='host',
            port=22)
        ftp._download_from_sftp(
            sftp_creds, 'test', 'test_remote_dir', 'test_local_dir')
        transport_obj.connect.assert_called_with(
            username='user',
            password='password',
            pkey=None)

        # assert IOError raised if sftp file size mismatch occurs
        sftp_obj.stat.return_value.st_size = 7
        with pytest.raises(IOError):
            ftp._download_from_sftp(
                sftp_creds, 'test', 'test_remote_dir', 'test_local_dir')


def test_download_from_ftp(monkeypatch):
    # Mock regular ftp client object
    ftp_obj = MagicMock()
    monkeypatch.setattr(ftp, 'FTP', MagicMock(return_value=ftp_obj))

    m = mock_open()
    with patch('builtins.open', m, create=True):
        ftp._download_from_ftp(
            CANADA_FTP_CREDS, 'test', 'test_remote_dir', 'test_local_dir')

        assert ftp_obj.close.called
        assert ftp_obj.retrbinary.called


def test_extract_bucket_path(monkeypatch):
    mock_bucket = 'bucket_name'
    mock_buck_path = 'bucket_path'
    actual_bucket_name, actual_bucket_path = ftp._extract_bucket_path(
        's3://bucket_name/bucket_path')
    assert actual_bucket_name == mock_bucket
    assert actual_bucket_path == mock_buck_path


def test_copy_from_s3_to_ftp(monkeypatch):
    monkeypatch.setattr(
        ftp, '_extract_bucket_path', MagicMock(
            return_value=('bucket_name', 'bucket_path')))
    monkeypatch.setattr(
        ftp, '_upload_from_local_to_ftp', MagicMock(return_value=''))
    monkeypatch.setattr(
        os, 'remove', MagicMock(return_value=''))
    mock_key = MagicMock()
    monkeypatch.setattr(
        mock_key, 'get_contents_to_filename', MagicMock(return_value=''))

    with patch('garcon.contrib.ftp.Bucket') as mock_bucket:
        monkeypatch.setattr(
            mock_bucket, 'get_key', MagicMock(return_value=mock_key))
        ftp.copy_from_s3_to_ftp(
            activity.Activity(), US_SFTP_CREDS, US_FILE_NAMES_LIST,
            's3://bucket_name/bucket_path', '/orchard/soundscan_us/2015-10-01',
            None)
        ftp._upload_from_local_to_ftp.assert_any_call(
            US_SFTP_CREDS, 'TRACK_ORCH.txt.gz',
            '/orchard/soundscan_us/2015-10-01/', './', None)
        ftp._upload_from_local_to_ftp.assert_any_call(
            US_SFTP_CREDS, 'PRODUCT_TRACK_ORCH.txt.gz',
            '/orchard/soundscan_us/2015-10-01/', './', None)
        ftp._upload_from_local_to_ftp.assert_any_call(
            US_SFTP_CREDS, 'PRODUCT_ORCH.txt.gz',
            '/orchard/soundscan_us/2015-10-01/', './', None)

        assert os.remove.called


def test_copy_from_ftp_to_s3(monkeypatch):
    monkeypatch.setattr(
        ftp, '_extract_bucket_path',
        MagicMock(return_value=('bucket_name', 'bucket_path')))
    monkeypatch.setattr(ftp, '_download_from_sftp', MagicMock())
    monkeypatch.setattr(ftp, '_download_from_ftp', MagicMock())
    monkeypatch.setattr(os, 'remove', MagicMock(return_value=''))
    monkeypatch.setattr(
        Key, 'set_contents_from_filename', MagicMock(return_value=[Key]))

    with patch('garcon.contrib.ftp.Bucket'):
        # Test sftp calls made when sftp creds are provided
        response = ftp.copy_from_ftp_to_s3(
            activity.Activity(), US_SFTP_CREDS, US_FILE_NAMES_LIST,
            's3://bucket_name/bucket_path',
            '/orchard/soundscan_us/2015-10-01',
            'key')
        ftp._download_from_sftp.assert_any_call(
            US_SFTP_CREDS, 'TRACK_ORCH.txt.gz',
            '/orchard/soundscan_us/2015-10-01/', './', 'key')
        ftp._download_from_sftp.assert_any_call(
            US_SFTP_CREDS, 'PRODUCT_TRACK_ORCH.txt.gz',
            '/orchard/soundscan_us/2015-10-01/', './', 'key')
        ftp._download_from_sftp.assert_any_call(
            US_SFTP_CREDS, 'PRODUCT_ORCH.txt.gz',
            '/orchard/soundscan_us/2015-10-01/', './', 'key')

        assert len(response['files']) == 3
        for resp in response['files']:
            assert resp['status'] is True
            assert resp.get('exception') is None
        assert os.remove.called
        os.remove.reset_mock()

        # Test ftp calls made when ftp creds are provided
        response = ftp.copy_from_ftp_to_s3(
            activity.Activity(), CANADA_FTP_CREDS, CANADA_FILE_NAMES_LIST,
            's3://bucket_name/bucket_path',
            '/orchard/soundscan_canada/2015-10-01',
            None)

        ftp._download_from_ftp.assert_any_call(
            CANADA_FTP_CREDS, 'ORCH_Canada_Track.txt.gz',
            '/orchard/soundscan_canada/2015-10-01/', './')
        ftp._download_from_ftp.assert_any_call(
            CANADA_FTP_CREDS, 'ORCH_Canada_Product.txt.gz',
            '/orchard/soundscan_canada/2015-10-01/', './')

        assert len(response['files']) == 2
        for resp in response['files']:
            assert resp['status'] is True
            assert resp.get('exception') is None
        assert os.remove.called


def _download_from_ftp(
        sftp_creds, file_name, remote_ftp_dir, local_dir):
    raise error_perm('No download for you')


def _download_from_sftp(
        sftp_creds, file_name, remote_ftp_dir, local_dir, pkey):
    raise FileNotFoundError('No download for you')


def test_copy_from_ftp_to_s3_fails(monkeypatch):

    monkeypatch.setattr(
        ftp, '_extract_bucket_path', MagicMock(
            return_value=('bucket_name', 'bucket_path')))
    monkeypatch.setattr(
        ftp, '_download_from_ftp', _download_from_ftp)
    monkeypatch.setattr(
        ftp, '_download_from_sftp', _download_from_sftp)
    monkeypatch.setattr(
        os, 'remove', MagicMock(return_value=''))
    monkeypatch.setattr(
        Key, 'set_contents_from_filename', MagicMock(return_value=[Key]))

    with patch('garcon.contrib.ftp.Bucket'):
        response = ftp.copy_from_ftp_to_s3(
            activity.Activity(), US_SFTP_CREDS, US_FILE_NAMES_LIST,
            's3://bucket_name/bucket_path', '/orchard/soundscan_us/2015-10-01',
            None)

        assert len(response['files']) == 3
        for resp in response['files']:
            assert resp['status'] is False
            assert resp['exception'] is not None
            assert isinstance(resp['exception'], FileNotFoundError)
        assert not os.remove.called
        assert not Key.set_contents_from_filename.called

        response = ftp.copy_from_ftp_to_s3(
            activity.Activity(), CANADA_FTP_CREDS, CANADA_FILE_NAMES_LIST,
            's3://bucket_name/bucket_path', '/orchard/soundscan_us/2015-10-01',
            None)

        assert len(response['files']) == 2
        for resp in response['files']:
            assert resp['status'] is False
            assert resp['exception'] is not None
            assert isinstance(resp['exception'], error_perm)
        assert not os.remove.called
        assert not Key.set_contents_from_filename.called


@patch('garcon.contrib.ftp.paramiko')
def test_file_exist_on_sftp(mock_paramiko):
    mock_sftp_client = MagicMock()
    mock_sftp_client.stat.return_value = None
    mock_sftp = MagicMock()
    mock_sftp.from_transport.return_value = mock_sftp_client
    mock_transport = MagicMock()
    mock_transport.connect.return_value = 'connect is called'
    mock_paramiko.SFTPClient.from_transport.return_value = mock_sftp_client
    mock_paramiko.Transport.return_value = mock_transport

    sftp_creds = dict(
        host='host',
        port=22,
        username='user',
        password='password'
    )
    actual = ftp.file_exist_on_sftp('test_path', **sftp_creds)
    assert mock_sftp_client.stat.called
    assert actual


@patch('garcon.contrib.ftp.paramiko')
def test_file_exist_on_sftp_file_not_found(mock_paramiko):
    mock_sftp_client = MagicMock()
    mock_sftp_client.stat.return_value = None
    mock_sftp_client.stat.side_effect = FileNotFoundError('No such file')
    mock_sftp = MagicMock()
    mock_sftp.from_transport.return_value = mock_sftp_client
    mock_transport = MagicMock()
    mock_transport.connect.return_value = 'connect is called'
    mock_paramiko.SFTPClient.from_transport.return_value = mock_sftp_client
    mock_paramiko.Transport.return_value = mock_transport

    sftp_creds = dict(
        host='host',
        port=22,
        username='user',
        password='password'
    )
    actual = ftp.file_exist_on_sftp('test_path', **sftp_creds)
    assert mock_sftp_client.stat.called
    assert not actual


def test_copy_file_from_ftp_to_s3(monkeypatch):
    monkeypatch.setattr(
        ftp, '_extract_bucket_path',
        MagicMock(return_value=('bucket_name', 'bucket_path')))
    monkeypatch.setattr(ftp, '_download_from_sftp', MagicMock())
    monkeypatch.setattr(ftp, '_download_from_ftp', MagicMock())
    monkeypatch.setattr(os, 'remove', MagicMock(return_value=''))
    monkeypatch.setattr(
        Key, 'set_contents_from_filename', MagicMock(return_value=[Key]))
    activity = MagicMock()

    with patch('garcon.contrib.ftp.Bucket'):
        # Test sftp calls made when sftp creds are provided
        ftp_path = '/ftp/path'
        ftp_file_name = 'ftp_file'
        s3_path = 's3://bucket_name/bucket_path'
        s3_file_name = 's3_file'
        pkey = 'key'

        sftp_response = ftp.copy_file_from_ftp_to_s3(
            activity, US_SFTP_CREDS, ftp_path, ftp_file_name, s3_path,
            s3_file_name, pkey)
        ftp._download_from_sftp.assert_any_call(
            US_SFTP_CREDS, ftp_file_name, ftp_path, '.', pkey)

        assert sftp_response['status'] is True
        assert sftp_response['file'] == s3_file_name
        assert sftp_response.get('exception') is None
        assert os.remove.called
        os.remove.reset_mock()

        # Test ftp calls made when ftp creds are provided
        ftp_response = ftp.copy_file_from_ftp_to_s3(
            activity, CANADA_FTP_CREDS, ftp_path, ftp_file_name, s3_path,
            s3_file_name, None)

        ftp._download_from_ftp.assert_any_call(
            CANADA_FTP_CREDS, ftp_file_name, ftp_path, '.')

        assert ftp_response['status'] is True
        assert sftp_response['file'] == s3_file_name
        assert ftp_response.get('exception') is None
        assert os.remove.called

        # Test target file name
        file_response = ftp.copy_file_from_ftp_to_s3(
            activity, CANADA_FTP_CREDS, ftp_path, ftp_file_name, s3_path)

        assert file_response['file'] == ftp_file_name


def test_copy_file_from_ftp_to_s3_fails(monkeypatch):
    monkeypatch.setattr(
        ftp, '_extract_bucket_path',
        MagicMock(return_value=('bucket_name', 'bucket_path')))
    monkeypatch.setattr(ftp, '_download_from_sftp', _download_from_sftp)
    monkeypatch.setattr(ftp, '_download_from_ftp', _download_from_ftp)
    monkeypatch.setattr(os, 'remove', MagicMock(return_value=''))
    monkeypatch.setattr(
        Key, 'set_contents_from_filename', MagicMock(return_value=[Key]))
    activity = MagicMock()

    with patch('garcon.contrib.ftp.Bucket'):
        # Test sftp calls made when sftp creds are provided
        ftp_path = '/ftp/path'
        ftp_file_name = 'ftp_file'
        s3_path = 's3://bucket_name/bucket_path'
        s3_file_name = 's3_file'
        pkey = 'key'

        sftp_response = ftp.copy_file_from_ftp_to_s3(
            activity, US_SFTP_CREDS, ftp_path, ftp_file_name, s3_path,
            s3_file_name, pkey)

        assert sftp_response['status'] is False
        assert sftp_response['exception'] is not None
        assert isinstance(sftp_response['exception'], FileNotFoundError)
        assert not os.remove.called
        assert not Key.set_contents_from_filename.called

        # Test ftp calls made when ftp creds are provided
        ftp_response = ftp.copy_file_from_ftp_to_s3(
            activity, CANADA_FTP_CREDS, ftp_path, ftp_file_name, s3_path,
            s3_file_name, None)

        assert ftp_response['status'] is False
        assert ftp_response['exception'] is not None
        assert isinstance(ftp_response['exception'], error_perm)
        assert not os.remove.called
        assert not Key.set_contents_from_filename.called