hswaw/machines/customs.hackerspace.pl/update_authorized_keys.py

#!/usr/bin/env nix-shell
#!nix-shell -i python3 -p python3Packages.ldap3

from ldap3 import Server, Connection, LEVEL
from ldap3.utils.dn import escape_rdn
import getpass
import logging
from pathlib import Path
import filecmp
import os

import argparse

parser = argparse.ArgumentParser()
parser.add_argument("hostname", help="hostname")
parser.add_argument("ldap_pass_file", type=Path, help="file containing lap password")

header_warning = """
################################### WARNING ####################################
# This file was created automatically from LDAP database and *WILL* be
# overwritten. If you need to add / remove keys make changes to
# {}-admin group / members sshPublicKey attributes in LDAP and rerun
# update_authorized_keys script   
################################################################################
""".lstrip()

def get_keys(connection: Connection, group: str):
    c = connection

    c.search(
        search_base="ou=People,dc=hackerspace,dc=pl",
        search_filter=(
            "(&"
            "(objectClass=hsMember)"
            f"(memberOf=cn={escape_rdn(group)},ou=Group,dc=hackerspace,dc=pl)"
            ")"
        ),
        search_scope=LEVEL,
        attributes=["sshPublicKey"],
    )

    admin_keys = []
    for entry in c.response:
        attributes = entry["attributes"]
        for key in entry["attributes"]["sshPublicKey"]:
            yield key.strip()


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    args = parser.parse_args()

    user = f"cn={escape_rdn(args.hostname)},ou=Boxen,dc=hackerspace,dc=pl"
    password = args.ldap_pass_file.read_text().strip()

    s = Server("ldap.hackerspace.pl", use_ssl=True)
    with Connection(s, user=user, password=password, raise_exceptions=True) as c:
        keys = list(get_keys(c, f"{args.hostname}-admin"))
        if len(keys) < 2:
            raise Exception("Less then two keys found - aborting")

        ssh_dir = Path("/", "root", ".ssh")
        ssh_dir.mkdir(mode=700, exist_ok=True)
        ssh_dir.chmod(0o700)
        new_file = ssh_dir.joinpath("authorized_keys_new")
        old_file = ssh_dir.joinpath("authorized_keys")
        try:
            new_file.unlink()
        except FileNotFoundError:
            pass
        new_file.write_bytes(
            header_warning.format(args.hostname).encode() + b"\n".join(keys)
        )
        if not old_file.exists():
            logging.info('Creating new "authorized_keys" file')
            os.rename(new_file, old_file)
        elif filecmp.cmp(new_file, old_file, shallow=False):
            logging.debug('Nothing changed - "authorized_keys" file is up to date')
            new_file.unlink()
        else:
            logging.info('Keys changed - overwriting "authorized_keys" file')
            os.rename(new_file, old_file)