#!/usr/bin/env python3
"""
Script to extract ip_prefix values from AWS IP ranges JSON file
and output in sing-box rule set format
"""

import json
import sys
from pathlib import Path

def extract_ip_prefixes(json_file_path):
    """
    Read the JSON file and extract all ip_prefix values into an array
    
    Args:
        json_file_path (str): Path to the JSON file
        
    Returns:
        list: Array of ip_prefix strings
    """
    try:
        # Read the JSON file
        with open(json_file_path, 'r', encoding='utf-8') as file:
            data = json.load(file)
        
        # Extract ip_prefix values from the prefixes array
        ip_prefixes = []
        
        # Check if 'prefixes' key exists and extract ip_prefix values
        if 'prefixes' in data:
            for prefix_obj in data['prefixes']:
                if 'ip_prefix' in prefix_obj:
                    ip_prefixes.append(prefix_obj['ip_prefix'])
        
        # Also check for ipv6_prefixes if they exist
        if 'ipv6_prefixes' in data:
            for prefix_obj in data['ipv6_prefixes']:
                if 'ipv6_prefix' in prefix_obj:
                    ip_prefixes.append(prefix_obj['ipv6_prefix'])
        
        return ip_prefixes
        
    except FileNotFoundError:
        print(f"Error: File '{json_file_path}' not found.")
        return []
    except json.JSONDecodeError as e:
        print(f"Error: Invalid JSON format - {e}")
        return []
    except Exception as e:
        print(f"Error: {e}")
        return []

def output_singbox_format(ip_prefixes, output_file=None):
    """
    Output IP prefixes in sing-box rule set format
    
    Args:
        ip_prefixes (list): List of IP prefix strings
        output_file (str): Optional output file path
    """
    # Create the sing-box format JSON
    singbox_rule = {
        "ip_cidr": ip_prefixes
    }
    
    # Convert to JSON string with proper formatting
    json_output = json.dumps(singbox_rule, indent=2, ensure_ascii=False)
    
    if output_file:
        # Save to file
        with open(output_file, 'w', encoding='utf-8') as f:
            f.write(json_output)
        print(f"Sing-box rule set saved to: {output_file}")
    else:
        # Print to console
        print(json_output)

def main():
    """Main function"""
    # Default file path
    json_file = "./ip-ranges.json"
    output_file = "aws_ip_cidr.json"
    
    # Check if file path is provided as command line argument
    if len(sys.argv) > 1:
        json_file = sys.argv[1]
    
    # Check if output file is provided as second argument
    if len(sys.argv) > 2:
        output_file = sys.argv[2]
    
    # Check if file exists
    if not Path(json_file).exists():
        print(f"Error: File '{json_file}' does not exist.")
        print("Usage: python extract_ip_prefixes.py [input_json_file] [output_file]")
        sys.exit(1)
    
    # Extract ip_prefixes
    print(f"Reading IP prefixes from: {json_file}")
    ip_prefixes = extract_ip_prefixes(json_file)
    
    if ip_prefixes:
        print(f"Found {len(ip_prefixes)} IP prefixes")
        
        # Output in sing-box format
        output_singbox_format(ip_prefixes, output_file)
        
        # Also show first few entries as preview
        print(f"\nPreview (first 5 entries):")
        for i, prefix in enumerate(ip_prefixes[:5]):
            print(f"  {prefix}")
        if len(ip_prefixes) > 5:
            print(f"  ... and {len(ip_prefixes) - 5} more")
        
    else:
        print("No IP prefixes found in the JSON file.")

if __name__ == "__main__":
    main()