#!/bin/sh
# Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"). You may not
# use this file except in compliance with the License. A copy of the
# License is located at
#
# http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file is distributed
# on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
# either express or implied. See the License for the specific language governing
# permissions and limitations under the License.

DIRECTORY_ID=""
DIRECTORY_NAME=""
DIRECTORY_OU=""
REALM=""
DNS_IP_ADDRESS1=""
DNS_IP_ADDRESS2=""
LINUX_DISTRO=""
CURTIME=""
REGION=""
# https://docs.aws.amazon.com/cli/latest/userguide/install-cliv2-linux.html
AWSCLI="/usr/local/bin/aws"
# Service Creds from Secrets Manager
DOMAIN_USERNAME=""
DOMAIN_PASSWORD=""
# Secrets Manager Secret ID needs to be of the form aws/directory-services/d-91673491b6/seamless-domain-join
SECRET_ID_PREFIX="aws/directory-services"
AWS_CLI_INSTALL_DIR="$PWD/"
# Optional arguments to set hostname with or without appended digits
SET_HOSTNAME=""
SET_HOSTNAME_APPEND_NUM_DIGITS=""
MAX_APPEND_DIGITS=5

# NetBIOS computer names consist of up to 15 bytes of OEM characters
# https://docs.microsoft.com/en-us/windows/win32/sysinfo/computer-names?redirectedfrom=MSDN
NETBIOS_COMPUTER_NAME_LEN=15

get_default_hostname() {
  INSTANCE_NAME=$(hostname --short) 2> /dev/null

  # Naming conventions in Active Directory
  # https://support.microsoft.com/en-us/help/909264/naming-conventions-in-active-directory-for-computers-domains-sites-and
  RANDOM_COMPUTER_NAME=$(head -c 512 /dev/urandom | tr -dc 'a-zA-Z0-9' | fold -w 6 | head -n 1)
  COMPUTER_NAME=$(echo EC2AMAZ-$RANDOM_COMPUTER_NAME)
}

##################################################
## Set hostname to NETBIOS computer name #########
##################################################
set_hostname() {
  HOSTNAME_LEN=$(expr length $COMPUTER_NAME)
  if [ $HOSTNAME_LEN -gt $NETBIOS_COMPUTER_NAME_LEN ]; then
    echo "**Failed: Hostname exceeds NetBIOS hostname length : $HOSTNAME_LEN"
    exit 1
  fi

  HOSTNAMECTL=$(which hostnamectl)
  if [ -n "$HOSTNAMECTL" ]; then
    hostnamectl set-hostname $COMPUTER_NAME.$DIRECTORY_NAME > /dev/null
  else
    hostname $COMPUTER_NAME.$DIRECTORY_NAME > /dev/null
  fi
  if [ $? -ne 0 ]; then echo "***Failed: set_hostname(): set hostname failed" && exit 1; fi

  # https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/set-hostname.html
  if [ -f /etc/sysconfig/network ]; then
    sed -i "s/HOSTNAME=.*$//g" /etc/sysconfig/network
    echo "HOSTNAME=$COMPUTER_NAME.$DIRECTORY_NAME" >> /etc/sysconfig/network
  fi
}

get_list_of_computers() {
  LDAP_BASE_DN=$(echo $DIRECTORY_NAME | awk -F\. '{ for (i = 1; i <= NF; i++) { if (i != NF) printf "DC=%s,",$i ; else printf "DC=%s", $i} }')
  DIRNAME_UPPER=$(echo "$DIRECTORY_NAME" | tr [:lower:] [:upper:])
  ldapsearch -H ldap://$DIRECTORY_NAME -b $LDAP_BASE_DN 'objectClass=computer' -D "$DOMAIN_USERNAME@$DIRNAME_UPPER" -w "$DOMAIN_PASSWORD" | grep "cn:" | sed 's/cn://g'
}

cleanup_temp_files() {
  rm -f $AWS_CLI_INSTALL_DIR/ldap_search.txt
  rm -f $AWS_CLI_INSTALL_DIR/all_suffixes.txt
  rm -f $AWS_CLI_INSTALL_DIR/ldap_filtered_hosts.txt
  rm -f $AWS_CLI_INSTALL_DIR/possible_suffixes.txt
  rm -f $AWS_CLI_INSTALL_DIR/awscliv2.zip
}

