#!/usr/bin/python
import ipaddress
import yaml
import os
from collections import defaultdict

def load_yaml_input(file_path):
    with open(file_path, "r") as f:
        return yaml.safe_load(f)

def group_by_subnet(data, mask=24):
    grouped = defaultdict(list)
    
    for record in data:
        ip_str = record.get("address")
        if not ip_str:
            continue
        try:
            ip = ipaddress.ip_address(ip_str)
            network = ipaddress.ip_network(f"{ip}/{mask}", strict=False)
            grouped[str(network)].append(record)
        except ValueError:
            print(f"Skipping invalid IP: {ip_str}")
    
    return grouped

def save_groups_to_yaml(grouped, output_dir="."):
    os.makedirs(output_dir, exist_ok=True)

    for subnet, entries in grouped.items():
        safe_subnet = subnet.replace("/", "_")
        filename = os.path.join(output_dir, f"subnet_{safe_subnet}.yaml")
        with open(filename, "w") as f:
            yaml.dump(entries, f, default_flow_style=False)
        print(f"Saved {len(entries)} entries to {filename}")

def main():
    input_file = "input.yaml"  # change as needed
    output_dir = "output_subnets"
    cidr_mask = 24  # change to desired subnet size

    records = load_yaml_input(input_file)
    grouped = group_by_subnet(records, cidr_mask)
    save_groups_to_yaml(grouped, output_dir)

if __name__ == "__main__":
    main()