#!/usr/bin/env python3
"""
IP2Location ASN & Geolocation Subnet Parser
Designed for research purposes. 

This script cross-references the IP2LOCATION-LITE-ASN and IP2LOCATION-LITE-DB11
databases to extract all CIDR subnets belonging to a specific country code (default: RU).
It outputs only the CIDR subnets to a plain text file.
"""

import csv
import bisect
import argparse
import logging
import sys
from typing import List, Tuple

# Configure logging for research tracking
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S"
)

def load_geolocation_data(db11_path: str) -> Tuple[List[int], List[Tuple[int, str]]]:
    """
    Loads the DB11 database into memory for fast binary search.
    
    Args:
        db11_path (str): Path to the IP2LOCATION-LITE-DB11.CSV file.
        
    Returns:
        Tuple containing:
            - A sorted list of starting IPs (ip_from).
            - A corresponding list of tuples (ip_to, country_code).
    """
    ip_starts: List[int] = []
    geo_data: List[Tuple[int, str]] = []
    
    logging.info(f"Loading Geolocation database: {db11_path}")
    
    try:
        with open(db11_path, mode='r', encoding='utf-8') as f:
            reader = csv.reader(f)
            for row_num, row in enumerate(reader, start=1):
                try:
                    ip_from = int(row[0])
                    ip_to = int(row[1])
                    country_code = row[2]
                    
                    ip_starts.append(ip_from)
                    geo_data.append((ip_to, country_code))
                except (ValueError, IndexError):
                    logging.warning(f"Skipping malformed row {row_num} in DB11.")
                    
    except FileNotFoundError:
        logging.error(f"File not found: {db11_path}")
        sys.exit(1)

    logging.info(f"Successfully loaded {len(ip_starts):,} geolocation records.")
    return ip_starts, geo_data

def extract_country_subnets(
    asn_path: str, 
    ip_starts: List[int], 
    geo_data: List[Tuple[int, str]], 
    target_country: str, 
    output_path: str
) -> None:
    """
    Parses the ASN database, cross-references with DB11 data, and saves matching CIDRs to a TXT file.
    
    Args:
        asn_path (str): Path to the IP2LOCATION-LITE-ASN.CSV file.
        ip_starts (List[int]): Sorted list of starting IPs from DB11.
        geo_data (List[Tuple[int, str]]): Corresponding DB11 data (ip_to, country_code).
        target_country (str): ISO 3166-1 alpha-2 country code (e.g., 'RU').
        output_path (str): Filepath for the resulting TXT file.
    """
    logging.info(f"Parsing ASN database ({asn_path}) for country code: '{target_country}'...")
    
    match_count = 0
    total_count = 0

    try:
        with open(asn_path, mode='r', encoding='utf-8') as infile, \
             open(output_path, mode='w', encoding='utf-8') as outfile:
             
            reader = csv.reader(infile)
            
            for row in reader:
                total_count += 1
                try:
                    ip_from = int(row[0])
                    cidr = row[2]
                    
                    # Binary search to find the matching geolocation block
                    # bisect_right returns the insertion point. The actual range is at index - 1
                    idx = bisect.bisect_right(ip_starts, ip_from) - 1
                    
                    if idx >= 0:
                        db_ip_to, country_code = geo_data[idx]
                        
                        # Verify that the IP falls within the bounds and matches the country
                        if ip_from <= db_ip_to and country_code == target_country:
                            outfile.write(f"{cidr}\n")
                            match_count += 1
                            
                except (ValueError, IndexError):
                    logging.warning(f"Skipping malformed row {total_count} in ASN DB.")

    except FileNotFoundError:
        logging.error(f"File not found: {asn_path}")
        sys.exit(1)

    logging.info(f"Processing complete. Parsed {total_count:,} ASN records.")
    logging.info(f"Found {match_count:,} subnets assigned to '{target_country}'.")
    logging.info(f"Results saved to: {output_path}")

def main():
    parser = argparse.ArgumentParser(
        description="Cross-reference IP2Location ASN and DB11 databases to extract specific country CIDRs."
    )
    
    parser.add_argument(
        "--asn", 
        required=True, 
        help="Path to the IP2LOCATION-LITE-ASN.CSV file"
    )
    parser.add_argument(
        "--db11", 
        required=True, 
        help="Path to the IP2LOCATION-LITE-DB11.CSV file"
    )
    parser.add_argument(
        "--output", 
        default="russian_subnets.txt", 
        help="Output TXT filename (default: russian_subnets.txt)"
    )
    parser.add_argument(
        "--country", 
        default="RU", 
        help="ISO 3166-1 alpha-2 country code to filter by (default: RU)"
    )
    
    args = parser.parse_args()

    # Step 1: Load DB11 into memory
    ip_starts, geo_data = load_geolocation_data(args.db11)
    
    # Step 2: Parse ASN, filter by country, and output to TXT
    extract_country_subnets(args.asn, ip_starts, geo_data, args.country.upper(), args.output)

if __name__ == "__main__":
    main()