#################################################
## find unused hostname after appending digits ##
#################################################
find_unused_hostname() {
  if [ -z $SET_HOSTNAME ]; then
    echo "**Failed: Argument SET_HOSTNAME is blank"
    exit 1
  fi

  if [ -z $SET_HOSTNAME_APPEND_NUM_DIGITS ]; then
    COMPUTER_NAME=$SET_HOSTNAME
    return
  fi

  COMPUTER_NAME=""

  NUM_DIGITS=1
  echo "SET_HOSTNAME_APPEND_NUM_DIGITS = $SET_HOSTNAME_APPEND_NUM_DIGITS"
  for i in $(seq $SET_HOSTNAME_APPEND_NUM_DIGITS); do
    NUM_DIGITS=$(expr $NUM_DIGITS \* 10)
  done
  NUM_DIGITS=$(expr $NUM_DIGITS - 1)

  # If the number of digits is 5, $AWS_CLI_INSTALL_DIR/all_suffixes.txt will
  # have numbers 00000-99999, one per line.
  # If the number of digits is 2, the file will have numbers from 00 to 99, one per line.
  for i in $(seq 0 $NUM_DIGITS); do
    if [ $SET_HOSTNAME_APPEND_NUM_DIGITS -eq 1 ]; then
      printf "%01d\n" $i
    fi
    if [ $SET_HOSTNAME_APPEND_NUM_DIGITS -eq 2 ]; then
      printf "%02d\n" $i
    fi
    if [ $SET_HOSTNAME_APPEND_NUM_DIGITS -eq 3 ]; then
      printf "%03d\n" $i
    fi
    if [ $SET_HOSTNAME_APPEND_NUM_DIGITS -eq 4 ]; then
      printf "%04d\n" $i
    fi
    if [ $SET_HOSTNAME_APPEND_NUM_DIGITS -eq 5 ]; then
      printf "%05d\n" $i
    fi
  done > $AWS_CLI_INSTALL_DIR/all_suffixes.txt

  MAX_RETRIES=5
  for i in $(seq 1 $MAX_RETRIES); do
    get_list_of_computers > $AWS_CLI_INSTALL_DIR/ldap_search.txt
    if [ ! -s $AWS_CLI_INSTALL_DIR/ldap_search.txt ]; then
      echo "**Failed: Could not get list of computers"
      exit 1
    fi

    # $AWS_CLI_INSTALL_DIR/ldap_filtered_hosts.txt has hosts
    # with the same hostname prefix
    # and then remove the hostname in each line so only the number suffixes remain.
    # This gives us the list of suffixes that already exist.
    cat $AWS_CLI_INSTALL_DIR/ldap_search.txt |
      grep -iP "^[ ]*$SET_HOSTNAME\d{$SET_HOSTNAME_APPEND_NUM_DIGITS}$" |
      sort | sed "s/$SET_HOSTNAME//gI" | tr -d ' ' > $AWS_CLI_INSTALL_DIR/ldap_filtered_hosts.txt

    # The 'comm' utility does 'set complement' between
    # $AWS_CLI_INSTALL_DIR/all_suffixes.txt
    # and $AWS_CLI_INSTALL_DIR/ldap_filtered_hosts.txt
    # to find suffixes that do not exist in ldap_filtered_hosts.
    comm -23 $AWS_CLI_INSTALL_DIR/all_suffixes.txt $AWS_CLI_INSTALL_DIR/ldap_filtered_hosts.txt > $AWS_CLI_INSTALL_DIR/possible_suffixes.txt
    if [ $? -ne 0 ]; then
      echo "**Failed: Could not execute comm utility"
      cleanup_temp_files
      exit 1
    fi

    num_lines=$(wc -l $AWS_CLI_INSTALL_DIR/possible_suffixes.txt | awk '{ print $1 }')

    if [ $num_lines -eq 0 ]; then
      echo "**Failed: Could not find unused hostnames"
      exit 1
    fi

    # When we have a list of possible hosts, one in each line, we can use
    # a random number to choose one of the unused suffixes.
    RANDOM_VALUE=$(head -c 512 /dev/urandom | xxd -p | tr -dc '0-9' | fold -w $MAX_APPEND_DIGITS | head -n 1)
    if [ -z $RANDOM_VALUE ]; then
      LINE_N=1
    else
      LINE_N=$(expr $RANDOM_VALUE % $num_lines)
      if [ $LINE_N -eq 0 -o $num_lines -eq 1 ]; then
        LINE_N=1
      fi
    fi
    HOSTNAME_SUFFIX=$(sed -n "$LINE_N"p < $AWS_CLI_INSTALL_DIR/possible_suffixes.txt)

    POSSIBLE_HOSTNAME=$SET_HOSTNAME$HOSTNAME_SUFFIX
    if [ "$POSSIBLE_HOSTNAME" = "$SET_HOSTNAME" ]; then
      echo "**Failed: Could not find a possible hostname"
      cleanup_temp_files
      exit 1
    fi

    echo "[$i] Trying possible hostname $POSSIBLE_HOSTNAME"
    cat $AWS_CLI_INSTALL_DIR/ldap_search.txt | grep -i "$POSSIBLE_HOSTNAME$"
    if [ $? -ne 0 ]; then
      echo "[$i] Found unused hostname as $POSSIBLE_HOSTNAME"
      COMPUTER_NAME=$POSSIBLE_HOSTNAME
      cleanup_temp_files
      break
    fi
  done

  cleanup_temp_files
}

