from ipaddress import ip_address

_split = lambda x: x.split(',') if x else []


class ProxyFix(object):
    """This is a slightly modified version of werkzeug's ProxyFix.

    Instead of using a fixed number of proxies it uses a list of trusted IP
    addresses.

    This middleware can be applied to add HTTP proxy support to an
    application that was not designed with HTTP proxies in mind.  It
    sets `REMOTE_ADDR`, `HTTP_HOST` from `X-Forwarded` headers.

    The original values of `REMOTE_ADDR` and `HTTP_HOST` are stored in
    the WSGI environment as `werkzeug.proxy_fix.orig_remote_addr` and
    `werkzeug.proxy_fix.orig_http_host`.

    :param app: the WSGI application
    :param num_proxies: the number of proxy servers in front of the app.
    """

    def __init__(self, app):
        self.wsgi = app.wsgi_app
        proxies = app.config.get('PROXYFIX_TRUSTED', ['127.0.0.1'])
        self.trusted = frozenset(ip_address(p.strip()) for p in proxies)

    def get_remote_addr(self, forwarded_for, remote):
        """Selects the new remote addr from the given list of ips in
        X-Forwarded-For. It picks the "rightmost" untrusted IP address.
        """
        for ip_str in reversed([remote] + forwarded_for):
            ip_str = ip_str.strip()
            try:
                if not ip_address(ip_str) in self.trusted:
                    return ip_str
            except ValueError:
                return ip_str
        return remote

    def update_environ(self, environ):
        env = environ.get
        remote_addr = env('REMOTE_ADDR')
        if not remote_addr:
            return

        try:
            remote_addr_ip = ip_address(remote_addr)
        except ValueError:
            remote_addr_ip = ip_address('127.0.0.1')

        forwarded_proto = env('HTTP_X_FORWARDED_PROTO', '')
        forwarded_for = _split(env('HTTP_X_FORWARDED_FOR', ''))
        forwarded_host = env('HTTP_X_FORWARDED_HOST', '')
        environ.update({
            'werkzeug.proxy_fix.orig_wsgi_url_scheme': env('wsgi.url_scheme'),
            'werkzeug.proxy_fix.orig_remote_addr':     env('REMOTE_ADDR'),
            'werkzeug.proxy_fix.orig_http_host':       env('HTTP_HOST')
        })

        if remote_addr_ip in self.trusted:
            if forwarded_host:
                environ['HTTP_HOST'] = forwarded_host
            if forwarded_proto:
                https = 'https' in forwarded_proto.lower()
                environ['wsgi.url_scheme'] = 'https' if https else 'http'

        remote_addr = self.get_remote_addr(forwarded_for, remote_addr)
        if remote_addr is not None:
            environ['REMOTE_ADDR'] = remote_addr

    def __call__(self, environ, start_response):
        self.update_environ(environ)
        return self.wsgi(environ, start_response)
