#!/bin/bash

# Copyright Security Onion Solutions LLC and/or licensed to Security Onion Solutions LLC under one
# or more contributor license agreements. Licensed under the Elastic License 2.0 as shown at 
# https://securityonion.net/license; you may not use this file except in compliance with the
# Elastic License 2.0.

. /usr/sbin/so-common

usage() {
    echo "Usage: $0 <operation> [args]"
    echo ""
    echo "Supported Operations:"
    echo "  setup          Loads all templates and creates all required buckets"
    echo "  templateapply  Applies a single template file, requires: <template-path>"
    echo "  userlist       Lists users"
    echo "  useradd        Adds a new user, requires: <email>"
    echo "  userdel        Removes an existing user, requires: <email>"
    echo "  userenable     Enables a user, requires: <email>"
    echo "  userdisable    Disables a user, requires: <email>"
    echo "  userpass       Updates a user's password, requires: <email>"
    echo ""
    echo "If required, the password will be read from STDIN."
    exit 1
}

if [ $# -lt 1 ]; then
  usage
fi

OP=$1
shift

set -eo pipefail

log() {
  echo -e "$(date) | InfluxDB | $@" >&2
}

check_response() {
  response=$1
  if [[ "$response" =~ "\"code\":" ]]; then
    log "Failed. Check the response for more details.\n$response"
    exit 1
  fi
}

lookup_user_id() {
  token=$1
  email=$2
  response=$(curl -sk https://localhost:8086/api/v2/users?limit=100 -H "Authorization: Token $token")
  check_response "$response"
  uid=$(echo "$response" | jq -r ".users[] | select(.name == \"$email\").id")
  if [[ -z "$uid" ]]; then
    log "User not found"
    exit 1
  fi
  echo "$uid"
}

lookup_org_id() {
  token=$1
  response=$(curl -sk https://localhost:8086/api/v2/orgs?limit=100 -H "Authorization: Token $token")
  check_response "$response"
  oid=$(echo "$response" | jq -r ".orgs[] | select(.name == \"Security Onion\").id")
  if [[ -z "$oid" ]]; then
    log "Organization not found"
    exit 1
  fi
  echo "$oid"
}

lookup_stack_id() {
  token=$1
  oid=$2
  response=$(curl -sk "https://localhost:8086/api/v2/stacks?orgID=$oid&name=Security+Onion" -H "Authorization: Token $token")
  check_response "$response"
  stackid=$(echo "$response" | jq -r ".stacks[0].id")
  if [[ -z "$stackid" || "$stackid" == null ]]; then
    response=$(curl -sk https://localhost:8086/api/v2/stacks -X POST -d "{\"name\":\"Security Onion\",\"orgID\":\"$oid\"}" -H "Authorization: Token $token")
    check_response "$response"
    stackid=$(echo "$response" | jq -r .id)
  fi
  echo "$stackid"
}

add_user_to_org() {
  token=$1
  uid=$2
  oid=$3
  log "Adding new user to organization"
  response=$(curl -sk https://localhost:8086/api/v2/orgs/$oid/members -X POST -d "{\"id\":\"$uid\"}" -H "Authorization: Token $token")
  check_response "$response"
}

change_password() {
  token=$1
  uid=$2
  set +e
  test -t 0
  if [[ $? == 0 ]]; then
    echo "Enter new password:"
  fi
  set -e
  read -rs pass
  check_password_and_exit "$pass"
  response=$(curl -sk https://localhost:8086/api/v2/users/$uid/password -X POST -d "{\"password\":\"$pass\"}" -H "Authorization: Token $token")
  check_response "$response"
}

apply_template() {
  token=$1
  oid=$2
  stackid=$3
  file=$4
  content=$(cat $file)
  body="{\"orgID\":\"$oid\",\"stackID\":\"$stackid\",\"template\":{\"contents\":$content}}"
  response=$(curl -sk https://localhost:8086/api/v2/templates/apply -X POST -d "$body" -H "Authorization: Token $token")
  check_response "$response"
}

create_bucket() {
  token=$1
  oid=$2
  name=$3
  age=$4
  shardduration=$5

  response=$(curl -sk "https://localhost:8086/api/v2/buckets?orgID=$oid&name=$name" -H "Authorization: Token $token")
  check_response "$response"
  bucketid=$(echo "$response" | jq -r ".buckets[0].id")
  if [[ -z "$stackid" || "$stackid" == null ]]; then
    response=$(curl -sk https://localhost:8086/api/v2/buckets -X POST -d "{\"name\":\"$name\",\"orgID\":\"oid\"}" -H "Authorization: Token $token")
    check_response "$response"
    bucketid=$(echo "$response" | jq -r .id)
  fi
  response=$(curl -sk "https://localhost:8086/api/v2/buckets/$bucketid" -d "{\"name\":\"$name\",\"retentionRules\":[{\"everySeconds\":$age,\"shardGroupDurationSeconds\":$shardduration,\"type\":\"expire\"}]}" -H "Authorization: Token $token")
  check_response "$response"
}

case "$OP" in

  templateload)
    [ $# -ne 1 ] && usage
    file=$1
    log "Applying template file; file=$file"
    token=$(lookup_pillar_secret influx_token)
    oid=$(lookup_org_id "$token")
    stackid=$(lookup_stack_id "$token" "$oid")
    apply_template "$token" "$oid" "$stackid" "$file"
    ;;

  setup)
    log "Ensuring organization is setup correctly"
    token=$(lookup_pillar_secret influx_token)
    oid=$(lookup_org_id "$token")

    # Load templates
    stackid=$(lookup_stack_id "$token" "$oid")
    for file in /opt/so/conf/influxdb/templates/*; do
      log "Ensuring template is loaded; template=$file"
      apply_template "$token" "$oid" "$stackid" "$file"
    done

    # Setup buckets and retention periods
    for rp in so_short_term so_long_term; do
      bucket=telegraf/$rp
      log "Ensuring bucket is created and configured; bucket=$bucket"
      age=$(cat /opt/so/conf/influxdb/buckets.json | jq -r .$rp.duration)
      shard_duration=$(cat /opt/so/conf/influxdb/buckets.json | jq -r .$rp.shard_duration)
      create_bucket "$token" "$oid" "$bucket" "$age" "$shard_duration"
    done
    ;;    

  userlist)
    log "Listing existing users"
    token=$(lookup_pillar_secret influx_token)
    response=$(curl -sk https://localhost:8086/api/v2/users -H "Authorization: Token $token")
    check_response "$response"
    echo "$response" | jq -r '.users[] | "\(.id): \(.name) (\(.status))"'
    ;;

  useradd)
    [ $# -ne 1 ] && usage
    email=$1
    log "Adding new user; email=$email"
    token=$(lookup_pillar_secret influx_token)
    oid=$(lookup_org_id "$token")
    response=$(curl -sk https://localhost:8086/api/v2/users -X POST -d "{\"name\":\"$email\"}" -H "Authorization: Token $token")
    check_response "$response"
    uid=$(echo "$response" | jq -r .id)
    add_user_to_org "$token" "$uid" "$oid"
    change_password "$token" "$uid"
    ;;

  userpass)
    [ $# -ne 1 ] && usage
    email=$1
    log "Updating user password; email=$email"
    token=$(lookup_pillar_secret influx_token)
    uid=$(lookup_user_id "$token" "$email")
    change_password "$token" "$uid"
    ;;
  
  userdel)
    [ $# -ne 1 ] && usage
    email=$1
    log "Deleting user; email=$email"
    token=$(lookup_pillar_secret influx_token)
    uid=$(lookup_user_id "$token" "$email")
    response=$(curl -sk https://localhost:8086/api/v2/users/$uid -X DELETE -H "Authorization: Token $token")
    check_response "$response"
    ;;

  userenable)
    [ $# -ne 1 ] && usage
    email=$1
    log "Enabling user; email=$email"
    token=$(lookup_pillar_secret influx_token)
    uid=$(lookup_user_id "$token" "$email")
    response=$(curl -sk https://localhost:8086/api/v2/users/$uid -X PATCH -d "{\"name\":\"$email\",\"status\":\"active\"}" -H "Authorization: Token $token")
    check_response "$response"
    ;;

  userdisable)
    [ $# -ne 1 ] && usage
    email=$1
    log "Disabling user; email=$email"
    token=$(lookup_pillar_secret influx_token)
    uid=$(lookup_user_id "$token" "$email")
    response=$(curl -sk https://localhost:8086/api/v2/users/$uid -X PATCH -d "{\"name\":\"$email\",\"status\":\"inactive\"}" -H "Authorization: Token $token")
    check_response "$response"
    ;;

  *)
    usage
    ;;
esac