129 lines
3.7 KiB
Python
129 lines
3.7 KiB
Python
#!/usr/bin/env python3
|
|
"""WHOIS subnet querying and IP consolidation."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import ipaddress
|
|
import json
|
|
import os
|
|
import re
|
|
from typing import Iterable
|
|
|
|
from .constants import (
|
|
SUBNET_THRESHOLD,
|
|
WHOIS_CACHE_PATH,
|
|
WHOIS_PORT,
|
|
WHOIS_SERVER,
|
|
WHOIS_TIMEOUT,
|
|
)
|
|
from .ip_utils import ip_version, parse_ip
|
|
|
|
_whois_cache: dict[str, str] = {}
|
|
|
|
|
|
def load_whois_cache() -> None:
|
|
global _whois_cache
|
|
if not _whois_cache and os.path.exists(WHOIS_CACHE_PATH):
|
|
try:
|
|
with open(WHOIS_CACHE_PATH, "r", encoding="utf-8") as f:
|
|
_whois_cache = json.load(f)
|
|
except Exception:
|
|
_whois_cache = {}
|
|
|
|
|
|
def save_whois_cache() -> None:
|
|
if _whois_cache:
|
|
os.makedirs(os.path.dirname(WHOIS_CACHE_PATH), exist_ok=True)
|
|
try:
|
|
with open(WHOIS_CACHE_PATH, "w", encoding="utf-8") as f:
|
|
json.dump(_whois_cache, f, indent=2)
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
def query_whois_subnet(ip: str) -> str | None:
|
|
"""Query whois.radb.net for the route/subnet of an IP address."""
|
|
addr = parse_ip(ip)
|
|
if addr is None:
|
|
return None
|
|
|
|
load_whois_cache()
|
|
if ip in _whois_cache:
|
|
return _whois_cache[ip] or None
|
|
|
|
import socket
|
|
try:
|
|
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
|
s.settimeout(WHOIS_TIMEOUT)
|
|
s.connect((WHOIS_SERVER, WHOIS_PORT))
|
|
s.sendall(f"{ip}\r\n".encode())
|
|
res = ""
|
|
while True:
|
|
chunk = s.recv(4096).decode(errors="ignore")
|
|
if not chunk:
|
|
break
|
|
res += chunk
|
|
s.close()
|
|
|
|
routes = re.findall(r"^(?:route|route6):\s*(\S+)",
|
|
res, re.IGNORECASE | re.MULTILINE)
|
|
if routes:
|
|
routes_parsed = []
|
|
for r in routes:
|
|
try:
|
|
routes_parsed.append(ipaddress.ip_network(r))
|
|
except ValueError:
|
|
continue
|
|
if routes_parsed:
|
|
best_route = max(routes_parsed, key=lambda n: n.prefixlen)
|
|
subnet_str = str(best_route)
|
|
_whois_cache[ip] = subnet_str
|
|
save_whois_cache()
|
|
return subnet_str
|
|
|
|
except Exception:
|
|
pass
|
|
|
|
subnet_str = f"{ip}/32" if addr.version == 4 else f"{ip}/128"
|
|
_whois_cache[ip] = subnet_str
|
|
save_whois_cache()
|
|
return subnet_str
|
|
|
|
|
|
def consolidate_by_subnet(ips: Iterable[str],
|
|
threshold: int = SUBNET_THRESHOLD) -> set[str]:
|
|
"""Group IPs by WHOIS subnet; if >= threshold IPs in a subnet, block the subnet.
|
|
|
|
Otherwise, block individual IPs.
|
|
"""
|
|
ip_list = list(ips)
|
|
subnet_to_ips: dict[str, set[str]] = {}
|
|
ip_to_subnet: dict[str, str] = {}
|
|
|
|
for ip in ip_list:
|
|
subnet = query_whois_subnet(ip)
|
|
if subnet:
|
|
try:
|
|
net = ipaddress.ip_network(subnet, strict=False)
|
|
addr = ipaddress.ip_address(ip)
|
|
if addr in net:
|
|
subnet_to_ips.setdefault(subnet, set()).add(ip)
|
|
ip_to_subnet[ip] = subnet
|
|
continue
|
|
except ValueError:
|
|
pass
|
|
fallback_net = f"{ip}/32" if ip_version(ip) == 4 else f"{ip}/128"
|
|
subnet_to_ips.setdefault(fallback_net, set()).add(ip)
|
|
ip_to_subnet[ip] = fallback_net
|
|
|
|
blocked_targets: set[str] = set()
|
|
for subnet, mapped_ips in subnet_to_ips.items():
|
|
if len(mapped_ips) >= threshold:
|
|
blocked_targets.add(subnet)
|
|
else:
|
|
for ip in mapped_ips:
|
|
blocked_targets.add(
|
|
f"{ip}/32" if ip_version(ip) == 4 else f"{ip}/128")
|
|
|
|
return blocked_targets
|