##################################################
## Download AWS CLI zip file #####################
##################################################
download_awscli_zipfile() {
  MAX_RETRIES=3
  CURL_DOWNLOAD_URL="$1"
  if [ -z "$1" ]; then
    echo "***Failed: No URL argument provided for curl"
    exit 1
  fi
  echo "CURL download url is $CURL_DOWNLOAD_URL"

  MAX_RETRIES=3
  STATUS=1
  for i in $(seq 1 $MAX_RETRIES); do
    echo "[$i] Attempt installing AWS CLI by using curl"
    curl --fail $CURL_DOWNLOAD_URL -o "$AWS_CLI_INSTALL_DIR/awscliv2.zip"
    STATUS=$?
    if [ $STATUS -eq 0 ]; then
      break
    else
      curl -1 --fail $CURL_DOWNLOAD_URL -o "$AWS_CLI_INSTALL_DIR/awscliv2.zip"
      STATUS=$?
      if [ $STATUS -eq 0 ]; then
        break
      fi
    fi
    sleep 3
  done

  if [ $STATUS -ne 0 ]; then
    echo "***Failed: curl $CURL_DOWNLOAD_URL failed."
    exit 1
  fi
}

###################################################
## Check permissions of AWS CLI install directory #
###################################################
check_awscli_install_dir() {
  if [ ! -d $AWS_CLI_INSTALL_DIR ]; then
    echo "***Failed: AWS CLI install dir $AWS_CLI_INSTALL_DIR does not exist"
    exit 1
  fi

  # POSIX ls specification - https://pubs.opengroup.org/onlinepubs/9699919799/
  # If the -l option is specified, the following information
  # shall be written for files other than character special
  # and block special files: "%s %u %s %s %u %s %s\n", <file mode>, <number of links>, <owner name>, <group name>, <size>, <date and time>, <pathname>
  ls -ld $AWS_CLI_INSTALL_DIR
  if [ $? -ne 0 ]; then
    echo "***Failed: ls -ld $AWS_CLI_INSTALL_DIR"
    exit 1
  fi
  AWS_CLI_INSTALL_DIR_LS_LD=$(ls -ld $AWS_CLI_INSTALL_DIR)
  AWS_CLI_INSTALL_DIR_OWNER=$(echo $AWS_CLI_INSTALL_DIR_LS_LD | awk '{print $3}')
  id -u $AWS_CLI_INSTALL_DIR_OWNER
  if [ $? -ne 0 ]; then
    echo "***Failed: id command"
    exit 1
  fi
  ID=$(id -u $AWS_CLI_INSTALL_DIR_OWNER)
  if [ $ID != "0" ]; then
    echo "***Failed: AWS CLI install dir user is not root"
    exit 1
  fi
  AWS_CLI_INSTALL_DIR_PERMISSIONS=$(echo $AWS_CLI_INSTALL_DIR_LS_LD | awk '{print $1}' | cut -c 5-)
  if echo "$AWS_CLI_INSTALL_DIR_PERMISSIONS" | grep -e "------"; then
    echo "Permissions check successful for AWS CLI install directory"
  else
    echo "***Failed: Wrong permissions for $AWS_CLI_INSTALL_DIR_PERMISSIONS"
    exit 1
  fi
}

