#!/usr/bin/env python3

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


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

    @staticmethod
    def create(struct: pwd.struct_passwd, ssh_keys: list[str]) -> 'UserData':
        return UserData(struct.pw_uid, struct.pw_name, struct.pw_gecos, struct.pw_shell, struct.pw_dir, [])

def get_userdata(user_file: dict[str, Any], verbose: bool) -> tuple[datetime.datetime, list[UserData]]:
    users: list[dict[str, Any]] = user_file["users"]
    created_at = datetime.datetime.fromisoformat(user_file["timestamp"])
    now = datetime.datetime.now(datetime.timezone.utc)

    if (now - created_at) > datetime.timedelta(days=1):
        print(f"WARNING: The user file is older than 1 day ({created_at=})", file=sys.stderr)
    
    if verbose:
        print(f"Loaded {len(users)} users, datafile age = {now - created_at}")

    result = []
    for user in users:
        uid = int(user["uid"])
        keys = user["ssh_keys"]
        try:
            struct = pwd.getpwuid(uid)
        except KeyError:
            print(f"User {user['name']}:{uid} does not exist", file=sys.stderr)
            continue
        result.append(UserData.create(struct, keys))
    return created_at, result

def copy_file(local_path, server, remote_path, verbose):
    cmd = [ "rsync", local_path, f"{server}:{remote_path}" ]
    if verbose:
        cmd.append("--verbose")
        print(f"Executing", *cmd)
    subprocess.check_call(cmd)

def execute_command(server, command, verbose):
    cmd = [ "ssh", server, command ]
    if verbose:
        print(f"Executing", *cmd)
    subprocess.check_call(cmd)

def main_args(keys_file: str, server: str, command: Optional[str], transfer_file: Optional[str], server_transfer_file: str, verbose: bool):
    if os.getuid() == 0:
        print("WARNING: This script should be not executed as root, only collect_keys.py requires that", file=sys.stderr)

    with open(keys_file, "r") as f:
        user_list = json.load(f)
    keys_ts, users = get_userdata(user_list, verbose)
    transfer_json = {
        "keys_timestamp": keys_ts.isoformat(),
        "timestamp": datetime.datetime.now(datetime.timezone.utc).isoformat(),
        "users": [ asdict(user) for user in users ]
    }

    if transfer_file is None:
        with tempfile.NamedTemporaryFile("w") as f:
            json.dump(transfer_json, f)
            f.flush()
            copy_file(f.name, server, server_transfer_file, verbose)
    else:
        with open(transfer_file, "w") as f:
            json.dump(transfer_json, f)
        copy_file(transfer_file, server, server_transfer_file, verbose)

    if command is not None:
        execute_command(server, command, verbose)

def main():
    parser = argparse.ArgumentParser(description='Copies list of users to a remote server')
    parser.add_argument("--keys-file", help="The file generated using collect_keys.py script", required=True)
    parser.add_argument("--server", help="SSH server where to copy the file and execute the command", required=True)
    parser.add_argument("--command", default=None, help="Command to execute on the remote server. Will be executed using default shell of ssh")
    parser.add_argument("--transfer-file", default=None, help="Write the transfer JSON file to this path. Otherwise, temporary file is used.")
    parser.add_argument("--server-transfer-file", help="The file generated using collect_keys.py script", default="/var/ksp-user-sync/new-users.json")
    parser.add_argument("--verbose", action="store_true", help="Print debug messages")
    args = parser.parse_args()
    main_args(args.keys_file, args.server, args.command, args.transfer_file, args.server_transfer_file, args.verbose)

main()