#!/usr/bin/env python3

import argparse
import base64
import binascii
import os
import pwd
import re
import subprocess


class UserAdder(object):
    def __init__(self, *args, **kwargs):
        # This doesn't *really* need to be a class, but I may do something with it in the future.
        _tmpargs = locals()
        del(_tmpargs['self'])
        for k, v in _tmpargs.items():
            setattr(self, k, v)
        self.users = {}

    def addUser(self, user, *args, **kwargs):
        # We ideally should do this purely pythonically, but libuser is external and not available everywhere...
        # We *could* do it by hand (add to /etc/passwd, etc.) but that's not guaranteed to be totally compatible.
        # Don't try to add a user if they exist. Doesn't support e.g. LDAP auth.
        try:
            u = pwd.getpwnam(user)
            homedir = u.pw_dir
        except KeyError:
            homedir = '/home/{0}'.format(user)
            subprocess.run(['useradd',
                            '-M',
                            '-c',
                            'Added by add-borguser.py',
                            '-d',
                            homedir,
                            user])
        sshdir = os.path.join(homedir, '.ssh')
        authkeys = os.path.join(sshdir, 'authorized_keys')
        userent = pwd.getpwnam(user)
        uid, gid = userent.pw_uid, userent.pw_gid
        os.makedirs(homedir, mode = 0o700, exist_ok = True)
        os.makedirs(sshdir, mode = 0o700, exist_ok = True)
        os.chown(homedir, uid, gid)
        os.chown(sshdir, uid, gid)
        if not os.path.isfile(authkeys):
            with open(authkeys, 'w') as f:
                f.write('')
        os.chmod(authkeys, 0o0400)
        os.chown(authkeys, uid, gid)
        self.users[user] = authkeys
        return()

    def addKey(self, ssh_key, *args, **kwargs):
        key_template = ('command='
                            #'"cd {homedir};'
                            #'borg serve --restrict-to-path {homedir}",'
                            '"/usr/local/bin/borg-restricted.py ${{SSH_ORIGINAL_COMMAND}}",'
                        'no-port-forwarding,'
                        'no-X11-forwarding,'
                        'no-pty,'
                        'no-agent-forwarding,'
                        'no-user-rc '
                        '{keystr}\n')
        for u, kp in self.users.items():
            userent = pwd.getpwnam(u)
            homedir = userent.pw_dir
            sshdir = os.path.join(homedir, '.ssh')
            key_insert = key_template.format(user = u,
                                             homedir = homedir,
                                             keystr = ssh_key)
            with open(kp, 'a') as f:
                f.write(key_insert)
            # When CentOS/RHEL move to python3 native, and port policycoreutils, do this natively.
            # But for now...
            subprocess.run(['chcon',
                            '-R unconfined_u:object_r:user_home_t:s0',
                            sshdir])
            subprocess.run(['semanage',
                            'fcontext',
                            '-a',
                                '-t',
                                    'ssh_home_t',
                            sshdir])
        return()

    def clean(self):
        self.users = {}
        return()

def parseArgs():
    def _valid_posix_user(username):
        # http://pubs.opengroup.org/onlinepubs/9699919799/basedefs/V1_chap03.html#tag_03_437
        # http://pubs.opengroup.org/onlinepubs/9699919799/basedefs/V1_chap03.html#tag_03_282
        # https://unix.stackexchange.com/a/435120/284004
        if not re.search('^[a-z_]([a-z0-9_-]{0,31}|[a-z0-9_-]{0,30}\$)$', username):
            raise argparse.ArgumentTypeError('user must be a POSIX-compliant username')
        return(username)
    def _valid_ssh_key(keystr):
        # This validation is *super* cursory. We could probably do some better parsing at some point.
        key_components = keystr.split()
        keytype = re.sub('^ssh-(.*)', '\g<1>', key_components[0])
        # We don't support anything but ED25519 or RSA, given that they used the hardening guide.
        if keytype not in ('ed25519', 'rsa'):
            raise argparse.ArgumentTypeError('Not a valid SSH pubkey type (must be RSA or ED25519)')
        try:
            base64.b64decode(key_components[1].encode('utf-8'))
        except binascii.Error:
            raise argparse.ArgumentTypeError('Not a valid SSH pubkey')
        return(keystr)
    args = argparse.ArgumentParser(description = ('Add local users to a borg server'))
    args.add_argument('user',
                      type = _valid_posix_user,
                      help = 'The username/machine name to add')
    args.add_argument('ssh_key',
                      type = _valid_ssh_key,
                      help = ('The full SSH pubkey (remember to enclose in quotes)'))
    return(args)

def main():
    if not os.geteuid() == 0:
        raise PermissionError('This script must be run as root or with root-like privileges')
    args = vars(parseArgs().parse_args())
    um = UserAdder(**args)
    um.addUser(**args)
    um.addKey(**args)
    um.clean()
    return()

if __name__ == '__main__':
    main()