##################################################
########## Install components ####################
##################################################
install_components() {
  LINUX_DISTRO=$(cat /etc/os-release | grep NAME | awk -F'=' '{print $2}')
  LINUX_DISTRO_VERSION_ID=$(cat /etc/os-release | grep VERSION_ID | awk -F'=' '{print $2}' | tr -d '"')
  if [ -z $LINUX_DISTRO_VERSION_ID ]; then
    echo "**Failed : Unsupported OS version $LINUX_DISTRO : $LINUX_DISTRO_VERSION_ID"
    exit 1
  fi

  if grep 'CentOS' /etc/os-release 1> /dev/null 2> /dev/null; then
    if [ "$LINUX_DISTRO_VERSION_ID" -lt "7" ]; then
      echo "**Failed : Unsupported OS version $LINUX_DISTRO : $LINUX_DISTRO_VERSION_ID"
      exit 1
    fi
    LINUX_DISTRO='CentOS'
    # yum -y update
    ## yum update takes too long
    yum -y install realmd adcli oddjob-mkhomedir oddjob samba-winbind-clients samba-winbind samba-common-tools samba-winbind-krb5-locator krb5-workstation unzip bind-utils python3 openldap-clients vim-common > /dev/null
    if [ $? -ne 0 ]; then echo "install_components(): yum install errors for CentOS" && return 1; fi
  elif grep -e 'Red Hat' /etc/os-release 1> /dev/null 2> /dev/null; then
    LINUX_DISTRO='RHEL'
    RHEL_MAJOR_VERSION=$(echo $LINUX_DISTRO_VERSION_ID | awk -F'.' '{print $1}')
    RHEL_MINOR_VERSION=$(echo $LINUX_DISTRO_VERSION_ID | awk -F'.' '{print $2}')
    if [ $RHEL_MAJOR_VERSION -eq "7" ] && [ ! -z $RHEL_MINOR_VERSION ] && [ $RHEL_MINOR_VERSION -lt "6" ]; then
      # RHEL 7.5 and below are not supported
      echo "**Failed : Unsupported OS version $LINUX_DISTRO : $LINUX_DISTRO_VERSION_ID"
      exit 1
    fi
    if [ $RHEL_MAJOR_VERSION -eq "7" ] && [ -z $RHEL_MINOR_VERSION ]; then
      # RHEL 7 is not supported
      echo "**Failed : Unsupported OS version $LINUX_DISTRO : $LINUX_DISTRO_VERSION_ID"
      exit 1
    fi
    # yum -y update
    ## yum update takes too long
    # https://access.redhat.com/documentation/en-us/red_hat_enterprise_linux/8/html-single/deploying_different_types_of_servers/index
    yum -y install realmd adcli oddjob-mkhomedir oddjob samba-winbind-clients samba-winbind samba-common-tools samba-winbind-krb5-locator krb5-workstation vim unzip bind-utils python3 openldap-clients NetworkManager > /dev/null
    if [ $? -ne 0 ]; then echo "install_components(): yum install errors for Red Hat" && return 1; fi
  elif grep -e 'Fedora' /etc/os-release 1> /dev/null 2> /dev/null; then
    LINUX_DISTRO='Fedora'
    ## yum update takes too long, but it is unavoidable here.
    yum -y update
    yum -y install realmd adcli oddjob-mkhomedir oddjob samba-winbind-clients samba-winbind samba-common-tools samba-winbind-krb5-locator krb5-workstation vim unzip bind-utils python3 openldap-clients vim-common NetworkManager > /dev/null
    if [ $? -ne 0 ]; then echo "install_components(): yum install errors for Fedora" && return 1; fi
    systemctl restart dbus
  elif grep 'Amazon Linux' /etc/os-release 1> /dev/null 2> /dev/null; then
    LINUX_DISTRO='AMAZON_LINUX'
    # yum -y update
    ## yum update takes too long
    yum -y install realmd adcli oddjob-mkhomedir oddjob samba-winbind-clients samba-winbind samba-common-tools samba-winbind-krb5-locator krb5-workstation unzip bind-utils python3 openldap-clients vim-common network-manager > /dev/null
    if [ $? -ne 0 ]; then echo "install_components(): yum install errors for Amazon Linux" && return 1; fi
  elif grep 'Ubuntu' /etc/os-release 1> /dev/null 2> /dev/null; then
    LINUX_DISTRO='UBUNTU'
    UBUNTU_MAJOR_VERSION=$(echo $LINUX_DISTRO_VERSION_ID | awk -F'.' '{print $1}')
    UBUNTU_MINOR_VERSION=$(echo $LINUX_DISTRO_VERSION_ID | awk -F'.' '{print $2}')
    if [ $UBUNTU_MAJOR_VERSION -lt "14" ]; then
      # Ubuntu versions below 14.04 are not supported
      echo "**Failed : Unsupported OS version $LINUX_DISTRO : $LINUX_DISTRO_VERSION_ID"
      exit 1
    fi
    # set DEBIAN_FRONTEND variable to noninteractive to skip any interactive post-install configuration steps.
    export DEBIAN_FRONTEND=noninteractive
    apt-get -y update
    if [ $? -ne 0 ]; then echo "install_components(): apt-get update errors for Ubuntu" && return 1; fi
    apt-get -yq install realmd adcli winbind samba libnss-winbind libpam-winbind libpam-krb5 krb5-config krb5-locales krb5-user packagekit ntp unzip dnsutils python3 ldap-utils network-manager > /dev/null
    if [ $? -ne 0 ]; then echo "install_components(): apt-get install errors for Ubuntu" && return 1; fi
    # Disable Reverse DNS resolution. Ubuntu Instances must be reverse-resolvable in DNS before the realm will work.
    sed -i "s/default_realm.*$/default_realm = $REALM\n\trdns = false/g" /etc/krb5.conf
    if [ $? -ne 0 ]; then
      echo "install_components(): access errors to /etc/krb5.conf"
      return 1
    fi
    if ! grep "Ubuntu 16.04" /etc/os-release 2> /dev/null; then
      pam-auth-update --enable mkhomedir
    fi
  elif grep 'SUSE Linux' /etc/os-release 1> /dev/null 2> /dev/null; then
    SUSE_MAJOR_VERSION=$(echo $LINUX_DISTRO_VERSION_ID | awk -F'.' '{print $1}')
    SUSE_MINOR_VERSION=$(echo $LINUX_DISTRO_VERSION_ID | awk -F'.' '{print $2}')
    if [ "$SUSE_MAJOR_VERSION" -lt "15" ]; then
      echo "**Failed : Unsupported OS version $LINUX_DISTRO : $LINUX_DISTRO_VERSION_ID"
      exit 1
    fi
    if [ "$SUSE_MAJOR_VERSION" -eq "15" ]; then
      sudo SUSEConnect -p PackageHub/15.1/x86_64
    fi
    LINUX_DISTRO='SUSE'
    sudo zypper update -y
    sudo zypper -n install realmd adcli sssd sssd-tools sssd-ad samba-client krb5-client samba-winbind krb5-client bind-utils python3 openldap2-client NetworkManager
    if [ $? -ne 0 ]; then
      return 1
    fi
  elif grep 'Debian' /etc/os-release; then
    DEBIAN_MAJOR_VERSION=$(echo $LINUX_DISTRO_VERSION_ID | awk -F'.' '{print $1}')
    DEBIAN_MINOR_VERSION=$(echo $LINUX_DISTRO_VERSION_ID | awk -F'.' '{print $2}')
    if [ "$DEBIAN_MAJOR_VERSION" -lt "9" ]; then
      echo "**Failed : Unsupported OS version $LINUX_DISTRO : $LINUX_DISTRO_VERSION_ID"
      exit 1
    fi
    apt-get -y update
    LINUX_DISTRO='DEBIAN'
    DEBIAN_FRONTEND=noninteractive apt-get -yq install realmd adcli winbind samba libnss-winbind libpam-winbind libpam-krb5 krb5-config krb5-locales krb5-user packagekit ntp unzip dnsutils python3 ldapscripts network-manager > /dev/null
    if [ $? -ne 0 ]; then
      return 1
    fi
  fi

  check_awscli_install_dir
  if uname -a | grep -e "x86_64" -e "amd64"; then
    download_awscli_zipfile "https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip"
  elif uname -a | grep "aarch64"; then
    download_awscli_zipfile "https://awscli.amazonaws.com/awscli-exe-linux-aarch64.zip"
  else
    echo "***Failed: install_components processor type is unsupported." && exit 1
  fi

  unzip -o "$AWS_CLI_INSTALL_DIR"/awscliv2.zip 1> /dev/null
  if [ $? -ne 0 ]; then echo "***Failed: unzip of awscliv2.zip" && exit 1; fi
  "$AWS_CLI_INSTALL_DIR"/aws/install -u 1> /dev/null
  if [ $? -ne 0 ]; then echo "***Failed: aws cli install" && exit 1; fi
  return 0
}

