#!/usr/bin/python
# -*- coding: utf-8 -*-

# (c) 2012, Stephen Fromm <sfromm@gmail.com>
#
# This file is part of Ansible
#
# Ansible is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Ansible is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Ansible.  If not, see <http://www.gnu.org/licenses/>.

DOCUMENTATION = '''
---
module: user
author: Stephen Fromm
version_added: 0.2
short_description: Manage user accounts
requirements: [ useradd, userdel, usermod ]
description:
    - Manage user accounts and user attributes.
options:
    name:
        required: true
        aliases: [ "user" ]
        description:
            - Name of the user to create, remove or modify.
    comment:
        required: false
        description:
            - Optionally sets the description (aka I(GECOS)) of user account.
    uid:
        required: false
        description:
            - Optionally sets the I(UID) of the user.
    group:
        required: false
        description:
            - Optionally sets the user's primary group (takes a group name).
    groups:
        required: false
        description:
            - Puts the user in this comma-delimited list of groups.
    append:
        required: false
        description:
            - If I(yes), will only add groups, not set them to just the list
              in I(groups).
    shell:
        required: false
        description:
            - Optionally set the user's shell.
    home:
        required: false
        description:
            - Optionally set the user's home directory.
    password:
        required: false
        description:
            - Optionally set the user's password to this crypted value.  See
              the user example in the github examples directory for what this looks
              like in a playbook.
    state:
        required: false
        default: "present"
        choices: [ present, absent ]
        description:
            - Whether the account should exist.  When I(absent), removes
              the user account.
    createhome:
        required: false
        default: "yes"
        choices: [ yes, no ]
        description:
            - Unless set to I(no), a home directory will be made for the user
              when the account is created.
    system:
        required: false
        default: "no"
        choices: [ yes, no ]
        description:
            - When creating an account, setting this to I(yes) makes the user a
              system account.  This setting cannot be changed on existing users.
    force:
        required: false
        default: "no"
        choices: [ yes, no ]
        description:
            - When used with I(state=absent), behavior is as with
              I(userdel --force).
    remove:
        required: false
        default: "no"
        choices: [ yes, no ]
        description:
            - When used with I(state=absent), behavior is as with
              I(userdel --remove).
    ssh_key:
        required: false
        choices: [ generate ]
        version_added: "0.9"
        description:
            - Whether to generate a SSH key for the user in question.
              This will B(not) overwrite an existing SSH key.
    ssh_key_bits:
        required: false
        default: 2048
        version_added: "0.9"
        description:
            - Optionally specify number of bits in SSH key to create.
    ssh_key_type:
        required: false
        default: rsa
        version_added: "0.9"
        description:
            - Optionally specify the type of SSH key to generate. 
              Available SSH key types will depend on implementation
              present on target host.
    ssh_key_file:
        required: false
        default: $HOME/.ssh/id_rsa
        version_added: "0.9"
        description:
            - Optionally specify the SSH key filename.
    ssh_key_comment:
        required: false
        default: ansible-generated
        version_added: "0.9"
        description:
            - Optionally define the comment for the SSH key.
    ssh_key_passphrase:
        required: false
        version_added: "0.9"
        description:
            - Set a passphrase for the SSH key.  If no
              passphrase is provided, the SSH key will default to
              having no passphrase.
examples:
    - code: 'user: name=johnd comment="John Doe" uid=1040'
      description: "Add the user 'johnd' with a specific uid and a primary group of 'admin'"
    - code: "user: name=johnd state=absent remove=yes"
      description: "Remove the user 'johnd'"
    - code: 'user: name=jsmith ssh_key=generate ssh_key_bits=2048'
      description: "Create a 2048-bit SSH key for user jsmith"
'''

import os
import pwd
import grp
try:
    import spwd
    HAVE_SPWD=True
except:
    HAVE_SPWD=False

SHADOWFILE = '/etc/shadow'
if os.path.exists('/etc/master.passwd'):
    SHADOWFILE = '/etc/master.passwd' # FreeBSD passwd
# Note: while the above has the correct location for where
# encrypted passwords are stored on FreeBSD, the code below doesn't
# invoke adduser in lieu of useradd, nor pw in lieu of usermod.
# That is, this won't work on FreeBSD.

