#!/usr/bin/env python3

import time
from contextlib import contextmanager
import re
import os
import os.path
import random
import string
import subprocess
import sqlite3
import hmac
from ipaddress import IPv4Address, IPv6Address

from flask import Flask, request, abort
import uwsgidecorators

app = Flask(__name__)
app.config.update(dict(
    RECORD_EXPIRY_S = 86400,
    NSD_CONTROL = 'nsd-control'
    ))
app.config.from_pyfile('config.py')


ZONEFILE_TEMPLATE = '''\
; #################################################### ;
; THIS FILE IS AUTOMATICALLY GENERATED! DO NOT MODIFY! ;
; #################################################### ;

$ORIGIN {zone}.
$TTL 1800
@       IN      SOA     {ns}.      {mail}. (
                        {serial}                ; serial number
                        60                      ; refresh
                        60                      ; retry
                        {expire}                ; expire
                        60                      ; ttl
                        )
; Name servers
                    IN      NS      {ns}.

; Additional A records from template
; @                   IN      A       192.0.2.3
; www                 IN      A       192.0.2.3

; Dynamically generated records
{dynamic_records}
'''

db = sqlite3.connect(app.config['SQLITE_DB'], check_same_thread=False)
with db as conn:
    conn.execute('''CREATE TABLE IF NOT EXISTS zone_versions (date TEXT)''')
    conn.execute('''CREATE TABLE IF NOT EXISTS records
            (name TEXT PRIMARY KEY, ipv4 TEXT, ipv6 TEXT, last_update INTEGER)''')

def purge_expired_records():
    with db as conn:
        conn.execute('DELETE FROM records WHERE last_update < ?',
                (int(time.time()) - app.config['RECORD_EXPIRY_S'],))

def update_record(record, ipv4=None, ipv6=None):
    with db as conn:
        old_v4, old_v6 = conn.execute('SELECT ipv4, ipv6 FROM records WHERE name=?', (record,)).fetchone() or (None, None)
        conn.execute('INSERT OR REPLACE INTO records VALUES (?, ?, ?, ?)', (record, ipv4, ipv6, int(time.time())))
        return ipv4 != old_v4 or ipv6 != old_v6

@contextmanager
def inplace_rewrite(filename, cleanup=True):
    print('Writing', filename)
    filename = os.path.abspath(filename)
    if cleanup:
        basename = os.path.basename(filename)
        for entry in os.scandir(os.path.dirname(filename)):
            if entry.name.startswith(basename) and re.match(r'\.tmp-[a-zA-Z0-9]{8}', entry.name[len(basename):]):
                os.remove(entry.path)

    tmp_fn = filename + f'.tmp-' + ''.join(random.choices(string.ascii_letters + string.digits, k=8))
    with open(tmp_fn, 'w') as tmp_f:
        yield tmp_f
        tmp_f.flush()
        os.fsync(tmp_f.fileno())
    os.rename(tmp_fn, filename)

def write_zonefile():
    # Find the next free zonefile version number
    with db as conn:
        conn.execute('INSERT INTO zone_versions VALUES (DATE())')
        date, version_num, = conn.execute('SELECT zone_versions.date, COUNT(*) FROM zone_versions WHERE zone_versions.date = DATE()').fetchone()
    zone_version = f'{date.replace("-", "")}{version_num:02d}'

    # Generate dynamic record block
    with db as conn:
        records = db.execute('SELECT name, "A", ipv4 FROM records UNION SELECT name, "AAAA", ipv6 FROM records')
        dynamic_records = '\n'.join(f'{name:<20} IN {rtype:<4} {value}' for name, rtype, value in records if value is not None)

    # Template zone file content
    content = ZONEFILE_TEMPLATE.format(
            zone = app.config['ZONE'],
            ns = app.config['NAMESERVER'],
            mail = app.config['NAMESERVER_MAIL'].replace('@', '.'),
            serial = zone_version,
            dynamic_records = dynamic_records,
            expire = app.config['RECORD_EXPIRY_S']
            )

    with inplace_rewrite(app.config['ZONEFILE'], cleanup=True) as f:
        f.write(content)

def kick_nsd():
    prog = app.config['NSD_CONTROL']
    if isinstance(prog, str):
        prog = [prog]
    subprocess.run([*prog, 'reload', app.config['ZONE']], check=True)

@app.before_first_request
@uwsgidecorators.timer(300)
def update_zonefile():
    purge_expired_records()
    write_zonefile()
    kick_nsd()

@app.route('/update', methods=['POST'])
def route_update():
    if request.authorization is None:
        abort(403)

    record = request.authorization['username']
    record_config = app.config['DYNAMIC_RECORDS'].get(record)
    if record_config is None:
        abort(403)

    *supported_formats, password = record_config
    if not hmac.compare_digest(request.authorization['password'], password):
        abort(403)

    ipv4 = request.args.get('ipv4', '127.0.0.1')
    ipv6 = request.args.get('ipv6', '::1')
    ipv4 = str(IPv4Address(ipv4)) if 'v4' in supported_formats else None
    ipv6 = str(IPv6Address(ipv6)) if 'v6' in supported_formats else None
    if update_record(record, ipv4=ipv4, ipv6=ipv6):
        update_zonefile()

    return 'success'


if __name__ == '__main__':
    app.run()