#!/usr/bin/python3
#
# ssh-based server to allow hosts to update their own TLS-related RRs
# intended to be used like
#   command="/srv/tls/bin/addrecord hepworth.siccegge.de" ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIOIlLx3R+Q5LgBZbJ6USuzam/uAEQITl6vzOn/ylk4fq christoph@mitoraj

from subprocess import call
from functools import reduce
import datetime
import fcntl
import os
import re
import sys
import yaml

class Addrecord:
    def __init__(self, host):
        self._host = host
        self._ports = {'www': [443],
                       'smtp': [25, 587, 465],
                       'imap': [993, 143],
                       'pop': [110, 995],
                       'xmpp': [5222, 5269],
                      }
        self._tlsa_re = re.compile('TLSA 3 1 1 [0-9a-f]{64}')
        self._acme_re = re.compile(r'[a-zA-Z0-9_-]{43}')
        with open('config/inventory.yaml') as inv:
            self._inventory = yaml.load(inv)


    def tlsa(self):
        host = input("Hostname: ")
        service = input("Service: ")
        value = input("Value: ")
        if host not in self._allowed_names(service):
            sys.stderr.write("Not authorized to update entries for service '%s' on '%s'\n"
                              % (service, host))
            return 1

        if re.fullmatch(self._tlsa_re, value) is None:
            sys.stderr.write("Not a valid TLSA record: '%s'\n" % value)
            return 2

        records = []
        for port in self._ports[service]:
            records.append("{0:<35s}\tIN\t{1}.\n".format("_%d._tcp.%s" % (port, host),
                                                    value))
        self._update_records('tlsa', host, records)
        print("OK")
        return 0


    def acme(self):
        host = input("Hostname: ")
        value = input("Value: ")

        if host not in self._allowed_names():
            sys.stderr.write("Not authorized to update entries for host '%s'\n" % host)
            return 1

        if re.fullmatch(self._acme_re, value) is None:
            sys.stderr.write("Not a valid ACME challenge record: '%s'\n" % value)
            return 2

        records = [ "{0:<35s}\tIN\tTXT\t\"{1}\"\n".format("_acme-challenge.%s." % host, value) ]
        self._update_records('acme', host, records)
        print("OK")
        return 0


    def _update_records(self, sort, host, records):
        def find_zone(host):
            assert len(host) != 0
            candidate = os.path.join(sort, '.'.join(host + ['m4']))
            if os.path.exists(candidate):
                return candidate
            else:
                return find_zone(host[1:])


        to_remove = [ i.split()[0] for i in records ]

        zone = find_zone(host.split('.'))
        with open(zone, 'r') as oldzone:
            fcntl.flock(oldzone, fcntl.LOCK_EX)
            lines = oldzone.readlines()
            lines = [ line for line in lines if line == '\n' or line.split()[0] not in to_remove ]

            lines.append('\n')
            lines.append("; Last updated %s by %s\n" % (datetime.datetime.utcnow().isoformat(),
                                                        self._host))
            lines = lines + records

            with open('%s.new' % (zone,), 'w') as newzone:
                newtext = ''.join(lines)
                newtext = re.sub(r'\n[\n]+', '\n\n', newtext)
                newtext = re.sub(r'\n;.*\n\n;', '\n;', newtext)
                newzone.write(newtext)

            os.rename('%s.new' % (zone,),
                      zone)
            fcntl.flock(oldzone, fcntl.LOCK_UN)
        call(["ssh", "root@localhost"])


    def _allowed_names(self, service=None):
        if service is None:
            perservice = self._inventory[self._host].values()
        else:
            perservice = self._inventory[self._host][service]

        names = set()
        for entry in perservice:
            if type(entry) is list:
                names = names.union(entry)
            elif type(entry) is dict:
                names = names.union(reduce(lambda x, y: x + y, entry.values(), []))
            else:
                sys.stderr.write("inventory format wrong\n")
        return names


def main():
    command = os.environ['SSH_ORIGINAL_COMMAND']
    host = sys.argv[1]
    addrecord = Addrecord(host)

    if command == 'acme':
        return addrecord.acme()
    elif command == 'tlsa':
        return addrecord.tlsa()


if __name__ == '__main__':
    sys.exit(main())