ksp-vm/user-sync/accept_users.py

#!/usr/bin/env python3

import os, sys, argparse, json, datetime, tempfile, pwd, grp, subprocess
from dataclasses import dataclass
from typing import Optional, Any

def report_persistent_problem(*msg):
    print(*msg, file=sys.stderr)
    with open("./errors", "a") as f:
        print(*msg, file=f)

@dataclass(frozen=True)
class UserData:
    uid: int
    name: str
    homedir: str
    full_name: str
    shell: str
    ssh_keys: list[str]

def read_file(path: str) -> Optional[dict[str, Any]]:
    if not os.path.exists(path):
        return None
    with open(path) as f:
        return json.load(f)
    
def validate_shell(user: UserData) -> str:
    if not os.path.exists(user.shell):
        report_persistent_problem(f"Shell {user.shell} is not installed on this system (user {user.name})!")
        return "/bin/bash"
    return user.shell

def execute_maybe(cmd: list[str], dry_run: bool, verbose: bool):
    if verbose or dry_run:
        print(f"Executing", *cmd)
    if not dry_run:
        subprocess.check_call(cmd)

def create_user(user: UserData, marker_group: str, extra_groups: list[str], dry_run: bool, verbose: bool):
    try:
        pwd.getpwnam(user.name)
        report_persistent_problem(f"Cannot create {user.name}:{user.uid} - username already taken!")
        return
    except KeyError:
        pass
    try:
        pwd.getpwuid(user.uid)
        report_persistent_problem(f"Cannot create {user.name}:{user.uid} - uid already taken!")
        return
    except KeyError:
        pass


    cmd = [ "useradd",
           "--uid", str(user.uid),
           "--user-group",
           "--create-home",
           "--comment", user.full_name,
           "--groups", ",".join([marker_group, *extra_groups]),
           "--shell", validate_shell(user),
            # TODO: subuid and subgid are handled?
    ]
    execute_maybe(cmd, dry_run, verbose)

def modify_user(old: UserData, new: UserData, marker_group: str, dry_run: bool, verbose: bool):
    if old.uid != new.uid or old.name != new.name:
        report_persistent_problem(f"Cannot modify {old.name}:{old.uid} to {new.name}:{new.uid} - uid or name changed!")
        return
    
    if new.name not in grp.getgrnam(marker_group).gr_mem:
        if verbose:
            print(f"Skipping {new.name}:{new.uid} - not in {marker_group} group")
        return

    cmd = [ "usermod" ]
    if old.full_name != new.full_name:
        cmd.extend([ "--comment", new.full_name ])
    if old.shell != new.shell:
        cmd.extend([ "--shell", validate_shell(new) ])

    if cmd == [ "usermod" ]:
        return
    execute_maybe(cmd, dry_run, verbose)

def add_group(user_name: str, group: str, dry_run: bool, verbose: bool):
    if user_name not in grp.getgrnam(group).gr_mem:
        cmd = [ "usermod", "--append", "--groups", group, user_name ]
        execute_maybe(cmd, dry_run, verbose)


def save_keys(users: list[UserData], ssh_keys_location: str, dry_run: bool, verbose: bool):
    remove_files = set(os.listdir(ssh_keys_location)).difference(u.name for u in users if len(u.ssh_keys) > 0)
    if remove_files:
        if verbose or dry_run:
            print(f"Removing SSH keys: {remove_files}")
        if not dry_run:
            for f in remove_files:
                os.remove(os.path.join(ssh_keys_location, f))

    for user in users:
        # doesn't matter if it changed on source, always overwrite if different
        # users should not overwrite these files
        if len(user.ssh_keys) == 0:
            continue
        file = os.path.join(ssh_keys_location, user.name)
        new_content = "".join(line.rstrip() + "\n" for line in  user.ssh_keys)
        if os.path.exists(file):
            with open(file) as f:
                old_content = f.read()
            if old_content == new_content:
                continue

        if verbose or dry_run:
            print(f"Saving {len(user.ssh_keys)} SSH keys for {user.name}:{user.uid}")
        if not dry_run:
            with open(file, "w") as f:
                f.write(new_content)

def main_args(new_file_path: str, old_file_path: str, ssh_keys_location, marker_group: str, extra_groups: list[str], removed_group: Optional[str], dry_run: bool, verbose: bool):
    if os.getuid() != 0:
        print("WARNING: This script should be executed as root", file=sys.stderr)

    new_file = read_file(new_file_path)
    if new_file is None:
        raise Exception(f"File {new_file_path} does not exist")
    old_file = read_file(old_file_path)

    if old_file and datetime.datetime.fromisoformat(old_file["timestamp"]) > datetime.datetime.fromisoformat(new_file["timestamp"]):
        raise Exception(f"Old file {old_file_path} is newer than {new_file_path}")

    new_users = [UserData(**u) for u in new_file["users"] ]
    old_users = {u["name"]: UserData(**u) for u in old_file["users"] } if old_file else {}

    created_users = [u for u in new_users if u.name not in old_users]
    modified_users = [u for u in new_users if u.name in old_users and u != old_users[u.name]]
    removed_users = set(u.name for u in old_users.values() if u.name).difference(u.name for u in new_users)

    allowed_users = set(grp.getgrnam(marker_group).gr_mem).union(u.name for u in created_users)

    for u in created_users:
        create_user(u, marker_group, extra_groups, dry_run, verbose)
    for u in modified_users:
        if u.name in allowed_users:
            modify_user(old_users[u.name], u, marker_group, dry_run, verbose)
    for u in removed_users:
        if u in allowed_users and removed_group is not None:
            add_group(u, removed_group, dry_run, verbose)

    save_keys(new_users, ssh_keys_location, dry_run, verbose)



def main():
    parser = argparse.ArgumentParser(description='Creates or modifies users on this machine based on data created by copy_users.py')
    parser.add_argument("--new-file", help="The file generated using collect-users.py script", required=True)
    parser.add_argument("--old-file", help="Files used to compare differences, when script finishes this file is overwritten with new_file", required=True)
    parser.add_argument("--ssh-keys-location", help="Directory where all ssh keys are stored", required=True)
    parser.add_argument("--marker-group", help="Group name to add to users created by this script. Only users with this group will be updated", required=True)
    parser.add_argument("--extra-groups", nargs="+", help="Additional groups")
    parser.add_argument("--removed-group", help="Add this group to users which disappeared from source")
    parser.add_argument("--dry_run", action="store_true", help="Only print which commands would be executed")
    parser.add_argument("--verbose", action="store_true", help="Print debug messages")
    args = parser.parse_args()
    main_args(args.new_file, args.old_file, args.ssh_keys_location, args.marker_group, args.extra_groups, args.removed_group, args.dry_run, args.verbose)

main()