#!/usr/bin/env python3
# File: generate-pihole-dns-dhcp.py
# 
# Description: This file generates the DNS and DHCP files for pihole.
# 
# Package: AniNIX/Ubiqtorate
# Copyright: WTFPL
# 
# Author: DarkFeather <darkfeather@aninix.net>

import os
import sys
import yaml

dnsfilepath="roles/Nazara/files/dns"
dhcpfilepath="roles/Nazara/files/dhcp"

def WriteDHCPEntry(content,hosttype,hostclass):
    ### Create the DHCP entry
    # param content: the yaml content to parse
    # param hosttype: managed or unmanaged
    # param hostclass: the type of host as classified in the yaml
    global dhcpfile

    with open(dhcpfilepath,'a') as dhcpfile: 
        for host in content['all']['children'][hosttype]['children'][hostclass]['hosts']:
            try:
                dhcpfile.write('dhcp-host=' + content['all']['children'][hosttype]['children'][hostclass]['hosts'][host]['vars']['mac'] + ',' + content['all']['children'][hosttype]['children'][hostclass]['hosts'][host]['vars']['ip'] + '\n')
            except:
                print(host + ' is not complete for DHCP.')

def WriteDNSEntry(content,hosttype,hostclass):
    ### Create the DNS entry
    # param content: the yaml content to parse
    # param hosttype: managed or unmanaged
    # param hostclass: the type of host as classified in the yaml
    global dnsfile

    with open(dnsfilepath,'a') as dnsfile: 
        for host in content['all']['children'][hosttype]['children'][hostclass]['hosts']:
            try:
                dnsfile.write(content['all']['children'][hosttype]['children'][hostclass]['hosts'][host]['vars']['ip'] + ' ' + host + '.' + content['all']['vars']['replica_domain'] + ' ' + host + '\n')
            except:
                print(host + ' is not complete for DNS.')

def GenerateFiles(file):
    ### Open the file and parse it
    # param file: the file to work on
    global dnsfile

    # Parse the yaml
    with open(file, 'r') as stream:
        content = yaml.safe_load(stream)

        # Clear the DNS file
        with open(dhcpfilepath,'w') as dhcpfile:
            dhcpfile.write('dhcp-range='+content['all']['vars']['dhcprange']+'\n')
            dhcpfile.write('dhcp-option=option:router,'+content['all']['vars']['router']+'\n')
            dhcpfile.write('dhcp-option=option:dns-server,'+content['all']['vars']['dns']+'\n\n')
            dhcpfile.write('dhcp-range='+content['all']['vars']['staticrange']+'\n')
        with open(dnsfilepath,'w') as dnsfile:
            dnsfile.write('')

        # Add DNS entries for each host
        hosttype = 'managed'
        for hostclass in ['physical','virtual','geth-hubs']: 
            WriteDNSEntry(content,hosttype,hostclass)
            WriteDHCPEntry(content,hosttype,hostclass)
        hosttype = 'unmanaged'
        for hostclass in ['ovas','hardware','iot']:
            WriteDNSEntry(content,hosttype,hostclass)
            WriteDHCPEntry(content,hosttype,hostclass)

if __name__ == '__main__':
    if len(sys.argv) != 2:
        print("You need to supply an inventory file.")
        sys.exit(1)
    GenerateFiles(sys.argv[1])
    sys.exit(0)