#!/bin/bash

# Copyright 2014-2022 Security Onion Solutions, LLC

#    This program 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.
#
#    This program 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 this program.  If not, see <http://www.gnu.org/licenses/>.


source $(dirname $0)/so-common

DEFAULT_ROLE=analyst

if [[ $# -lt 1 || $# -gt 3 ]]; then
  echo "Usage: $0 <operation> [email] [role]"
  echo ""
  echo " where <operation> is one of the following:"
  echo ""
  echo "     list: Lists all user email addresses currently defined in the identity system"
  echo "      add: Adds a new user to the identity system; requires 'email' parameter, while 'role' parameter is optional and defaults to $DEFAULT_ROLE"
  echo "  addrole: Grants a role to an existing user; requires 'email' and 'role' parameters"
  echo "  delrole: Removes a role from an existing user; requires 'email' and 'role' parameters"
  echo "   update: Updates a user's password and disables MFA; requires 'email' parameter"
  echo "   enable: Enables a user; requires 'email' parameter"
  echo "  disable: Disables a user; requires 'email' parameter"
  echo " validate: Validates that the given email address and password are acceptable; requires 'email' parameter"
  echo " valemail: Validates that the given email address is acceptable; requires 'email' parameter"
  echo "  valpass: Validates that a password is acceptable"
  echo ""
  echo " Note that the password can be piped into STDIN to avoid prompting for it"
  exit 1
fi

operation=$1
email=$2
role=$3

kratosUrl=${KRATOS_URL:-http://127.0.0.1:4434}
databasePath=${KRATOS_DB_PATH:-/opt/so/conf/kratos/db/db.sqlite}
databaseTimeout=${KRATOS_DB_TIMEOUT:-5000}
bcryptRounds=${BCRYPT_ROUNDS:-12}
elasticUsersFile=${ELASTIC_USERS_FILE:-/opt/so/saltstack/local/salt/elasticsearch/files/users}
elasticRolesFile=${ELASTIC_ROLES_FILE:-/opt/so/saltstack/local/salt/elasticsearch/files/users_roles}
socRolesFile=${SOC_ROLES_FILE:-/opt/so/conf/soc/soc_users_roles}
esUID=${ELASTIC_UID:-930}
esGID=${ELASTIC_GID:-930}
soUID=${SOCORE_UID:-939}
soGID=${SOCORE_GID:-939}

function lock() {
  # Obtain file descriptor lock
  exec 99>/var/tmp/so-user.lock || fail "Unable to create lock descriptor; if the system was not shutdown gracefully you may need to remove /var/tmp/so-user.lock manually."
  flock -w 10 99 || fail "Another process is using so-user; if the system was not shutdown gracefully you may need to remove /var/tmp/so-user.lock manually."
  trap 'rm -f /var/tmp/so-user.lock' EXIT
}

function fail() {
  msg=$1
  echo "$1"
  exit 1
}

function require() {
  cmd=$1
  which "$1" 2>&1 > /dev/null
  [[ $? != 0 ]] && fail "This script requires the following command be installed: ${cmd}"
}

# Verify this environment is capable of running this script
function verifyEnvironment() {
  require "htpasswd"
  require "jq"
  require "curl"
  require "openssl"
  require "sqlite3"
  [[ ! -f $databasePath ]] && fail "Unable to find database file; specify path via KRATOS_DB_PATH environment variable"
  response=$(curl -Ss -L ${kratosUrl}/)
  [[ "$response" != "404 page not found" ]] && fail "Unable to communicate with Kratos; specify URL via KRATOS_URL environment variable"
}

function findIdByEmail() {
  email=$1

  response=$(curl -Ss -L ${kratosUrl}/identities)
  identityId=$(echo "${response}" | jq -r ".[] | select(.verifiable_addresses[0].value == \"$email\") | .id")
  echo $identityId
}

function validatePassword() {
  password=$1

  len=$(expr length "$password")
  if [[ $len -lt 8 ]]; then
    fail "Password does not meet the minimum requirements"
  fi
  if [[ $len -gt 72 ]]; then
    fail "Password is too long (max: 72)"
  fi
  check_password_and_exit "$password"
}

function validateEmail() {
  email=$1
  # (?:[a-z0-9!#$%&'*+/=?^_`{|}~-]+(?:\.[a-z0-9!#$%&'*+/=?^_`{|}~-]+)*|"(?:[\x01-\x08\x0b\x0c\x0e-\x1f\x21\x23-\x5b\x5d-\x7f]|\\[\x01-\x09\x0b\x0c\x0e-\x7f])*")@(?:(?:[a-z0-9](?:[a-z0-9-]*[a-z0-9])?\.)+[a-z0-9](?:[a-z0-9-]*[a-z0-9])?|\[(?:(?:(2(5[0-5]|[0-4][0-9])|1[0-9][0-9]|[1-9]?[0-9]))\.){3}(?:(2(5[0-5]|[0-4][0-9])|1[0-9][0-9]|[1-9]?[0-9])|[a-z0-9-]*[a-z0-9]:(?:[\x01-\x08\x0b\x0c\x0e-\x1f\x21-\x5a\x53-\x7f]|\\[\x01-\x09\x0b\x0c\x0e-\x7f])+)\])
  if [[ ! "$email" =~ ^[[:alnum:]._%+-]+@[[:alnum:].-]+\.[[:alpha:]]{2,}$ ]]; then
    fail "Email address is invalid"
  fi

  if [[ "$email" =~ [A-Z] ]]; then
    fail "Email addresses cannot contain uppercase letters"
  fi
}

function hashPassword() {
  password=$1

  passwordHash=$(echo "${password}" | htpasswd -niBC $bcryptRounds SOUSER)
  passwordHash=$(echo "$passwordHash" | cut -c 11-)
  passwordHash="\$2a${passwordHash}" # still waiting for https://github.com/elastic/elasticsearch/issues/51132
  echo "$passwordHash"
}


function updatePassword() {
  identityId=$1
  
  if [ -z "$password" ]; then
    # Read password from stdin (show prompt only if no stdin was piped in)
    test -t 0
    if [[ $? == 0 ]]; then
      echo "Enter new password:"
    fi
    read -rs password

    validatePassword "$password"
  fi

  if [[ -n "$identityId" ]]; then
    # Generate password hash
    passwordHash=$(hashPassword "$password")
    # Update DB with new hash
    echo "update identity_credentials set config=CAST('{\"hashed_password\":\"$passwordHash\"}' as BLOB), updated_at=datetime('now') where identity_id='${identityId}' and identity_credential_type_id=(select id from identity_credential_types where name='password');" | sqlite3 -cmd ".timeout ${databaseTimeout}" "$databasePath"
    # Deactivate MFA
    echo "delete from identity_credential_identifiers where identity_credential_id=(select id from identity_credentials where identity_id='${identityId}' and identity_credential_type_id=(select id from identity_credential_types where name='totp'));" | sqlite3 -cmd ".timeout ${databaseTimeout}" "$databasePath"
    echo "delete from identity_credentials where identity_id='${identityId}' and identity_credential_type_id=(select id from identity_credential_types where name='totp');" | sqlite3 -cmd ".timeout ${databaseTimeout}" "$databasePath"
    [[ $? != 0 ]] && fail "Unable to update password"
  fi
}

function createFile() {
  filename=$1
  uid=$2
  gid=$3

  mkdir -p $(dirname "$filename")
  truncate -s 0 "$filename"
  chmod 600 "$filename"
  chown "${uid}:${gid}" "$filename"
}

function ensureRoleFileExists() {
  if [[ ! -f "$socRolesFile" || ! -s "$socRolesFile" ]]; then
    # Generate the new users file
    rolesTmpFile="${socRolesFile}.tmp"
    createFile "$rolesTmpFile" "$soUID" "$soGID"

    if [[ -f "$databasePath" ]]; then
      echo "Migrating roles to new file: $socRolesFile"

      echo "select 'superuser:' || id from identities;" | sqlite3 -cmd ".timeout ${databaseTimeout}" "$databasePath" \
        >> "$rolesTmpFile"
      [[ $? != 0 ]] && fail "Unable to read identities from database"

      echo "The following users have all been migrated with the super user role:"
      cat "${rolesTmpFile}"
    else
      echo "Database file does not exist yet, installation is likely not yet complete."
    fi

    if [[ -d "$socRolesFile" ]]; then
      echo "Removing invalid roles directory created by Docker"
      rm -fr "$socRolesFile"
    fi
    mv "${rolesTmpFile}" "${socRolesFile}"
  fi
}

function syncElasticSystemUser() {
  json=$1
  userid=$2
  usersFile=$3

  user=$(echo "$json" | jq -r ".local.users.$userid.user")
  pass=$(echo "$json" | jq -r ".local.users.$userid.pass")
  
  [[ -z "$user" || -z "$pass" ]] && fail "Elastic auth credentials for system user '$userid' are missing"
  hash=$(hashPassword "$pass")

  echo "${user}:${hash}" >> "$usersFile"
}

function syncElasticSystemRole() {
  json=$1
  userid=$2
  role=$3
  rolesFile=$4

  user=$(echo "$json" | jq -r ".local.users.$userid.user")
  
  [[ -z "$user" ]] && fail "Elastic auth credentials for system user '$userid' are missing"

  echo "${role}:${user}" >> "$rolesFile"
}

function syncElastic() {
  echo "Syncing users and roles between SOC and Elastic..."

  usersTmpFile="${elasticUsersFile}.tmp"
  createFile "${usersTmpFile}" "$esUID" "$esGID"
  rolesTmpFile="${elasticRolesFile}.tmp"
  createFile "${rolesTmpFile}" "$esUID" "$esGID"

  authPillarJson=$(lookup_salt_value "auth" "elasticsearch" "pillar" "json")

  syncElasticSystemUser "$authPillarJson" "so_elastic_user"  "$usersTmpFile"
  syncElasticSystemUser "$authPillarJson" "so_kibana_user"   "$usersTmpFile"
  syncElasticSystemUser "$authPillarJson" "so_logstash_user" "$usersTmpFile"
  syncElasticSystemUser "$authPillarJson" "so_beats_user"    "$usersTmpFile"
  syncElasticSystemUser "$authPillarJson" "so_monitor_user"  "$usersTmpFile"

  syncElasticSystemRole "$authPillarJson" "so_elastic_user"  "superuser" "$rolesTmpFile"
  syncElasticSystemRole "$authPillarJson" "so_kibana_user"   "superuser" "$rolesTmpFile"
  syncElasticSystemRole "$authPillarJson" "so_logstash_user" "superuser" "$rolesTmpFile"
  syncElasticSystemRole "$authPillarJson" "so_beats_user"    "superuser" "$rolesTmpFile"
  syncElasticSystemRole "$authPillarJson" "so_monitor_user"  "remote_monitoring_collector" "$rolesTmpFile"
  syncElasticSystemRole "$authPillarJson" "so_monitor_user"  "remote_monitoring_agent" "$rolesTmpFile"
  syncElasticSystemRole "$authPillarJson" "so_monitor_user"  "monitoring_user" "$rolesTmpFile"

  if [[ -f "$databasePath" && -f "$socRolesFile" ]]; then
    # Append the SOC users 
    userData=$(echo "select '{\"user\":\"' || ici.identifier || '\", \"data\":' || ic.config || '}'" \
      "from identity_credential_identifiers ici, identity_credentials ic, identities i, identity_credential_types ict " \
      "where " \
      "      ici.identity_credential_id=ic.id " \
      "  and ic.identity_id=i.id " \
      "  and ict.id=ic.identity_credential_type_id " \
      "  and ict.name='password' " \
      "  and instr(ic.config, 'hashed_password') " \
      "  and i.state == 'active' " \
      "order by ici.identifier;" | \
      sqlite3 -cmd ".timeout ${databaseTimeout}" "$databasePath")
    [[ $? != 0 ]] && fail "Unable to read credential hashes from database"
    echo "${userData}" | \
      jq -r '.user + ":" + .data.hashed_password' \
      >> "$usersTmpFile"

    # Append the user roles
    while IFS="" read -r rolePair || [ -n "$rolePair" ]; do
      userId=$(echo "$rolePair" | cut -d: -f2)
      role=$(echo "$rolePair" | cut -d: -f1)
      echo "select '$role:' || ici.identifier " \
        "from identity_credential_identifiers ici, identity_credentials ic, identity_credential_types ict " \
        "where ici.identity_credential_id=ic.id " \
        "  and ict.id=ic.identity_credential_type_id " \
        "  and ict.name='password' " \
        "  and ic.identity_id = '$userId';" | \
        sqlite3 -cmd ".timeout ${databaseTimeout}" "$databasePath" >> "$rolesTmpFile"
      [[ $? != 0 ]] && fail "Unable to read role identities from database"
    done < "$socRolesFile"

  else
    echo "Database file or soc roles file does not exist yet, skipping users export"
  fi

  if [[ -s "${usersTmpFile}" ]]; then
    mv "${usersTmpFile}" "${elasticUsersFile}"
    mv "${rolesTmpFile}" "${elasticRolesFile}"

    if [[ -z "$SKIP_STATE_APPLY" ]]; then
      echo "Elastic state will be re-applied to affected minions. This may take several minutes..."
      echo "Applying elastic state to elastic minions at $(date)" >> /opt/so/log/soc/sync.log 2>&1
      salt -C 'G@role:so-standalone or G@role:so-eval or G@role:so-import or G@role:so-manager or G@role:so-managersearch or G@role:so-node or G@role:so-heavynode' state.apply elasticsearch queue=True >> /opt/so/log/soc/sync.log 2>&1
    fi
  else
    echo "Newly generated users/roles files are incomplete; aborting."
  fi
}

function syncAll() {
  ensureRoleFileExists

  # Check if a sync is needed. Sync is not needed if the following are true:
  # - user database entries are all older than the elastic users file
  # - soc roles file last modify date is older than the elastic roles file
  if [[ -z "$FORCE_SYNC" && -f "$databasePath" && -f "$elasticUsersFile" ]]; then
    usersFileAgeSecs=$(echo $(($(date +%s) - $(date +%s -r "$elasticUsersFile"))))
    staleCount=$(echo "select count(*) from identity_credentials where updated_at >= Datetime('now', '-${usersFileAgeSecs} seconds');" \
      | sqlite3 -cmd ".timeout ${databaseTimeout}" "$databasePath")
    [[ $? != 0 ]] && fail "Unable to read user count from database"
    if [[ "$staleCount" == "0" && "$elasticRolesFile" -nt "$socRolesFile" ]]; then
      return 1
    fi
  fi

  syncElastic

  return 0
}

function listUsers() {
  response=$(curl -Ss -L ${kratosUrl}/identities)
  [[ $? != 0 ]] && fail "Unable to communicate with Kratos"

  users=$(echo "${response}" | jq -r ".[] | .verifiable_addresses[0].value" | sort)
  for user in $users; do
    roles=$(grep ":$user\$" "$elasticRolesFile" | cut -d: -f1 | tr '\n' ' ')
    echo "$user: $roles"
  done
}

function addUserRole() {
  email=$1
  role=$2

  adjustUserRole "$email" "$role" "add"
}

function deleteUserRole() {
  email=$1
  role=$2

  adjustUserRole "$email" "$role" "del"
}

function adjustUserRole() {
  email=$1
  role=$2
  op=$3

  identityId=$(findIdByEmail "$email")
  [[ ${identityId} == "" ]] && fail "User not found"

  ensureRoleFileExists

  filename="$socRolesFile"
  hasRole=0
  grep "^$role:" "$socRolesFile" | grep -q "$identityId" && hasRole=1
  if [[ "$op" == "add" ]]; then
    if [[ "$hasRole" == "1" ]]; then
      echo "User '$email' already has the role: $role"
      return 1
    else
      echo "$role:$identityId" >> "$filename"
    fi
  elif [[ "$op" == "del" ]]; then
    if [[ "$hasRole" -ne 1 ]]; then
      fail "User '$email' does not have the role: $role"
    else
      sed "/^$role:$identityId\$/d" "$filename" > "$filename.tmp"
      cat "$filename".tmp > "$filename"
      rm -f "$filename".tmp
    fi
  else
    fail "Unsupported role adjustment operation: $op"
  fi
  return 0
}

function createUser() {
  email=$1
  role=$2

  now=$(date -u +%FT%TZ)
  addUserJson=$(cat <<EOF
{
  "traits": {"email":"${email}"},
  "schema_id": "default"
}
EOF
  )
  
  response=$(curl -Ss -L ${kratosUrl}/identities -d "$addUserJson")
  [[ $? != 0 ]] && fail "Unable to communicate with Kratos"

  identityId=$(echo "${response}" | jq -r ".id")
  if [[ "${identityId}" == "null" ]]; then
    code=$(echo "${response}" | jq ".error.code")
    [[ "${code}" == "409" ]] && fail "User already exists"

    reason=$(echo "${response}" | jq ".error.message")
    [[ $? == 0 ]] && fail "Unable to add user: ${reason}"
  else
    updatePassword "$identityId"
    addUserRole "$email" "$role"
  fi
}

function migrateLockedUsers() {
  # This is a migration function to convert locked users from prior to 2.3.90
  # to inactive users using the newer Kratos functionality. This should only
  # find locked users once.
  lockedEmails=$(curl -s http://localhost:4434/identities | jq -r '.[] | select(.traits.status == "locked") | .traits.email')
  if [[ -n "$lockedEmails" ]]; then
    echo "Disabling locked users..."
    for email in $lockedEmails; do
      updateStatus "$email" locked
    done
  fi
}

function updateStatus() {
  email=$1
  status=$2

  identityId=$(findIdByEmail "$email")
  [[ ${identityId} == "" ]] && fail "User not found"

  response=$(curl -Ss -L "${kratosUrl}/identities/$identityId")
  [[ $? != 0 ]] && fail "Unable to communicate with Kratos"

  schemaId=$(echo "$response" | jq -r .schema_id)

  # Capture traits and remove obsolete 'status' trait if exists
  traitBlock=$(echo "$response" | jq -c .traits | sed -re 's/,?"status":".*?"//')

  state="active"
  if [[ "$status" == "locked" ]]; then
    state="inactive"
  fi
  body="{ \"schema_id\": \"$schemaId\", \"state\": \"$state\", \"traits\": $traitBlock }"
  response=$(curl -fSsL -XPUT "${kratosUrl}/identities/$identityId" -d "$body")
  [[ $? != 0 ]] && fail "Unable to update user"
}

function updateUser() {
  email=$1

  identityId=$(findIdByEmail "$email")
  [[ ${identityId} == "" ]] && fail "User not found"

  updatePassword "$identityId"
}

function deleteUser() {
  email=$1

  identityId=$(findIdByEmail "$email")
  [[ ${identityId} == "" ]] && fail "User not found"

  response=$(curl -Ss -XDELETE -L "${kratosUrl}/identities/$identityId")
  [[ $? != 0 ]] && fail "Unable to communicate with Kratos"

  rolesTmpFile="${socRolesFile}.tmp"
  createFile "$rolesTmpFile" "$soUID" "$soGID"
  grep -v "$identityId" "$socRolesFile" > "$rolesTmpFile"
  mv "$rolesTmpFile" "$socRolesFile"
}

case "${operation}" in
  "add")
    verifyEnvironment
    [[ "$email" == "" ]] && fail "Email address must be provided"

    lock
    validateEmail "$email"
    updatePassword
    createUser "$email" "${role:-$DEFAULT_ROLE}"
    syncAll
    echo "Successfully added new user to SOC"
    check_container fleet && echo "$password" | so-fleet-user-add "$email"
    ;;

  "list")
    verifyEnvironment
    listUsers
    ;;

  "addrole")
    verifyEnvironment
    [[ "$email" == "" ]] && fail "Email address must be provided"
    [[ "$role" == "" ]] && fail "Role must be provided"

    lock
    validateEmail "$email"
    if addUserRole "$email" "$role"; then
      syncElastic
      echo "Successfully added role to user"
    fi
    ;;

  "delrole")
    verifyEnvironment
    [[ "$email" == "" ]] && fail "Email address must be provided"
    [[ "$role" == "" ]] && fail "Role must be provided"

    lock
    validateEmail "$email"
    deleteUserRole "$email" "$role"
    syncElastic
    echo "Successfully removed role from user"
    ;;

  "update")
    verifyEnvironment
    [[ "$email" == "" ]] && fail "Email address must be provided"

    lock
    updateUser "$email"
    syncAll
    echo "Successfully updated user"
    ;;

  "enable")
    verifyEnvironment
    [[ "$email" == "" ]] && fail "Email address must be provided"

    lock
    updateStatus "$email" 'active'
    syncAll
    echo "Successfully enabled user"
    echo "Fleet user will need to be recreated manually with so-fleet-user-add"
    ;;

  "disable")
    verifyEnvironment
    [[ "$email" == "" ]] && fail "Email address must be provided"

    lock
    updateStatus "$email" 'locked'
    syncAll
    echo "Successfully disabled user"
    check_container fleet && so-fleet-user-delete "$email"   
    ;;    

  "delete")
    verifyEnvironment
    [[ "$email" == "" ]] && fail "Email address must be provided"

    lock
    deleteUser "$email"
    syncAll
    echo "Successfully deleted user"
    check_container fleet && so-fleet-user-delete "$email"
    ;;

  "sync")
    lock
    syncAll
    ;;

  "validate")
    validateEmail "$email"
    updatePassword
    echo "Email and password are acceptable"
    ;;

  "valemail")
    validateEmail "$email"
    echo "Email is acceptable"
    ;;

  "valpass")
    updatePassword
    echo "Password is acceptable"
    ;;

  "migrate")
    migrateLockedUsers
    echo "User migration complete"
    ;;

  *)
    fail "Unsupported operation: $operation"
    ;;
esac

exit 0