#!/usr/bin/env python3
"""
NetBox DHCP Pool Setup Script
Marks IPs .100-.150 in each subnet as DHCP pool range
"""

import requests
import ipaddress
import sys
import os

# Configuration
NETBOX_URL = "https://netbox.pdmarf.co.uk"
NETBOX_TOKEN = os.environ.get("NETBOX_TOKEN") or sys.argv[1] if len(sys.argv) > 1 else None

if not NETBOX_TOKEN:
    print("Error: NetBox API token required")
    print("Usage: NETBOX_TOKEN='your-token' python3 sys-apps-netbox-dhcp-setup.py")
    print("   or: python3 sys-apps-netbox-dhcp-setup.py 'your-token'")
    sys.exit(1)

# API headers
headers = {
    "Authorization": f"Token {NETBOX_TOKEN}",
    "Content-Type": "application/json",
    "Accept": "application/json"
}

def get_all_prefixes():
    """Get all IPv4 prefixes from NetBox"""
    url = f"{NETBOX_URL}/api/ipam/prefixes/"
    params = {"family": 4, "limit": 1000}  # IPv4 only

    try:
        response = requests.get(url, headers=headers, params=params, verify=True)
        response.raise_for_status()
        return response.json()["results"]
    except requests.exceptions.RequestException as e:
        print(f"Error fetching prefixes: {e}")
        sys.exit(1)

def create_dhcp_range(prefix_str, prefix_id):
    """Create DHCP pool IP range (.100-.150) for a prefix"""
    try:
        network = ipaddress.ip_network(prefix_str, strict=False)

        # Calculate .100 and .150 addresses
        # For /24: x.x.x.100 and x.x.x.150
        # For other sizes, adjust accordingly

        base_ip = str(network.network_address)
        octets = base_ip.split('.')

        # Start: .100
        start_octets = octets[:-1] + ['100']
        start_ip = '.'.join(start_octets)

        # End: .150
        end_octets = octets[:-1] + ['150']
        end_ip = '.'.join(end_octets)

        # Validate IPs are within the network
        start_addr = ipaddress.ip_address(start_ip)
        end_addr = ipaddress.ip_address(end_ip)

        if start_addr not in network or end_addr not in network:
            print(f"  ⚠️  Skipping {prefix_str}: DHCP range (.100-.150) not within network")
            return False

        # Create IP Range via API
        url = f"{NETBOX_URL}/api/ipam/ip-ranges/"
        data = {
            "start_address": start_ip,
            "end_address": end_ip,
            "status": "active",
            "description": "DHCP Pool - Dynamic Assignment Range (.100-.150)",
            "comments": "Auto-created: IPs reserved for DHCP dynamic assignment"
        }

        response = requests.post(url, headers=headers, json=data, verify=True)

        if response.status_code == 201:
            print(f"  ✓ Created DHCP range: {start_ip} - {end_ip}")
            return True
        elif response.status_code == 400:
            error_detail = response.json()
            if "already exists" in str(error_detail).lower() or "overlap" in str(error_detail).lower():
                print(f"  ℹ️  Range already exists: {start_ip} - {end_ip}")
                return True
            else:
                print(f"  ✗ Failed: {error_detail}")
                return False
        else:
            print(f"  ✗ HTTP {response.status_code}: {response.text}")
            return False

    except Exception as e:
        print(f"  ✗ Error processing {prefix_str}: {e}")
        return False

def main():
    print(f"\nConnecting to NetBox: {NETBOX_URL}")
    print("=" * 60)

    # Get all prefixes
    print("\nFetching prefixes...")
    prefixes = get_all_prefixes()
    print(f"Found {len(prefixes)} prefixes\n")

    # Process each prefix
    success_count = 0
    skip_count = 0
    fail_count = 0

    for prefix in prefixes:
        prefix_str = prefix["prefix"]
        prefix_id = prefix["id"]
        vlan_name = prefix.get("vlan", {}).get("name", "No VLAN") if prefix.get("vlan") else "No VLAN"

        print(f"Processing: {prefix_str} ({vlan_name})")

        result = create_dhcp_range(prefix_str, prefix_id)

        if result:
            success_count += 1
        elif result is None:
            skip_count += 1
        else:
            fail_count += 1

    # Summary
    print("\n" + "=" * 60)
    print("Summary:")
    print(f"  ✓ Created/Exists: {success_count}")
    print(f"  ⚠️  Skipped: {skip_count}")
    print(f"  ✗ Failed: {fail_count}")
    print("=" * 60)

    print("\n✓ DHCP pool ranges configured!")
    print("   IPs .100-.150 are now marked as DHCP pool in NetBox")
    print(f"   View at: {NETBOX_URL}/ipam/ip-ranges/")

if __name__ == "__main__":
    main()