#!/usr/bin/env python3
"""
check_nft_blacklist.py
Checks if an IP address is in the nftables blacklist configuration.
Usage:
  check_nft_blacklist.py nft_bl.conf 192.168.1.1
  check_nft_blacklist.py nft_bl.conf 2001:db8::1
"""

import sys
import re
from ipaddress import ip_address, ip_network, AddressValueError
from pathlib import Path

def iter_set_blocks(content):
    current_name = None
    current_lines = []
    brace_depth = 0

    for line in content.splitlines():
        if current_name is None:
            match = re.match(r"\s*set\s+([A-Za-z0-9_]+)\s*\{", line)
            if match:
                current_name = match.group(1)
                current_lines = [line]
                brace_depth = line.count("{") - line.count("}")
            continue

        current_lines.append(line)
        brace_depth += line.count("{") - line.count("}")
        if brace_depth == 0:
            yield current_name, "\n".join(current_lines)
            current_name = None
            current_lines = []

def parse_nft_config(config_path):
    """Extract IPv4 and IPv6 prefixes from nftables config."""
    p = Path(config_path)
    if not p.exists():
        raise FileNotFoundError(f"Config file not found: {config_path}")
    
    content = p.read_text(encoding="utf-8")
    v4_prefixes = []
    v6_prefixes = []

    for _, block in iter_set_blocks(content):
        if "type ipv4_addr" in block:
            for match in re.finditer(r"(\d+\.\d+\.\d+\.\d+(?:/\d+)?)", block):
                try:
                    v4_prefixes.append(ip_network(match.group(1), strict=False))
                except Exception as e:
                    print(f"Warning: Could not parse IPv4 prefix '{match.group(1)}': {e}", file=sys.stderr)
        elif "type ipv6_addr" in block:
            for match in re.finditer(r"([0-9a-fA-F:]+(?:/\d+)?)", block):
                try:
                    v6_prefixes.append(ip_network(match.group(1), strict=False))
                except Exception:
                    pass
    
    return v4_prefixes, v6_prefixes

def check_ip_in_blacklist(ip_addr, v4_prefixes, v6_prefixes):
    """Check if IP address is in any of the blacklist prefixes."""
    try:
        addr = ip_address(ip_addr)
    except AddressValueError as e:
        raise ValueError(f"Invalid IP address: {ip_addr} ({e})")
    
    prefixes = v4_prefixes if addr.version == 4 else v6_prefixes
    
    for prefix in prefixes:
        if addr in prefix:
            return True, prefix
    
    return False, None

def main(argv):
    if len(argv) < 3:
        print("Usage: python3 check_nft_blacklist.py <nft_config.conf> <ip_address>")
        print("Examples:")
        print("  check_nft_blacklist.py nft_bl.conf 192.168.1.1")
        print("  check_nft_blacklist.py nft_bl.conf 2001:db8::1")
        return 2
    
    config_file = argv[1]
    ip_to_check = argv[2]
    
    # Parse the nftables config
    try:
        print(f"Loading blacklist from: {config_file}")
        v4_prefixes, v6_prefixes = parse_nft_config(config_file)
        print(f"Loaded {len(v4_prefixes)} IPv4 prefixes and {len(v6_prefixes)} IPv6 prefixes")
    except Exception as e:
        print(f"ERROR: Could not parse config file: {e}", file=sys.stderr)
        return 3
    
    # Check if IP is in blacklist
    try:
        is_blocked, matching_prefix = check_ip_in_blacklist(ip_to_check, v4_prefixes, v6_prefixes)
        
        print(f"\nChecking IP: {ip_to_check}")
        print("-" * 50)
        
        if is_blocked:
            print(f"✗ BLOCKED - IP is in blacklist")
            print(f"  Matching prefix: {matching_prefix}")
            return 1
        else:
            print(f"✓ OK - IP is NOT in blacklist")
            return 0
            
    except ValueError as e:
        print(f"ERROR: {e}", file=sys.stderr)
        return 4

if __name__ == "__main__":
    sys.exit(main(sys.argv))