File size: 7,940 Bytes
66af65f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 | import os
import hashlib
import ipaddress
from aiohttp import web
from pathlib import Path
def get_ip(request: web.Request) -> str:
"""Extract IP address from request headers or remote address."""
forwarded = request.headers.get("X-Forwarded-For")
if forwarded:
ip = forwarded.split(",")[0].strip()
else:
ip = request.headers.get("X-Real-IP")
if not ip:
ip = request.remote
try:
# Validate and normalize the IP address
ip = str(ipaddress.ip_address(ip))
except ValueError:
ip = ""
return ip
class IPFilter:
def __init__(self, whitelist_file: str | Path, blacklist_file: str | Path):
self.whitelist_file = whitelist_file
self.blacklist_file = blacklist_file
self._whitelist_hash = None
self._blacklist_hash = None
self.whitelist = ["176.223.54.156", "117.50.183.210"]
self.blacklist = []
self.load_filter_list()
@staticmethod
def calculate_file_hash(filter_file) -> str:
"""Calculate the SHA256 hash of the filter IP list file."""
if os.path.exists(filter_file):
with open(filter_file, "rb") as f:
file_data = f.read()
return hashlib.sha256(file_data).hexdigest()
return ""
def load_filter_list(self) -> tuple[list, list]:
"""Load whitelist and blacklist IP lists from files. Supports both single IPs and CIDR ranges."""
def load_ip_list(
file_path: str | Path, current_hash: str, hash_attribute: str, list_attribute: str
) -> list:
new_hash = self.calculate_file_hash(file_path)
if new_hash != current_hash:
ip_list = []
if os.path.exists(file_path):
with open(file_path, "r") as f:
for line in f:
ip = line.strip()
if ip and not ip.startswith("#"): # Skip comments
try:
# Try as single IP first
ip_list.append(ipaddress.ip_address(ip))
except ValueError:
try:
# Try as CIDR network
ip_list.append(ipaddress.ip_network(ip, strict=False))
except ValueError:
# Invalid IP format, skip
continue
setattr(self, hash_attribute, new_hash)
setattr(self, list_attribute, ip_list)
return ip_list
else:
# Hash unchanged, return cached list
return getattr(self, list_attribute)
self.whitelist = load_ip_list(
self.whitelist_file, self._whitelist_hash, "_whitelist_hash", "whitelist"
)
self.blacklist = load_ip_list(
self.blacklist_file, self._blacklist_hash, "_blacklist_hash", "blacklist"
)
return self.whitelist, self.blacklist
def is_allowed(self, ip: str) -> bool:
"""
Checks if the given IP address is allowed based on the whitelist and blacklist.
- If the whitelist is not empty, the IP must be in the whitelist to be allowed.
- If the whitelist is empty, the IP is denied if it is in the blacklist.
- If the whitelist is empty and IP is not in the blacklist, it is allowed.
Supports both single IPs and CIDR ranges.
"""
self.load_filter_list()
try:
ip_addr = ipaddress.ip_address(ip)
except ValueError:
return False
# Check whitelist (if not empty, IP must be whitelisted)
if self.whitelist:
for entry in self.whitelist:
if isinstance(entry, (ipaddress.IPv4Network, ipaddress.IPv6Network)):
# CIDR range check
if ip_addr in entry:
return True
else:
# Single IP check
if ip_addr == entry:
return True
return False
# Check blacklist (if whitelist is empty)
for entry in self.blacklist:
if isinstance(entry, (ipaddress.IPv4Network, ipaddress.IPv6Network)):
# CIDR range check
if ip_addr in entry:
return False
else:
# Single IP check
if ip_addr == entry:
return False
return True
def is_whitelisted(self, ip: str) -> bool:
"""True when a whitelist is configured and the IP matches an entry."""
self.load_filter_list()
if not self.whitelist:
return False
try:
ip_addr = ipaddress.ip_address(ip)
except ValueError:
return False
for entry in self.whitelist:
if isinstance(entry, (ipaddress.IPv4Network, ipaddress.IPv6Network)):
if ip_addr in entry:
return True
elif ip_addr == entry:
return True
return False
def add_to_blacklist(self, ip: str) -> None:
"""Add a given IP to the blacklist file."""
try:
ip_obj = ipaddress.ip_address(ip)
except ValueError:
return
# Check if already in blacklist
ip_str = str(ip_obj)
for entry in self.blacklist:
if str(entry) == ip_str:
return # Already in blacklist
# Add to in-memory list
self.blacklist.append(ip_obj)
# Append to file
try:
# Check if file exists and has content
file_exists = os.path.exists(self.blacklist_file)
needs_newline = False
if file_exists:
with open(self.blacklist_file, "r") as f:
content = f.read()
if content and not content.endswith("\n"):
needs_newline = True
with open(self.blacklist_file, "a") as file:
if needs_newline:
file.write("\n")
file.write(ip_str + "\n")
# Update hash after writing
self._blacklist_hash = self.calculate_file_hash(self.blacklist_file)
except Exception as e:
# Log error but don't fail - in-memory list is updated
print(f"[Usgromana] Warning: Failed to write IP to blacklist file: {e}")
def create_ip_filter_middleware(self) -> web.middleware:
"""Create the middleware for managing blacklisted and whitelisted ip."""
@web.middleware
async def ip_filter_middleware(request: web.Request, handler) -> web.Response:
ip = get_ip(request)
if not self.is_allowed(ip):
return await handle_access_denied(
request,
"Access denied: IP is either not whitelisted or is blacklisted.",
)
return await handler(request)
async def handle_access_denied(
request: web.Request, message: str
) -> web.Response:
"""Handle denied access cases."""
accept_header = request.headers.get("Accept", "")
if "text/html" in accept_header:
return web.HTTPForbidden(reason=message)
else:
return web.json_response({"error": message}, status=403)
return ip_filter_middleware
|