######################################################################
##### Retrieve Service Account Credentials from Secrets Manager   ####
######################################################################
get_servicecreds() {
  SECRET_ID="${SECRET_ID_PREFIX}/$DIRECTORY_ID/seamless-domain-join"
  SECRET_VALUE=$($AWSCLI secretsmanager get-secret-value --secret-id "$SECRET_ID" --region $REGION --query "SecretString" --output text 2> /dev/null)
  if [ $? -ne 0 ]; then
    PARENT_DIRECTORY_ID=$($AWSCLI ds describe-directories --region $REGION --query "DirectoryDescriptions[?DirectoryId =='$DIRECTORY_ID'].OwnerDirectoryDescription.DirectoryId | [0]" | sed 's/"//g')
    if [ $? -ne 0 ] || [ -z "$PARENT_DIRECTORY_ID" ] || [ "$PARENT_DIRECTORY_ID" = null ]; then
      echo "***Failed: Cannot find parent directory Id"
      exit 1
    fi
    PARENT_ACCOUNT_ID=$($AWSCLI ds describe-directories --region $REGION --query "DirectoryDescriptions[?DirectoryId =='$DIRECTORY_ID'].OwnerDirectoryDescription.AccountId | [0]" | sed 's/"//g')
    if [ $? -ne 0 ] || [ -z "$PARENT_ACCOUNT_ID" ] || [ "$PARENT_ACCOUNT_ID" = null ]; then
      echo "***Failed: Cannot find parent account Id"
      exit 1
    fi
    SECRET_ID="arn:aws:secretsmanager:${REGION}:${PARENT_ACCOUNT_ID}:secret:aws/directory-services/${PARENT_DIRECTORY_ID}/seamless-domain-join"
    SECRET_VALUE=$($AWSCLI secretsmanager get-secret-value --secret-id "$SECRET_ID" --region $REGION --query 'SecretString' --output text 2> /dev/null)
    if [ $? -ne 0 ] || [ -z "$SECRET_VALUE" ]; then
      echo "***Failed: aws secretsmanager get-secret-value"
      exit 1
    fi
  fi

  if [ -n "$(command -v python3)" ]; then
    PYTHON=$(which python3)
  else
    PYTHON=$(which python)
  fi
  DOMAIN_USERNAME=$(echo "$SECRET_VALUE" | $PYTHON -c 'import sys, json; obj=json.load(sys.stdin); print(obj["awsSeamlessDomainUsername"])')
  if [ $? -ne 0 ] || [ -z "$DOMAIN_USERNAME" ]; then
    echo "***Failed: Invalid DOMAIN_USERNAME"
    exit 1
  fi
  DOMAIN_PASSWORD=$(echo "$SECRET_VALUE" | $PYTHON -c 'import sys, json; obj=json.load(sys.stdin); print(obj["awsSeamlessDomainPassword"])')
  if [ $? -ne 0 ] || [ -z "$DOMAIN_PASSWORD" ]; then
    echo "***Failed: Invalid DOMAIN_PASSWORD"
    exit 1
  fi
}

