#!/bin/sh
# Copyright 2017 Google Inc. 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.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License 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.

is_freebsd() {
  [ "$(uname)" = "FreeBSD" ]
  return $?
}

nss_config="/etc/nsswitch.conf"
pam_sshd_config="/etc/pam.d/sshd"
pam_su_config="/etc/pam.d/su"
sshd_config="/etc/ssh/sshd_config"
group_config="/etc/security/group.conf"
sudoers_dir="/var/google-sudoers.d"
users_dir="/var/google-users.d"
added_comment="# Added by Google Compute Engine OS Login."
sshd_block="#### Google OS Login control. Do not edit this section. ####"
sshd_end_block="#### End Google OS Login control section. ####"
sudoers_file="/etc/sudoers.d/google-oslogin"
if is_freebsd; then
  sudoers_file="/usr/local/etc/sudoers.d/google-oslogin"
fi

# Update nsswitch.conf to include OS Login NSS module for passwd.
modify_nsswitch_conf() {
  local nss_config="${1:-${nss_config}}"

  if ! grep -q '^passwd:.*oslogin' "$nss_config"; then
    $sed -i"" '/^passwd:/ s/$/ cache_oslogin oslogin/' "$nss_config"
    $sed -i"" '/^group:/ s/$/ cache_oslogin oslogin/' "$nss_config"
  fi

  if is_freebsd && grep -q '^passwd:.*compat' "$nss_config"; then
    $sed -i"" '/^passwd:/ s/compat/files/' "$nss_config"
  fi
}

restore_nsswitch_conf() {
  local nss_config="${1:-${nss_config}}"

  $sed -i"" '/^passwd:/ s/ cache_oslogin oslogin//' "$nss_config"
  $sed -i"" '/^group:/ s/ cache_oslogin oslogin//' "$nss_config"
  if is_freebsd; then
    $sed -i"" '/^passwd:/ s/files/compat/' "$nss_config"
  fi
}

modify_sshd_conf() (
  set -e

  local sshd_config="${1:-${sshd_config}}"

  local sshd_auth_keys_command="AuthorizedKeysCommand /usr/bin/google_authorized_keys"
  local sshd_auth_keys_command_user="AuthorizedKeysCommandUser root"
  local sshd_auth_methods="AuthenticationMethods publickey,keyboard-interactive"
  local sshd_challenge="ChallengeResponseAuthentication yes"

  # Update google_authorized_keys path in FreeBSD.
  if is_freebsd; then
    sshd_auth_keys_command="AuthorizedKeysCommand /usr/local/bin/google_authorized_keys"
  fi

  # Update directives for EL 6.
  if grep -qs "release 6" /etc/redhat-release; then
    sshd_auth_keys_command_user="AuthorizedKeysCommandRunAs root"
    sshd_auth_methods="RequiredAuthentications2 publickey,keyboard-interactive"
  fi

  add_or_update_sshd() {
    local entry="$1"
    local sshd_config="$2"
    local directive="$(echo "$entry" | cut -d' ' -f1)"
    local value="$(echo "$entry" | cut -d' ' -f2-)"

    # Check if directive is present.
    if grep -Eq "^\s*${directive}" "$sshd_config"; then
      # Check if value is incorrect.
      if ! grep -Eq "^\s*${directive}(\s|=)+${value}" "$sshd_config"; then
        # Comment out the line (because sshd_config is first-directive-found)
        # and add to end section.
        $sed -i"" -E "/^\s*${directive}/ s/^/${added_comment}\n#/" "$sshd_config"
        $sed -i"" "/$sshd_end_block/ i${entry}" "$sshd_config"
      fi
    else
      $sed -i"" "/$sshd_end_block/ i${entry}" "$sshd_config"
    fi
  }

  # Setup Google config block.
  if ! grep -q "$sshd_block" "$sshd_config"; then
    # Remove old-style additions.
    $sed -i"" "/${added_comment}/,+1d" "$sshd_config"
    printf "\n\n${sshd_block}\n${sshd_end_block}" >> "$sshd_config"
  fi

  for entry in "$sshd_auth_keys_command" "$sshd_auth_keys_command_user"; do
    add_or_update_sshd "$entry" "$sshd_config"
  done

  if [ -n "$two_factor" ]; then
    for entry in "$sshd_auth_methods" "$sshd_challenge"; do
      add_or_update_sshd "$entry" "$sshd_config"
    done
  fi
)