def user_del(module, user, **kwargs):
    cmd = [module.get_bin_path('userdel', True)]
    for key in kwargs:
        if key == 'force' and module.boolean(kwargs[key]):
            cmd.append('-f')
        elif key == 'remove' and module.boolean(kwargs[key]):
            cmd.append('-r')
    cmd.append(user)
    p = subprocess.Popen(cmd, shell=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    (out, err) = p.communicate()
    rc = p.returncode
    return (rc, out, err)

def user_add(module, user, **kwargs):
    cmd = [module.get_bin_path('useradd', True)]
    for key in kwargs:
        if key == 'uid' and kwargs[key] is not None:
            cmd.append('-u')
            cmd.append(kwargs[key])
        elif key == 'group' and kwargs[key] is not None:
            if not group_exists(kwargs[key]):
                module.fail_json(msg="Group %s does not exist" % (kwargs[key]))
            cmd.append('-g')
            cmd.append(kwargs[key])
        elif key == 'groups' and kwargs[key] is not None:
            for g in kwargs[key].split(','):
                if not group_exists(g):
                    module.fail_json(msg="Group %s does not exist" % (g))
            cmd.append('-G')
            cmd.append(kwargs[key])
        elif key == 'comment' and kwargs[key] is not None:
            cmd.append('-c')
            cmd.append(kwargs[key])
        elif key == 'home' and kwargs[key] is not None:
            cmd.append('-d')
            cmd.append(kwargs[key])
        elif key == 'shell' and kwargs[key] is not None:
            cmd.append('-s')
            cmd.append(kwargs[key])
        elif key == 'password' and kwargs[key] is not None:
            cmd.append('-p')
            cmd.append(kwargs[key])
        elif key == 'createhome':
            if kwargs[key] is not None:
                value = module.boolean(kwargs[key])
                if value:
                    cmd.append('-m')
                else:
                    cmd.append('-M')
        elif key == 'system' and module.boolean(kwargs[key]):
            cmd.append('-r')
    cmd.append(user)
    p = subprocess.Popen(cmd, shell=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    (out, err) = p.communicate()
    rc = p.returncode
    return (rc, out, err)

"""
Without spwd, we would have to resort to reading /etc/shadow
to get the encrypted string.  For now, punt on idempotent password changes.
"""
def user_mod(module, user, **kwargs):
    cmd = [module.get_bin_path('usermod', True)]
    info = user_info(user)
    for key in kwargs:
        if key == 'uid':
            if kwargs[key] is not None and info[2] != int(kwargs[key]):
                cmd.append('-u')
                cmd.append(kwargs[key])
        elif key == 'group' and kwargs[key] is not None:
            if not group_exists(kwargs[key]):
                module.fail_json(msg="Group %s does not exist" % (kwargs[key]))
            ginfo = group_info(kwargs[key])
            if info[3] != ginfo[2]:
                cmd.append('-g')
                cmd.append(kwargs[key])
        elif key == 'groups' and kwargs[key] is not None:
            current_groups = user_group_membership(user)
            groups = kwargs[key].split(',')
            for g in groups:
                if not group_exists(g):
                    module.fail_json(msg="Group %s does not exist" % (g))
            group_diff = set(sorted(current_groups)).symmetric_difference(set(sorted(groups)))
            groups_need_mod = False

            if group_diff:
                if kwargs['append'] is not None and module.boolean(kwargs['append']):
                    for g in groups:
                        if g in group_diff:
                            cmd.append('-a')
                            groups_need_mod = True
                else:
                    groups_need_mod = True

            if groups_need_mod:
                cmd.append('-G')
                cmd.append(','.join(groups))

        elif key == 'comment':
            if kwargs[key] is not None and info[4] != kwargs[key]:
                cmd.append('-c')
                cmd.append(kwargs[key])
        elif key == 'home':
            if kwargs[key] is not None and info[5] != kwargs[key]:
                cmd.append('-d')
                cmd.append(kwargs[key])
        elif key == 'shell':
            if kwargs[key] is not None and info[6] != kwargs[key]:
                cmd.append('-s')
                cmd.append(kwargs[key])
        elif key == 'password':
            if kwargs[key] is not None and info[1] != kwargs[key]:
                cmd.append('-p')
                cmd.append(kwargs[key])
    # skip if no changes to be made
    if len(cmd) == 1:
        return (None, '', '')
    cmd.append(user)
    p = subprocess.Popen(cmd, shell=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    (out, err) = p.communicate()
    rc = p.returncode
    return (rc, out, err)

def group_exists(group):
    try:
        if group.isdigit():
            if grp.getgrgid(group):
                return True
        else:
            if grp.getgrnam(group):
                return True
    except KeyError:
        return False

def group_info(group):
    if not group_exists(group):
        return False
    if group.isdigit():
        return list(grp.getgrgid(group))
    else:
        return list(grp.getgrnam(group))

def user_group_membership(user):
    groups = []
    info = get_pwd_info(user)
    for group in grp.getgrall():
        if user in group[3]:
            groups.append(group[0])
    return groups

def user_exists(user):
    try:
        if pwd.getpwnam(user):
            return True
    except KeyError:
        return False

def get_pwd_info(user):
    if not user_exists(user):
        return False
    return list(pwd.getpwnam(user))

def user_info(user):
    if not user_exists(user):
        return False
    info = get_pwd_info(user)
    if len(info[1]) == 1 or len(info[1]) == 0:
        info[1] = user_password(user)
    return info

def user_password(user):
    passwd = ''
    if not user_exists(user):
        return passwd
    if HAVE_SPWD:
        try:
            passwd = spwd.getspnam(user)[1]
        except KeyError:
            return passwd
    else:
        # Read shadow file for user's encrypted password string
        if os.path.exists(SHADOWFILE) and os.access(SHADOWFILE, os.R_OK):
            for line in open(SHADOWFILE).readlines():
                if line.startswith('%s:' % user):
                    passwd = line.split(':')[1]
    return passwd

def get_ssh_key_path(user, ssh_file):
    info = user_info(user)
    if os.path.isabs(ssh_file):
        ssh_key_file = ssh_file
    else:
        ssh_key_file = os.path.join(info[5], ssh_file)
    return ssh_key_file

def ssh_key_gen(module, user, ssh):
    info = user_info(user)
    if not os.path.exists(info[5]):
        return (1, '', 'User %s home directory does not exist' % user)
    ssh_key_file = get_ssh_key_path(user, ssh['file'])
    ssh_dir = os.path.dirname(ssh_key_file)
    if not os.path.exists(ssh_dir):
        try:
            os.mkdir(ssh_dir, 0700)
        except OSError, e:
            return (1, '', 'Failed to create %s: %s' % (ssh_dir, str(e)))
    if os.path.exists(ssh_key_file):
        return (None, 'Key already exists', '')
    cmd = [module.get_bin_path('ssh-keygen', True)]
    for key in ssh:
        if key == 'type' and ssh[key] is not None:
            cmd.append('-t')
            cmd.append(ssh[key])
        elif key == 'bits' and ssh[key] is not None:
            cmd.append('-b')
            cmd.append(ssh[key])
        elif key == 'comment' and ssh[key] is not None:
            cmd.append('-C')
            cmd.append(ssh[key])
        elif key == 'file' and ssh[key] is not None:
            cmd.append('-f')
            cmd.append(ssh_key_file)
        elif key == 'passphrase':
            cmd.append('-N')
            if ssh[key] is not None:
                cmd.append(ssh['passphrase'])
            else:
                cmd.append('')
    p = subprocess.Popen(cmd, shell=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    (out, err) = p.communicate()
    rc = p.returncode
    return (rc, out, err)

def ssh_key_fingerprint(module, user, ssh):
    ssh_key_file = get_ssh_key_path(user, ssh['file'])
    if not os.path.exists(ssh_key_file):
        return (1, 'SSH Key file %s does not exist' % ssh_key_file, '')
    cmd = [module.get_bin_path('ssh-keygen', True)]
    cmd.append('-l')
    cmd.append('-f')
    cmd.append(ssh_key_file)
    p = subprocess.Popen(cmd, shell=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    (out, err) = p.communicate()
    rc = p.returncode
    return (rc, out, err)



# ===========================================

def main():
    ssh_defaults = {
            'bits': '2048',
            'type': 'rsa',
            'passphrase': None,
            'comment': 'ansible-generated'
    }
    ssh_defaults['file'] = os.path.join('.ssh', 'id_%s' % ssh_defaults['type'])
    ssh = dict(ssh_defaults)
    module = AnsibleModule(
        argument_spec = dict(
            state=dict(default='present', choices=['present', 'absent']),
            name=dict(required=True, aliases=['user']),
            uid=dict(default=None),
            group=dict(default=None),
            groups=dict(default=None),
            comment=dict(default=None),
            home=dict(default=None),
            shell=dict(default=None),
            password=dict(default=None),
            # following options are specific to userdel
            force=dict(default='no', choices=BOOLEANS),
            remove=dict(default='no', choices=BOOLEANS),
            # following options are specific to useradd
            createhome=dict(default='yes', choices=BOOLEANS),
            system=dict(default='no', choices=BOOLEANS),
            # following options are specific to usermod
            append=dict(default='no', choices=BOOLEANS),
            # following are specific to ssh key generation
            ssh_key=dict(choices=['generate']),
            ssh_key_bits=dict(default=ssh_defaults['bits']),
            ssh_key_type=dict(default=ssh_defaults['type']),
            ssh_key_file=dict(default=ssh_defaults['file']),
            ssh_key_comment=dict(default=ssh_defaults['comment']),
            ssh_key_passphrase=dict(default=None)
        )
    )

    state      = module.params['state']
    name       = module.params['name']
    uid        = module.params['uid']
    group      = module.params['group']
    groups     = module.params['groups']
    comment    = module.params['comment']
    home       = module.params['home']
    shell      = module.params['shell']
    password   = module.params['password']
    force      = module.params['force']
    remove     = module.params['remove']
    createhome = module.params['createhome']
    system     = module.params['system']
    append     = module.params['append']
    sshkeygen  = module.params['ssh_key']

    ssh['bits'] = module.params['ssh_key_bits']
    ssh['type'] = module.params['ssh_key_type']
    ssh['file'] = module.params['ssh_key_file']
    ssh['comment'] = module.params['ssh_key_comment']
    ssh['passphrase'] = module.params['ssh_key_passphrase']
    # If using default filename, make sure it is named appropriately
    if ssh['file'] == ssh_defaults['file']:
        ssh['file'] = os.path.join('.ssh', 'id_%s' % ssh_defaults['type'])


    rc = None
    out = ''
    err = ''
    result = {}
    result['name'] = name
    result['state'] = state
    if state == 'absent':
        if user_exists(name):
            (rc, out, err) = user_del(module, name, force=force, remove=remove)
            if rc != 0:
                module.fail_json(name=name, msg=err, rc=rc)
            result['force'] = force
            result['remove'] = remove
    elif state == 'present':
        if not user_exists(name):
            (rc, out, err) = user_add(module,
                                      name, uid=uid, group=group, groups=groups,
                                      comment=comment, home=home, shell=shell,
                                      password=password, createhome=createhome,
                                      system=system)
            result['system'] = system
            result['createhome'] = createhome
        else:
            (rc, out, err) = user_mod(module,
                                      name, uid=uid, group=group, groups=groups,
                                      comment=comment, home=home, shell=shell,
                                      password=password, append=append)
            result['append'] = append
        if rc is not None and rc != 0:
            module.fail_json(name=name, msg=err, rc=rc)
        if password is not None:
            result['password'] = 'NOT_LOGGING_PASSWORD'

    if rc is None:
        result['changed'] = False
    else:
        result['changed'] = True
    if out:
        result['stdout'] = out
    if err:
        result['stderr'] = err
    if user_exists(name):
        info = user_info(name)
        if info == False:
            result['msg'] = "failed to look up user name: %s" % name
            result['failed'] = True
        result['uid'] = info[2]
        result['group'] = info[3]
        result['comment'] = info[4]
        result['home'] = info[5]
        result['shell'] = info[6]
        groups = user_group_membership(name)
        result['uid'] = info[2]
        if len(groups) > 0:
            result['groups'] = groups
        if sshkeygen:
            (rc, out, err) = ssh_key_gen(module, name, ssh)
            if rc is not None and rc != 0:
                module.fail_json(name=name, msg=err, rc=rc)
            if rc == 0:
                result['changed'] = True
            (rc, out, err) = ssh_key_fingerprint(module, name, ssh)
            if rc == 0:
                result['ssh_fingerprint'] = out.strip()
            else:
                result['ssh_fingerprint'] = err.strip()
            result['ssh_key_file'] = get_ssh_key_path(name, ssh['file'])


    module.exit_json(**result)

# include magic from lib/ansible/module_common.py
#<<INCLUDE_ANSIBLE_MODULE_COMMON>>
main()