"""SSH tunnel helper.

If Environment == dev, it initializes a SSH tunnel to a remote host which
have access to psql database.
"""

from functools import wraps
import os

from sshtunnel import SSHTunnelForwarder

from sme_labelcopy_loader import config


def ssh_tunnel(func):
    """Initialize a SSH tunnel if Environment == dev.

    Args:
        func (func): unload_labelcopy_data function.

    Returns:
        func: Decorated unload_labelcopy_data function.
    """
    @wraps(func)
    def wrapper(date, host, port):
        if os.environ.get('Environment') == 'dev':
            with SSHTunnelForwarder(
                    (config.ssh_host, config.ssh_port),
                    ssh_username=config.ssh_username,
                    ssh_private_key=config.ssh_private_key,
                    remote_bind_address=(host, port)) as server:

                return func(
                    date, server.local_bind_host, server.local_bind_port)

        else:
            return func(date, host, port)

    return wrapper