restore_sshd_conf() {
  local sshd_config="${1:-${sshd_config}}"

  if ! grep -q "$sshd_block" "$sshd_config"; then
    # Remove old-style additions.
    $sed -i"" "/${added_comment}/,+1d" "$sshd_config"
  else
    # Uncomment commented-out fields and remove Google config block.
    $sed -i"" "/${added_comment}/{n;s/^#//}" "$sshd_config"
    $sed -i"" "/${added_comment}/d" "$sshd_config"
    $sed -i"" "/${sshd_block}/,/${sshd_end_block}/d" "$sshd_config"
  fi
}

# Inserts pam modules to relevant pam stacks if missing.
modify_pam_config() (
  # TODO: idempotency of this function would be better assured if it wiped out
  # and applied desired changes each time rather than detecting deltas.

  set -e

  local pam_sshd_config="${1:-${pam_sshd_config}}"
  local pam_su_config="${1:-${pam_su_config}}"

  local pam_auth_oslogin="auth       [success=done perm_denied=die default=ignore] pam_oslogin_login.so"
  local pam_auth_group="auth       [default=ignore] pam_group.so"
  local pam_account_oslogin="account    [success=ok default=ignore] pam_oslogin_admin.so"
  local pam_account_admin="account    [success=ok ignore=ignore default=die] pam_oslogin_login.so"
  local pam_session_homedir="session    [success=ok default=ignore] pam_mkhomedir.so"
  local pam_account_su="account    [success=bad ignore=ignore] pam_oslogin_login.so"

  # In FreeBSD, the used flags are not supported, replacing them with the
  # previous ones (requisite and optional). This is not an exact feature parity
  # with Linux.
  if is_freebsd; then
    pam_auth_oslogin="auth       optional pam_oslogin_login.so"
    pam_auth_group="auth       optional pam_group.so"
    pam_account_oslogin="account    optional pam_oslogin_admin.so"
    pam_account_admin="account    requisite pam_oslogin_login.so"
    pam_session_homedir="session    optional pam_mkhomedir.so"
  fi

  local added_config=""
  local added_su_config=""

  # For COS this file is solely includes, so simply prepend the new config,
  # making each entry the top of its stack.
  if [ -e /etc/os-release ] && grep -q "ID=cos" /etc/os-release; then
    added_config="${added_comment}\n"
    for cfg in "$pam_account_admin" "$pam_account_oslogin" \
        "$pam_session_homedir" "$pam_auth_group"; do
      grep -qE "^${cfg%% *}.*${cfg##* }" ${pam_sshd_config} || added_config="${added_config}${cfg}\n"
    done

    if [ -n "$two_factor" ]; then
      grep -q "$pam_auth_oslogin" "$pam_sshd_config" || added_config="${added_config}${pam_auth_oslogin}\n"
    fi

    $sed -i"" "1i ${added_config}\n\n" "$pam_sshd_config"

    added_su_config="${added_comment}\n${pam_account_su}"
    $sed -i"" "1i ${added_su_config}" "$pam_su_config"

    return 0
  fi

  # Find the distro-specific insertion point for auth and su.
  if [ -e /etc/debian_version ]; then
    # Get location of common-auth and check if preceding line is a comment.
    insert=$($sed -rn "/^@include\s+common-auth/=" "$pam_sshd_config")
    $sed -n "$((insert-1))p" "$pam_sshd_config" | grep -q '^#' && insert=$((insert-1))
    su_insert=$($sed -rn "/^@include\s+common-account/=" "$pam_su_config")
  elif [ -e /etc/redhat-release ]; then
    # Get location of password-auth.
    insert=$($sed -rn "/^auth\s+(substack|include)\s+password-auth/=" \
      "$pam_sshd_config")
    # Get location of system-auth.
    su_insert=$($sed -rn "/^account\s+include\s+system-auth/=" "$pam_su_config")
  elif [ -e /etc/os-release ] && grep -q 'ID="sles"' /etc/os-release; then
    # Get location of common-auth.
    insert=$($sed -rn "/^auth\s+include\s+common-auth/=" "$pam_sshd_config")
    # Get location of common-account.
    su_insert=$($sed -rn "/^account\s+include\s+common-account/=" "$pam_su_config")
  elif [ -e /etc/arch-release ]; then
    # Get location of system-remote-login.
    insert=$($sed -rn "/^auth\s+include\s+system-remote-login/=" "$pam_sshd_config")
    # TODO: find su_insert point for arch linux.
  elif is_freebsd; then
    # Get location of the first auth occurrence
    insert=$($sed -rn '/^auth/=' "$pam_sshd_config" | head -1)
  fi

  added_config="$added_comment"
  if ! grep -qE '^auth.*pam_group' "$pam_sshd_config"; then
    added_config="${added_config}\n${pam_auth_group}"
  fi

  # This auth entry for OS Login+two factor MUST be added last, as it will
  # short-circuit processing of the auth stack via [success=ok]. auth stack
  # entries after this one will not be processed.
  if [ -n "$two_factor" ] && ! grep -qE '^auth.*oslogin' "$pam_sshd_config"; then
    added_config="${added_config}\n${pam_auth_oslogin}"
  fi

  # Insert auth modules at top of `sshd:auth` stack.
  if [ -n "$insert" ] && [ "$added_config" != "$added_comment" ]; then
    $sed -i"" "${insert}i ${added_config}" "$pam_sshd_config"
  fi

  # Insert su blocker at top of `su:account` stack.
  if [ -n "$su_insert" ] && ! grep -qE "$pam_account_su" "$pam_su_config"; then
    added_su_config="${added_comment}\n${pam_account_su}"
    $sed -i"" "${su_insert}i ${added_su_config}" "$pam_su_config"
  fi

  # Append account modules at end of `sshd:account` stack.
  if ! grep -qE '^account.*oslogin' "$pam_sshd_config"; then
    added_config="\\\n${added_comment}\n${pam_account_admin}\n${pam_account_oslogin}"
    account_end=$($sed -n '/^account/=' "$pam_sshd_config" | tail -1)
    $sed -i"" "${account_end}a ${added_config}" "$pam_sshd_config"
  fi

  # Append mkhomedir module at end of `sshd:session` stack.
  if ! grep -qE '^session.*mkhomedir' "$pam_sshd_config"; then
    added_config="\\\n${added_comment}\n${pam_session_homedir}"
    session_end=$($sed -n '/^session/=' "$pam_sshd_config" | tail -1)
    $sed -i"" "${session_end}a ${added_config}" "$pam_sshd_config"
  fi
)

