"""Middleware to fetch real IP."""

from ipaddress import ip_address

from aiohttp.web import middleware
from aiohttp.hdrs import X_FORWARDED_FOR

from homeassistant.core import callback

from .const import KEY_REAL_IP


@callback
def setup_real_ip(app, use_x_forwarded_for, trusted_networks):
    """Create IP Ban middleware for the app."""
    @middleware
    async def real_ip_middleware(request, handler):
        """Real IP middleware."""
        connected_ip = ip_address(
            request.transport.get_extra_info('peername')[0])
        request[KEY_REAL_IP] = connected_ip

        # Only use the XFF header if enabled, present, and from a trusted proxy
        if (use_x_forwarded_for and
                X_FORWARDED_FOR in request.headers and
                any(connected_ip in trusted_network
                    for trusted_network in trusted_networks)):
            request[KEY_REAL_IP] = ip_address(
                request.headers.get(X_FORWARDED_FOR).split(',')[0])

        return await handler(request)

    async def app_startup(app):
        """Initialize bans when app starts up."""
        app.middlewares.append(real_ip_middleware)

    app.on_startup.append(app_startup)