#!/usr/libexec/platform-python

#
# lookup CERN LDAP for user info
# v. 0.7 07 Oct 2022 Manuel Guijarro <manuel.guijarro@cern.ch>
#                   - allowed apostrophe in gecos
# v. 0.6 20 Jul 2022 Manuel Guijarro <manuel.guijarro@cern.ch>
#		    - added option "-s" to install user default shell if not present
# v. 0.5 29 Apr 2022 Ben Morrice <ben.morrice@cern.ch>
#                   - port from perl to python
# v. 0.4 13 Feb 2017 Jaroslaw.Polok <jaroslaw.polok@cern.ch>
#                   - add check for shell existence.
# v. 0.1 27 Oct 2008 Jaroslaw.Polok <jaroslaw.polok@cern.ch>

import argparse
import ldap
import os
import subprocess
import sys

ldap_server = 'xldap.cern.ch'
ldap_user_base = 'ou=users,ou=organic units,dc=cern,dc=ch'
ldap_group_base = 'ou=unix,ou=workgroups,dc=cern,dc=ch'
ldap_user_attrs = [ 'cn', 'loginShell', 'uidNumber', 'gidNumber', 'unixHomeDirectory', 'gecos' ]
ldap_group_attrs = [ 'cn', 'gidNumber' ]

def do_execute (cmd, verbose):
  if verbose:
    print(f"running: {cmd}")
  try:
    p = subprocess.Popen(cmd, stderr=subprocess.PIPE, shell=True, stdout=subprocess.PIPE)
  except:
    print(f"Failed to execute: {cmd}")
    sys.exit(1)

  out, err = p.communicate()
  if len(err) > 0:
    print(err.decode().strip())
  if p.wait() != 0:
    print('Error executing command.')
    sys.exit(p.wait())

def do_ldap_init():
  try:
    conn = ldap.initialize(f"ldap://{ldap_server}:389")
  except conn.LDAPError as error_message:
    print(f"Cannot connect to ldap. Error: {error_message}")
    sys.exit(1)
  return conn

def do_ldap_search (conn, base, search, attrs, verbose):
  if verbose:
    print(f"LDAP search: {search}")
  result = conn.search_s(base, ldap.SCOPE_SUBTREE, search, attrs)
  if verbose:
    print(f"LDAP return: {len(result)} entries.")
  if len(result) == 0:
    if 'gecos' in attrs:
      print('No match in user database')
    else:
      print('No match in group database')
    sys.exit(1)

  incomplete_record = 0
  for attr in attrs:
    if attr not in result[0][1]:
      incomplete_record = 1
  if incomplete_record:
    if 'cn' in result[0][1]:
      print(f"Incomplete LDAP data for: ${result[0][1]['cn']}, skipping record.")
    else:
      if 'gecos' in attrs:
        print("LDAP data incoherent (object is NOT an 'organizationalPerson' ?), skipping record.")
      else:
        print("LDAP data incoherent (object is NOT a 'group' ?), skipping record.")
    sys.exit(1)
  for k,v in result[0][1].items():
    result[0][1][k] = v[0].decode()

  return result[0][1]

def main():

  if os.geteuid() != 0:
    print('You must be root in order to add users')
    sys.exit(1)

  parser = argparse.ArgumentParser()
  parser.add_argument('-l', '--login', required=False, type=str, action='store', help='Add user with given LOGINID')
  parser.add_argument('-v', '--verbose', required=False,action='store_true', help=argparse.SUPPRESS)
  parser.add_argument('-s', '--shell', required=False, action='store_true', help='Install user shell if not present')
  parser.add_argument('-d', '--directory', required=False, action='store_true', help='Use homedir from LDAP (default is a local homedir)')
  parser.add_argument('-u', '--update', required=False, action='store_true', help='Update local user with content from LDAP (usermod versus useradd)')
  parser.add_argument('-f', '--forceprimarygid', required=False, type=int, default=-1, help='Use this gid for the primary group instead of content from LDAP')
  parser.add_argument('--home', required=False, type=str, default='/home', help='The base of the home area')
  parser.add_argument('--home-under-lower-first-letter', required=False, action='store_true', help='Use a subdir of the format "/home/t/testuser"')

  # This logic ensures that single dash shortname arguments work (eg: '-log morrice')
  new_argv = []
  for arg in sys.argv:
    if arg.startswith('-') and not arg.startswith('--') and len(arg) > 2:
      arg = '-' + str(arg)
    new_argv.append(str(arg))
  sys.argv = new_argv

  args, unknown = parser.parse_known_args()
  if not args.login and len(unknown) == 1:
    args.login = unknown[0]

  # --directory was not passed, let's see if /etc/sysconfig/locmap-initialsetup
  # contains LOCMAP_HOMEDIRECTORY_LOCAL
  if not args.directory:
    try:
      lfh = open('/etc/sysconfig/locmap-initialsetup', 'r', encoding='utf-8')
      for line in lfh:
        name, var = line.partition("=")[::2]
        if 'LOCMAP_HOMEDIRECTORY_LOCAL' in name.strip() and 'False' in var.strip():
          # User has selected /afs, which is different from the default
          print("INFO: 'LOCMAP_HOMEDIRECTORY_LOCAL=False' found in '/etc/sysconfig/locmap-initialsetup', forcing --directory flag")
          args.directory = True
    except FileNotFoundError:
      pass

  if args.login:
    conn = do_ldap_init()
    user_results = do_ldap_search(conn, ldap_user_base, f"(&(cn={args.login})(objectClass=organizationalPerson))", ldap_user_attrs, args.verbose)
    if args.forceprimarygid >= 0:
      user_results['gidNumber'] = args.forceprimarygid
    group_results = do_ldap_search(conn, ldap_group_base, f"(&(gidNumber={user_results['gidNumber']})(objectClass=group))", ldap_group_attrs, args.verbose)
    do_execute(f"/usr/sbin/groupadd -g \"{group_results['gidNumber']}\" -f \"{group_results['cn']}\"", args.verbose)
    if not os.path.isfile(user_results['loginShell']):
      print(f"WARN: user shell {user_results['loginShell']} is not installed on this system.")
      if args.shell:
        print(f"INFO: installing user shell {user_results['loginShell']}")
        do_execute(f"/usr/bin/yum -y install \"{user_results['loginShell']}\"", args.verbose)
    common_arguments = f"-c \"{user_results['gecos']}\" -g \"{user_results['gidNumber']}\" -s \"{user_results['loginShell']}\" -u \"{user_results['uidNumber']}\""
    user = f"{user_results['cn']}"
    if args.update:
      binary = '/usr/sbin/usermod'
      unique_arguments = ''
    else:
      binary = '/usr/sbin/useradd'
      if args.directory:
        unique_arguments = "-N -M"
      else:
        unique_arguments = '-N -m'
    if args.directory:
      homedir = f"-d {user_results['unixHomeDirectory']}"
    else:
      if args.home_under_lower_first_letter:
        shard = user_results['cn'][0].lower()
        homedir = f"-d {args.home}/{shard}/{user_results['cn']}"
      else:
        homedir = f"-d {args.home}/{user_results['cn']}"
    do_execute(f"{binary} {unique_arguments} {common_arguments} {homedir} {user}", args.verbose)
  else:
    parser.print_help(sys.stdout)

if __name__ == "__main__":
    main()