restore_pam_config() {
  local pam_sshd_config="${1:-${pam_sshd_config}}"
  local pam_su_config="${1:-${pam_su_config}}"

  $sed -i"" "/${added_comment}/d" "$pam_sshd_config"
  $sed -i"" "/pam_oslogin/d" "$pam_sshd_config"
  $sed -i"" "/^session.*mkhomedir/d" "$pam_sshd_config"
  $sed -i"" "/^auth.*pam_group/d" "$pam_sshd_config"

  $sed -i"" "/${added_comment}/d" "$pam_su_config"
  $sed -i"" "/pam_oslogin/d" "$pam_su_config"
}

modify_group_conf() {
  # In FreeBSD there is no pam_group config file similar to
  # /etc/security/group.conf.
  if is_freebsd; then
    return
  fi

  local group_config="${1:-${group_config}}"
  local group_conf_entry="sshd;*;*;Al0000-2400;video"

  if ! grep -Fq "$group_conf_entry" "$group_config"; then
    $sed -i"" "\$a ${added_comment}\n${group_conf_entry}" "$group_config"
  fi
}

restore_group_conf() {
  # In FreeBSD there is no pam_group config file similar to
  # /etc/security/group.conf.
  if is_freebsd; then
    return
  fi

  local group_config="${1:-${group_config}}"

  $sed -i"" "/${added_comment}/{n;d}" "$group_config"
  $sed -i"" "/${added_comment}/d" "$group_config"
}

