#!/bin/bash

# Copyright 2014,2015,2016,2017,2018,2019,2020,2021 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

if [[ $# -lt 1 || $# -gt 2 ]]; then
  echo "Usage: $0 <list|add|update|enable|disable|validate|valemail|valpass> [email]"
  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"
  echo "   update: Updates a user's password; 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

kratosUrl=${KRATOS_URL:-http://127.0.0.1:4434}
databasePath=${KRATOS_DB_PATH:-/opt/so/conf/kratos/db/db.sqlite}
bcryptRounds=${BCRYPT_ROUNDS:-12}
elasticUsersFile=${ELASTIC_USERS_FILE:-/opt/so/conf/elasticsearch/users}
elasticRolesFile=${ELASTIC_ROLES_FILE:-/opt/so/conf/elasticsearch/users_roles}

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 ".[] | select(.verifiable_addresses[0].value == \"$email\") | .id")
  echo $identityId
}

function validatePassword() {
  password=$1

  len=$(expr length "$password")
  if [[ $len -lt 6 ]]; then
    echo "Password does not meet the minimum requirements"
    exit 2
  fi
}

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
    echo "Email address is invalid"
    exit 3
  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) where identity_id=${identityId};" | sqlite3 "$databasePath"
    [[ $? != 0 ]] && fail "Unable to update password"
  fi
}

function createElasticTmpFile() {
  filename=$1
  tmpFile=${filename}.tmp
  truncate -s 0 "$tmpFile"
  chmod 600 "$tmpFile"
  chown elasticsearch:elasticsearch "$tmpFile"
  echo "$tmpFile"
}

function syncElastic() {
  usersFileTmp=$(createElasticTmpFile "${elasticUsersFile}")
  rolesFileTmp=$(createElasticTmpFile "${elasticRolesFile}")

  sysUser=$(lookup_pillar "auth:user" "elasticsearch")
  sysPass=$(lookup_pillar "auth:pass" "elasticsearch")
  sysHash=$(hashPassword "$sysPass")

  # Generate the new users file
  echo "${sysUser}:${sysHash}" >> "$usersFileTmp"
  echo "select '{\"user\":\"' || ici.identifier || '\", \"data\":' || ic.config || '}'" \
    "from identity_credential_identifiers ici, identity_credentials ic " \
    "where ici.identity_credential_id=ic.id and ic.config like '%hashed_password%' " \
    "order by ici.identifier;" | \
    sqlite3 "$databasePath" | \
    jq -r '.user + ":" + .data.hashed_password' \
    >> "$usersFileTmp"
  mv -f "$usersFileTmp" "$elasticUsersFile"

  # Generate the new users_roles file
  echo "superuser:${sysUser}" >> "$rolesFileTmp"
  echo "select 'superuser:' || ici.identifier " \
    "from identity_credential_identifiers ici, identity_credentials ic " \
    "where ici.identity_credential_id=ic.id and ic.config like '%hashed_password%' " \
    "order by ici.identifier;" | \
    sqlite3 "$databasePath" \
    >> "$rolesFileTmp"
  mv -f "$rolesFileTmp" "$elasticRolesFile"
}

function syncAll() {
  syncElastic
}

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

  echo "${response}" | jq -r ".[] | .verifiable_addresses[0].value" | sort
}

function createUser() {
  email=$1

  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 ".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}"
  fi

  updatePassword $identityId
}

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"

  oldConfig=$(echo "select config from identity_credentials where identity_id=${identityId};" | sqlite3 "$databasePath")
  if [[ "$status" == "locked" ]]; then
    config=$(echo $oldConfig | sed -e 's/hashed/locked/')
    echo "update identity_credentials set config=CAST('${config}' as BLOB) where identity_id=${identityId};" | sqlite3 "$databasePath"
    [[ $? != 0 ]] && fail "Unable to lock credential record"

    echo "delete from sessions where identity_id=${identityId};" | sqlite3 "$databasePath"
    [[ $? != 0 ]] && fail "Unable to invalidate sessions"    
  else
    config=$(echo $oldConfig | sed -e 's/locked/hashed/')
    echo "update identity_credentials set config=CAST('${config}' as BLOB) where identity_id=${identityId};" | sqlite3 "$databasePath"
    [[ $? != 0 ]] && fail "Unable to unlock credential record"
  fi  

  updatedJson=$(echo "$response" | jq ".traits.status = \"$status\" | del(.verifiable_addresses) | del(.id) | del(.schema_url)")
  response=$(curl -Ss -XPUT -L ${kratosUrl}/identities/$identityId -d "$updatedJson")
  [[ $? != 0 ]] && fail "Unable to mark user as locked"

}

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"
}

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

    validateEmail "$email"
    updatePassword
    createUser "$email"
    syncAll
    echo "Successfully added new user to SOC"
    check_container thehive && echo $password | so-thehive-user-add "$email"
    check_container fleet && echo $password | so-fleet-user-add "$email"
    ;;

  "list")
    verifyEnvironment
    listUsers
    ;;

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

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

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

    updateStatus "$email" 'active'
    syncAll
    echo "Successfully enabled user"
    check_container thehive && so-thehive-user-enable "$email" true
    check_container fleet && so-fleet-user-enable "$email" true   
    ;;

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

    updateStatus "$email" 'locked'
    syncAll
    echo "Successfully disabled user"
    check_container thehive && so-thehive-user-enable "$email" false
    check_container fleet && so-fleet-user-enable "$email" false   
    ;;    

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

    deleteUser "$email"
    syncAll
    echo "Successfully deleted user"
    check_container thehive && so-thehive-user-enable "$email" false
    check_container fleet && so-fleet-user-enable "$email" false
    ;;

  "sync")
    syncAll
    echo "Synchronization complete"
    ;;

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

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

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

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

exit 0