##################################################
## Setup resolv.conf and also dhclient.conf ######
## to prevent overwriting of resolv.conf    ######
##################################################
setup_resolv_conf_and_dhclient_conf() {
  if [ ! -z "$DNS_IP_ADDRESS1" ] && [ ! -z "$DNS_IP_ADDRESS2" ]; then
    touch /etc/resolv.conf
    mv /etc/resolv.conf /etc/resolv.conf.backup."$CURTIME"
    echo ";Generated by Domain Join SSMDocument" > /etc/resolv.conf
    echo "search $DIRECTORY_NAME" >> /etc/resolv.conf
    echo "nameserver $DNS_IP_ADDRESS1" >> /etc/resolv.conf
    echo "nameserver $DNS_IP_ADDRESS2" >> /etc/resolv.conf
    touch /etc/dhcp/dhclient.conf
    mv /etc/dhcp/dhclient.conf /etc/dhcp/dhclient.conf.backup."$CURTIME"
    echo "supersede domain-name-servers $DNS_IP_ADDRESS1, $DNS_IP_ADDRESS2;" > /etc/dhcp/dhclient.conf
  elif [ ! -z "$DNS_IP_ADDRESS1" ] && [ -z "$DNS_IP_ADDRESS2" ]; then
    touch /etc/resolv.conf
    mv /etc/resolv.conf /etc/resolv.conf.backup."$CURTIME"
    echo ";Generated by Domain Join SSMDocument" > /etc/resolv.conf
    echo "search $DIRECTORY_NAME" >> /etc/resolv.conf
    echo "nameserver $DNS_IP_ADDRESS1" >> /etc/resolv.conf
    touch /etc/dhcp/dhclient.conf
    mv /etc/dhcp/dhclient.conf /etc/dhcp/dhclient.conf.backup."$CURTIME"
    echo "supersede domain-name-servers $DNS_IP_ADDRESS1;" > /etc/dhcp/dhclient.conf
  elif [ -z "$DNS_IP_ADDRESS1" ] && [ ! -z "$DNS_IP_ADDRESS2" ]; then
    touch /etc/resolv.conf
    mv /etc/resolv.conf /etc/resolv.conf.backup."$CURTIME"
    echo ";Generated by Domain Join SSMDocument" > /etc/resolv.conf
    echo "search $DIRECTORY_NAME" >> /etc/resolv.conf
    echo "nameserver $DNS_IP_ADDRESS2" >> /etc/resolv.conf
    touch /etc/dhcp/dhclient.conf
    mv /etc/dhcp/dhclient.conf /etc/dhcp/dhclient.conf.backup."$CURTIME"
    echo "supersede domain-name-servers $DNS_IP_ADDRESS2;" > /etc/dhcp/dhclient.conf
  else
    echo "***Failed: No DNS IPs available" && exit 1
  fi
}

##################################################
## Set PEER_DNS to yes ###########################
##################################################
set_peer_dns() {
  for f in $(ls /etc/sysconfig/network-scripts/ifcfg-*); do
    if echo $f | grep "lo"; then
      continue
    fi
    if ! grep PEERDNS $f; then
      echo "" >> $f
      echo PEERDNS=yes >> $f
    fi
  done
}

##################################################
## Print shell variables #########################
##################################################
print_vars() {
  echo "REGION = $REGION"
  echo "DIRECTORY_ID = $DIRECTORY_ID"
  echo "DIRECTORY_NAME = $DIRECTORY_NAME"
  echo "DIRECTORY_OU = $DIRECTORY_OU"
  echo "REALM = $REALM"
  echo "DNS_IP_ADDRESS1 = $DNS_IP_ADDRESS1"
  echo "DNS_IP_ADDRESS2 = $DNS_IP_ADDRESS2"
  echo "COMPUTER_NAME = $COMPUTER_NAME"
  echo "hostname = $(hostname)"
  echo "LINUX_DISTRO = $LINUX_DISTRO"
}

#########################################################
## Add FQDN and Hostname to Hosts file for below error ##
# No DNS domain configured for ip-172-31-12-23.         #
# Unable to perform DNS Update.                         #
#########################################################
configure_hosts_file() {
  fullhost="${COMPUTER_NAME}.${DIRECTORY_NAME}" # ,, means lowercase since bash v4
  ip_address="$(ip -o -4 addr show eth0 | awk '{print $4}' | cut -d/ -f1)"
  cleanup_comment='# Generated by Domain Join SSMDocument'
  sed -i".orig" -r "/^.*${cleanup_comment}/d;\
        /^127.0.0.1\s+localhost\s*/a\\${ip_address} ${fullhost} ${COMPUTER_NAME} ${cleanup_comment}" /etc/hosts
}

#########################################################
## Restart NetworkManager if present ####################
#########################################################
restart_network_manager() {
  if [ -d /etc/NetworkManager/conf.d/ ]; then
    systemctl restart NetworkManager
    if [ $? -ne 0 ]; then
      service NetworkManager restart
    fi
    if [ $? -ne 0 ]; then
      echo "**Failed: NetworkManager failed"
      exit 1
    fi
  fi
}