restart_service() {
  local service="$1"

  # The other options will be wrappers to systemctl on
  # systemd-enabled systems, so stop if found.
  if readlink -f /sbin/init|grep -q systemd; then
    if systemctl is-active --quiet "$service"; then
      systemctl restart "$service"
      return $?
    else
      return 0
    fi
  fi

  # Use the service helper if it exists.
  if command -v service > /dev/null; then
    if ! service "$service" status 2>&1 | grep -Eq "unrecognized|does not exist"; then
      service "$service" restart
      return $?
    else
      return 0
    fi
  fi

  # Fallback to trying sysvinit script of the same name.
  if command -v /etc/init.d/"$service" > /dev/null; then
    if /etc/init.d/"$service" status > /dev/null 2>&1; then
      /etc/init.d/"$service" restart
      return $?
    else
      return 0
    fi
  fi

  # We didn't find any way to restart this service.
  return 1
}

# Restart sshd unless --norestartsshd flag is set.
restart_sshd() {
  if [ -n "$no_restart_sshd" ]; then
    return 0
  fi
  echo "Restarting SSHD"
  for svc in "ssh" "sshd"; do
    restart_service "$svc"
  done
}

restart_svcs() {
  echo "Restarting optional services."
  for svc in "nscd" "unscd" "systemd-logind" "cron" "crond"; do
    restart_service "$svc"
  done
}

setup_google_dirs() {
  for dir in "$sudoers_dir" "$users_dir"; do
    [ -d "$dir" ] && continue
    mkdir -p "$dir"
    chmod 750 "$dir"
    if fixfiles=$(command -v fixfiles); then
      $fixfiles restore "$dir"
    fi
  done
  echo "#includedir ${sudoers_dir}" > "$sudoers_file"
  chmod 0440 "$sudoers_file"
}

remove_google_dirs() {
  for dir in "$sudoers_dir" "$users_dir"; do
    rm -rf "$dir"
  done
  rm -f "$sudoers_file"
}

activate() {
  for func in modify_sshd_conf modify_nsswitch_conf \
              modify_pam_config setup_google_dirs restart_svcs restart_sshd \
              modify_group_conf; do
    $func
    [ $? -eq 0 ] || return 1
  done
}

deactivate() {
  for func in remove_google_dirs restore_nsswitch_conf \
              restore_sshd_conf restore_pam_config restart_svcs restart_sshd \
              restore_group_conf; do
    $func
  done
}

# get_status checks each file for appropriate updates and exits on first
# failure. Checks for two factor config changes only if requested.
get_status() (
  set -e

  grep -Eq '^account.*oslogin' "$pam_sshd_config"
  grep -Eq 'google_authorized_keys' "$sshd_config"
  grep -Eq 'passwd:.*oslogin' "$nss_config"
  if [ -n "$two_factor" ]; then
    grep -Eq '^auth.*oslogin' "$pam_sshd_config"
    grep -Eq '^(AuthenticationMethods|RequiredAuthentications2).*publickey,keyboard-interactive' "$sshd_config"
  fi
)

usage() {
  echo "Usage: $(basename "$0") {activate|deactivate|status} [--norestartsshd] [--twofactor]"
  echo "This script will activate or deactivate the features for"
  echo "Google Compute Engine OS Login and (optionally) two-factor authentication."
  echo "This script must be run as root."
  exit 1
}


# Main
if [ $(id -u) -ne 0 ] || [ $# -lt 1 ]; then
  usage
fi

sed="sed"
is_freebsd && sed="gsed"

while [ $# -gt 0 ]; do
  case "$1" in
    --norestartsshd)
      no_restart_sshd="true"
      shift
      ;;
    --twofactor)
      two_factor="true"
      shift
      ;;
    activate)
      action="activate"
      shift
      ;;
    deactivate)
      action="deactivate"
      shift
      ;;
    status)
      action="status"
      shift
      ;;
    *)
      shift
      ;;
  esac
done

case "$action" in
  activate)
    echo "Activating Google Compute Engine OS Login."
    activate
    if [ $? -ne 0 ]; then
      echo "Failed to apply changes, rolling back"
      deactivate
      exit 1
    fi
    ;;
  deactivate)
    echo "Deactivating Google Compute Engine OS Login."
    deactivate
    ;;
  status)
    get_status
    exit $?
    ;;
  *)
    usage
    ;;
esac