core/homeassistant/components/http/real_ip.py

41 lines
1.2 KiB
Python

"""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_proxies):
"""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
try:
if (use_x_forwarded_for and
X_FORWARDED_FOR in request.headers and
any(connected_ip in trusted_proxy
for trusted_proxy in trusted_proxies)):
request[KEY_REAL_IP] = ip_address(
request.headers.get(X_FORWARDED_FOR).split(', ')[-1])
except ValueError:
pass
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)