| 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:
|
|
|
| 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("#"):
|
| try:
|
|
|
| ip_list.append(ipaddress.ip_address(ip))
|
| except ValueError:
|
| try:
|
|
|
| ip_list.append(ipaddress.ip_network(ip, strict=False))
|
| except ValueError:
|
|
|
| continue
|
| setattr(self, hash_attribute, new_hash)
|
| setattr(self, list_attribute, ip_list)
|
| return ip_list
|
| else:
|
|
|
| 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
|
|
|
|
|
| if self.whitelist:
|
| for entry in self.whitelist:
|
| if isinstance(entry, (ipaddress.IPv4Network, ipaddress.IPv6Network)):
|
|
|
| if ip_addr in entry:
|
| return True
|
| else:
|
|
|
| if ip_addr == entry:
|
| return True
|
| return False
|
|
|
|
|
| for entry in self.blacklist:
|
| if isinstance(entry, (ipaddress.IPv4Network, ipaddress.IPv6Network)):
|
|
|
| if ip_addr in entry:
|
| return False
|
| else:
|
|
|
| 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
|
|
|
|
|
| ip_str = str(ip_obj)
|
| for entry in self.blacklist:
|
| if str(entry) == ip_str:
|
| return
|
|
|
|
|
| self.blacklist.append(ip_obj)
|
|
|
|
|
| try:
|
|
|
| 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")
|
|
|
|
|
| self._blacklist_hash = self.calculate_file_hash(self.blacklist_file)
|
| except Exception as e:
|
|
|
| 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
|
|
|
