localhostgen/localhostgen.py

#! /usr/bin/env python3

import argparse
import json
import os
from time import time
import psycopg2
import sys
import re


class LocalHostGen:
    db_config = None
    db_instance = None
    dhcpd_replace_string = "# %%HOST%%"

    def __init__(self, db_config):
        self.db_config = db_config

    def connect(self):
        self.db_instance = psycopg2.connect("dbname='%s' user='%s' host='%s' password='%s' port='%s'" %
                                            (self.db_config['database'], self.db_config['username'],
                                             self.db_config['host'], self.db_config['password'],
                                             self.db_config['port']))

    def disconnect(self):
        if self.db_instance is not None:
            self.db_instance.close()
            self.db_instance = None

    def get_domain_id(self, domain):
        cur = self.db_instance.cursor()
        cur.execute("SELECT id FROM domains WHERE name='%s'" % domain)
        id = cur.fetchone()
        cur.close()
        if id is not None:
            return id[0]
        return None

    def get_record_id(self, record):
        cur = self.db_instance.cursor()
        cur.execute("SELECT id FROM records WHERE name='%s'" % record)
        id = cur.fetchone()
        cur.close()
        if id is not None:
            return id[0]
        return None

    def get_record_content(self, record):
        cur = self.db_instance.cursor()
        cur.execute("SELECT content FROM records WHERE name='%s'" % record)
        content = cur.fetchone()
        cur.close()
        if content is not None:
            return content[0]
        return None

    def get_reversed_ip(self, ip):
        return '.'.join(reversed(ip.split(".")))

    def get_reverse_domain_record(self, reverse_domain, ip):
        reversed_ip = self.get_reversed_ip(ip)
        reverse_domain_part = reversed_ip
        while not reverse_domain.startswith(reverse_domain_part):
            reverse_domain_part = reverse_domain_part[reverse_domain_part.index('.') + 1:]
        return reversed_ip[0:len(reversed_ip) - len(reverse_domain_part) - 1]

    def insert_host(self, domain, record, content, ttl, type):
        record_full = "%s.%s" % (record, domain)

        record_id = self.get_record_id(record_full)
        if record_id is None:
            domain_id = self.get_domain_id(domain)
            cur = self.db_instance.cursor()
            cur.execute("INSERT INTO records (domain_id, name, type, content, ttl, prio, change_date,"
                        "disabled, ordername, auth) VALUES ('%s', '%s', '%s', '%s', '%s', '0', '%s', 'f', NULL, 't')"
                        "RETURNING id" % (domain_id, record_full, type, content, ttl, int(time())))
            record_id = cur.fetchone()[0]
            cur.close()
            self.db_instance.commit()
        return record_id

    def delete_host(self, record):
        cur = self.db_instance.cursor()
        cur.execute("DELETE FROM records WHERE name='%s' RETURNING id" % record)
        record_id = cur.fetchone()
        cur.close()
        self.db_instance.commit()
        if record_id is not None:
            return record_id[0]
        return None

    def remove_host_dhcp(self, host, dhcpd_file_content):
        return re.sub(" *host +%s +\{[^}]+} *\n" % host, "", dhcpd_file_content)

    def create_host(self, domain, reverse_domain, host, ip, mac, ttl, dhcpd_file_path, wildcard):
        record_full = "%s.%s" % (host, domain)
        host_wildcard = "*.%s" % host
        reverse_domain_record = self.get_reverse_domain_record(reverse_domain, ip)

        self.insert_host(domain, host, ip, ttl, 'A')
        if wildcard:
            self.insert_host(domain, host_wildcard, record_full, ttl, 'CNAME')
        self.insert_host(reverse_domain, reverse_domain_record, record_full, ttl, 'PTR')

        with open(dhcpd_file_path, "r") as f:
            dhcpd_file_content = f.read()

        dhcpd_host_config = "host %s {\n" \
                            "        hardware ethernet %s;\n" \
                            "        fixed-address %s;\n" \
                            "    }" % (host, mac, record_full)

        dhcpd_file_content = self.remove_host_dhcp(host, dhcpd_file_content)
        dhcpd_file_content = dhcpd_file_content.replace(self.dhcpd_replace_string, "%s\n    %s" %
                                                        (dhcpd_host_config, self.dhcpd_replace_string))

        with open(dhcpd_file_path, "w") as f:
            f.write(dhcpd_file_content)

    def remove_host(self, domain, reverse_domain, host, dhcpd_file_path):
        record_full = "%s.%s" % (host, domain)
        record_full_wildcard = "*.%s" % record_full
        ip = self.get_record_content(record_full)
        if ip is not None:
            reverse_domain_record = self.get_reverse_domain_record(reverse_domain, ip)
            reverse_domain_record_full = "%s.%s" %(reverse_domain_record, reverse_domain)
            self.delete_host(record_full_wildcard)
            self.delete_host(record_full)
            self.delete_host(reverse_domain_record_full)

        with open(dhcpd_file_path, "r") as f:
            dhcpd_file_content = f.read()

        dhcpd_file_content = self.remove_host_dhcp(host, dhcpd_file_content)

        with open(dhcpd_file_path, "w") as f:
            f.write(dhcpd_file_content)


def eprint(*args, **kwargs):
    print(*args, file=sys.stderr, **kwargs)


def main():
    parser = argparse.ArgumentParser(description='Manage local hosts for PowerDNS and isc-dhcp-server')
    parser.add_argument('--host', help='The host to manage', required=True)
    parser.add_argument('--config', dest='config', default='/etc/localhostgen/localhostgen.json',
                        help='Configuration file path')

    parser.add_argument('--create', help='Create a host', action='store_true')
    parser.add_argument('--remove', help='Remove the host', action='store_true')

    parser.add_argument('--mac', help='The MAC address of the new host')
    parser.add_argument('--ip', help='The IP address of the new host')

    args = parser.parse_args()

    with open(args.config, "r") as f:
        config = json.load(f)

    host = args.host
    ip = args.ip
    mac = args.mac

    local_host_gen = LocalHostGen(config['database'])
    local_host_gen.connect()

    if args.create:
        if ip is None or mac is None:
            eprint("IP and MAC are required to create a host")
            exit(1)
        local_host_gen.create_host(config['dns']['domain'], config['dns']['reverseDomain'], host, ip, mac,
                                   config['dns']['recordTTL'], config['dhcpd']['filePath'], config['dns']['wildcard'])
    elif args.remove:
        local_host_gen.remove_host(config['dns']['domain'], config['dns']['reverseDomain'], host,
                                   config['dhcpd']['filePath'])

    local_host_gen.disconnect()

    os.system(config['dhcpd']['reloadCommand'])


main()