localhostgen/localhostgen.py

#! /usr/bin/env python3

import argparse
import json
from time import time

import psycopg2
import sys


class LocalHostGen:
    db_config = None
    db_instance = None

    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()[0]
        cur.close()
        self.db_instance.commit()
        return record_id

    def create_host(self, domain, reverse_domain, host, ip, mac, ttl):
        if self.db_instance is None:
            self.connect()
        reverse_domain_record = self.get_reverse_domain_record(reverse_domain, ip)
        record_full = "%s.%s" % (host, domain)

        self.insert_host(domain, host, ip, ttl, 'A')
        self.insert_host(reverse_domain, reverse_domain_record, record_full, ttl, 'PTR')

    def remove_host(self, domain, reverse_domain, host):
        if self.db_instance is None:
            self.connect()
        record_full = "%s.%s" % (host, domain)
        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)
            self.delete_host(reverse_domain_record_full)


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
    ttl = config['recordTTL']

    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['domain'], config['reverseDomain'], host, ip, mac, ttl)
    elif args.remove:
        local_host_gen.remove_host(config['domain'], config['reverseDomain'], host)

    local_host_gen.disconnect()


main()