##################################################
## Add AWS Directory Service DNS IP Addresses as #
## primary to the resolv.conf and dhclient       #
## configuration files.                          #
##################################################
do_dns_config() {
  # Prevent Network Manager from overwriting resolv.conf
  # https://access.redhat.com/documentation/en-us/red_hat_enterprise_linux/8/html/configuring_and_managing_networking/manually-configuring-the-etc-resolv-conf-file_configuring-and-managing-networking
  if [ -d /etc/NetworkManager/conf.d/ ]; then
    echo "[main]" > /etc/NetworkManager/conf.d/90-dns-none.conf
    echo "dns=none" >> /etc/NetworkManager/conf.d/90-dns-none.conf
    if [ $? -ne 0 ]; then
      echo "**Failed: Writing to /etc/NetworkManager/conf.d/90-dns-none.conf"
      exit 1
    fi
    restart_network_manager
  fi

  setup_resolv_conf_and_dhclient_conf
  if [ $LINUX_DISTRO = 'AMAZON_LINUX' ]; then
    set_peer_dns
  fi

  if [ $LINUX_DISTRO = "UBUNTU" ]; then
    if [ -d /etc/netplan ]; then
      # Ubuntu 18.04
      cat << EOF | tee /etc/netplan/99-custom-dns.yaml
network:
    version: 2
    ethernets:
        eth0:
            nameservers:
                addresses: [$DNS_IP_ADDRESS1, $DNS_IP_ADDRESS2]
            dhcp4-overrides:
                use-dns: false
EOF
      netplan apply
      if [ $? -ne 0 ]; then echo "***Failed: do_dns_config(): netplan apply failed" && exit 1; fi
      # Seems to fail otherwise
      sleep 15
    fi
  fi

  if [ $LINUX_DISTRO = "RHEL" ] || [ $LINUX_DISTRO = "Fedora" ]; then
    set_peer_dns
    if [ -f /etc/NetworkManager/NetworkManager.conf ]; then
      cp /etc/NetworkManager/NetworkManager.conf /etc/NetworkManager/NetworkManager.conf."$CURTIME"
      cat /etc/NetworkManager/NetworkManager.conf."$CURTIME" | sed "s/\[main\]/[main]\ndns=none/g" > /etc/NetworkManager/NetworkManager.conf
    fi
  fi

  if [ $LINUX_DISTRO = "CentOS" ]; then
    set_peer_dns
  fi

  restart_network_manager
}

##################################################
## Resolve domain name to IP address(es)        ##
## by using nslookup command                    ##
##################################################
resolve_name_to_ip() {
  (
    nslookup "$1" | tail -n +3 | sed -n 's/Address:\s*//p'
  ) && return 0 || return 1
}

##################################################
## DNS may already be reachable if DHCP option  ##
## sets are used.                               ##
##################################################
is_directory_reachable() {
  grep "Generated by Domain Join SSMDocument" /etc/resolv.conf
  if [ $? -ne 0 ]; then
    echo "**Failed: /etc/resolv.conf was overwritten"
    exit 1
  fi
  DNS_IPS=$(resolve_name_to_ip $DIRECTORY_NAME)
  if [ $? -ne 0 ]; then
    echo "***Failed: Cannot resolve domain name $DIRECTORY_NAME" && return 1
  fi
  echo -e "Successfully resolve domain name $DIRECTORY_NAME to IP address(es):\n$DNS_IPS"

  return 0
}

##################################################
## Join Linux instance to AWS Directory Service ##
##################################################
do_domainjoin() {
  MAX_RETRIES=10
  for i in $(seq 1 $MAX_RETRIES); do
    if [ -z "$DIRECTORY_OU" ]; then
      LOG_MSG=$(echo $DOMAIN_PASSWORD | realm join --client-software=winbind -U ${DOMAIN_USERNAME}@${DIRECTORY_NAME} "$DIRECTORY_NAME" -v 2>&1)
    else
      LOG_MSG=$(echo $DOMAIN_PASSWORD | realm join --client-software=winbind -U ${DOMAIN_USERNAME}@${DIRECTORY_NAME} "$DIRECTORY_NAME" --computer-ou="$DIRECTORY_OU" -v 2>&1)
    fi
    STATUS=$?
    if [ $STATUS -eq 0 ]; then
      break
    else
      if echo "$LOG_MSG" | grep -q "Already joined to this domain"; then
        echo "do_domainjoin(): Already joined to this domain : $LOG_MSG"
        STATUS=0
        break
      fi
    fi
    sleep 10
  done

  if [ $STATUS -ne 0 ]; then
    echo "***Failed: realm join failed" && exit 1
  fi
  echo "########## SUCCESS: realm join successful $LOG_MSG ##########"
}

##############################
## Configure nsswitch.conf  ##
##############################
config_nsswitch() {
  # Edit nsswitch config
  NSSWITCH_CONF_FILE=/etc/nsswitch.conf
  sed -i 's/^\s*passwd:.*$/passwd:     compat winbind/' $NSSWITCH_CONF_FILE
  sed -i 's/^\s*group:.*$/group:      compat winbind/' $NSSWITCH_CONF_FILE
  sed -i 's/^\s*shadow:.*$/shadow:     compat winbind/' $NSSWITCH_CONF_FILE
}

###################################################
## Configure id-mappings in Samba                ##
###################################################
config_samba() {
  AD_INFO=$(adcli info ${DIRECTORY_NAME} | grep '^domain-short = ' | awk '{print $3}')
  sed -i".pre-join" -r "/^\[global\]/a\\
        idmap config * : backend = autorid\n\
        idmap config * : range = 100000000-2100000000\n\
        idmap config * : rangesize = 100000000\n\
        idmap config ${AD_INFO} : backend = rid\n\
        idmap config ${AD_INFO} : range = 65536 - 99999999\n\
        winbind refresh tickets = yes\n\
        kerberos method = secrets and keytab\n\
        winbind enum groups = no\n\
        winbind enum users = no
        /^\s*idmap/d;\
        /^\s*kerberos\s+method/d;\
        /^\s*winbind\s+refresh/d;\
        /^\s*winbind\s+enum/d" \
    /etc/samba/smb.conf

  cp /etc/samba/smb.conf /tmp

  # Flushing Samba Winbind databases
  net cache flush

  # Restarting Winbind daemon
  service winbind restart
}

reconfigure_samba() {
  sed -i 's/kerberos method = system keytab/kerberos method = secrets and keytab/g' /etc/samba/smb.conf
  service winbind restart
  if [ $? -ne 0 ]; then
    systemctl restart winbind
    if [ $? -ne 0 ]; then
      service winbind restart
    fi
  fi
}

##################################################
## Main entry point ##############################
##################################################
CURTIME=$(date | sed 's/ //g')

if [ $# -eq 0 ]; then
  exit 1
fi

for i in "$@"; do
  case "$i" in
    --directory-id)
      shift
      DIRECTORY_ID="$1"
      continue
      ;;
    --directory-name)
      shift
      DIRECTORY_NAME="$1"
      continue
      ;;
    --directory-ou)
      shift
      DIRECTORY_OU="$1"
      continue
      ;;
    --instance-region)
      shift
      REGION="$1"
      continue
      ;;
    --dns-addresses)
      shift
      DNS_ADDRESSES="$1"
      DNS_IP_ADDRESS1=$(echo $DNS_ADDRESSES | awk -F',' '{ print $1 }')
      DNS_IP_ADDRESS2=$(echo $DNS_ADDRESSES | awk -F',' '{ print $2 }')
      continue
      ;;
    --proxy-address)
      shift
      PROXY_ADDRESS="$1"
      continue
      ;;
    --no-proxy)
      shift
      NO_PROXY="$1"
      continue
      ;;
    --set-hostname)
      shift
      SET_HOSTNAME="$1"
      continue
      ;;
    --set-hostname-append-num-digits)
      shift
      SET_HOSTNAME_APPEND_NUM_DIGITS="$1"
      continue
      ;;
  esac
  shift
done

if [ -z $REGION ]; then
  echo "***Failed: No Region found" && exit 1
fi

if [ ! -z $SET_HOSTNAME_APPEND_NUM_DIGITS ]; then
  if [ $SET_HOSTNAME_APPEND_NUM_DIGITS -le 0 -o \
    $SET_HOSTNAME_APPEND_NUM_DIGITS -gt $MAX_APPEND_DIGITS ]; then
    echo "**Failed: Invalid number of append digits $SET_HOSTNAME_APPEND_NUM_DIGITS"
    exit 1
  fi
fi

# Deal with scenario where this script is run again after the domain is already joined.
# We want to avoid rerunning as the set_hostname function can change the hostname of a server that is already
# domain joined and cause a mismatch.
realm list 2> /dev/null | grep -q "domain-name: ${DIRECTORY_NAME}\$"
if [ $? -eq 0 ]; then
  echo "########## SKIPPING Domain Join: ${DIRECTORY_NAME} already joined  ##########"
  exit 0
fi

REALM=$(echo "$DIRECTORY_NAME" | tr [a-z] [A-Z])

MAX_RETRIES=8
for i in $(seq 1 $MAX_RETRIES); do
  echo "[$i] Attempt installing components"
  install_components
  if [ $? -eq 0 ]; then
    break
  fi
  sleep 30
done

if [ -z $DNS_IP_ADDRESS1 ] && [ -z $DNS_IP_ADDRESS2 ]; then
  DNS_ADDRESSES=$($AWSCLI ds describe-directories --region $REGION --directory-id $DIRECTORY_ID --output text | grep DNSIPADDR | awk '{print $2}')
  if [ $? -ne 0 ]; then
    echo "***Failed: DNS IPs not found" && exit 1
  fi
  DNS_IP_ADDRESS1=$(echo $DNS_ADDRESSES | awk '{ print $1 }')
  DNS_IP_ADDRESS2=$(echo $DNS_ADDRESSES | awk '{ print $2 }')
fi

## Configure DNS even if DHCP option set is used.
do_dns_config
get_servicecreds

COMPUTER_NAME=$(hostname --short)

if [ ! -z $SET_HOSTNAME ]; then
  find_unused_hostname "$SET_HOSTNAME"
  if [ -z $COMPUTER_NAME ]; then
    echo "**Failed: computername is empty"
    exit 1
  fi
else
  get_default_hostname
fi

set_hostname
configure_hosts_file

sed -i 's/PasswordAuthentication no/PasswordAuthentication yes/g' /etc/ssh/sshd_config
systemctl restart sshd
if [ $? -ne 0 ]; then
  systemctl restart ssh
  if [ $? -ne 0 ]; then
    service sshd restart
  fi
  if [ $? -ne 0 ]; then
    service ssh restart
  fi
fi

print_vars
is_directory_reachable
if [ $? -eq 0 ]; then
  config_nsswitch
  config_samba
  do_domainjoin
  reconfigure_samba
else
  echo "**Failed: Unable to reach DNS server"
  exit 1
fi

echo "Success"
exit 0
