#!/bin/bash
#
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.
#

################################
# constants
################################

FLUME_AGENT_CLASS="org.apache.flume.node.Application"
FLUME_AVRO_CLIENT_CLASS="org.apache.flume.client.avro.AvroCLIClient"
FLUME_VERSION_CLASS="org.apache.flume.tools.VersionInfo"
FLUME_TOOLS_CLASS="org.apache.flume.tools.FlumeToolsMain"

CLEAN_FLAG=1
################################
# functions
################################

info() {
  if [ ${CLEAN_FLAG} -ne 0 ]; then
    local msg=$1
    echo "Info: $msg" >&2
  fi
}

warn() {
  if [ ${CLEAN_FLAG} -ne 0 ]; then
    local msg=$1
    echo "Warning: $msg" >&2
  fi
}

error() {
  local msg=$1
  local exit_code=$2

  echo "Error: $msg" >&2

  if [ -n "$exit_code" ] ; then
    exit $exit_code
  fi
}

# If avail, add Hadoop paths to the FLUME_CLASSPATH and to the
# FLUME_JAVA_LIBRARY_PATH env vars.
# Requires Flume jars to already be on FLUME_CLASSPATH.
add_hadoop_paths() {
  local HADOOP_IN_PATH=$(PATH="${HADOOP_HOME:-${HADOOP_PREFIX}}/bin:$PATH" \
      which hadoop 2>/dev/null)

  if [ -f "${HADOOP_IN_PATH}" ]; then
    info "Including Hadoop libraries found via ($HADOOP_IN_PATH) for HDFS access"

    # determine hadoop java.library.path and use that for flume
    local HADOOP_CLASSPATH=""
    local HADOOP_JAVA_LIBRARY_PATH=$(HADOOP_CLASSPATH="$FLUME_CLASSPATH" \
        ${HADOOP_IN_PATH} org.apache.flume.tools.GetJavaProperty \
        java.library.path)

    # look for the line that has the desired property value
    # (considering extraneous output from some GC options that write to stdout)
    # IFS = InternalFieldSeparator (set to recognize only newline char as delimiter)
    IFS=$'\n'
    for line in $HADOOP_JAVA_LIBRARY_PATH; do
      if [[ $line =~ ^java\.library\.path=(.*)$ ]]; then
        HADOOP_JAVA_LIBRARY_PATH=${BASH_REMATCH[1]}
        break
      fi
    done
    unset IFS

    if [ -n "${HADOOP_JAVA_LIBRARY_PATH}" ]; then
      FLUME_JAVA_LIBRARY_PATH="$FLUME_JAVA_LIBRARY_PATH:$HADOOP_JAVA_LIBRARY_PATH"
    fi

    # determine hadoop classpath
    HADOOP_CLASSPATH=$($HADOOP_IN_PATH classpath)

    FLUME_CLASSPATH="$FLUME_CLASSPATH:$HADOOP_CLASSPATH"
  fi
}
add_HBASE_paths() {
  local HBASE_IN_PATH=$(PATH="${HBASE_HOME}/bin:$PATH" \
      which hbase 2>/dev/null)

  if [ -f "${HBASE_IN_PATH}" ]; then
    info "Including HBASE libraries found via ($HBASE_IN_PATH) for HBASE access"

    # determine HBASE java.library.path and use that for flume
    local HBASE_CLASSPATH=""
    local HBASE_JAVA_LIBRARY_PATH=$(HBASE_CLASSPATH="$FLUME_CLASSPATH" \
        ${HBASE_IN_PATH} org.apache.flume.tools.GetJavaProperty \
        java.library.path)

    # look for the line that has the desired property value
    # (considering extraneous output from some GC options that write to stdout)
    # IFS = InternalFieldSeparator (set to recognize only newline char as delimiter)
    IFS=$'\n'
    for line in $HBASE_JAVA_LIBRARY_PATH; do
      if [[ $line =~ ^java\.library\.path=(.*)$ ]]; then
        HBASE_JAVA_LIBRARY_PATH=${BASH_REMATCH[1]}
        break
      fi
    done
    unset IFS

    if [ -n "${HBASE_JAVA_LIBRARY_PATH}" ]; then
      FLUME_JAVA_LIBRARY_PATH="$FLUME_JAVA_LIBRARY_PATH:$HBASE_JAVA_LIBRARY_PATH"
    fi

    # determine HBASE classpath
    HBASE_CLASSPATH=$($HBASE_IN_PATH classpath)

    FLUME_CLASSPATH="$FLUME_CLASSPATH:$HBASE_CLASSPATH"
    FLUME_CLASSPATH="$FLUME_CLASSPATH:$HBASE_HOME/conf"

  fi
}

add_hive_paths(){
  if [ -d "${HIVE_HOME}/lib" ]; then
    info "Including Hive libraries found via ($HIVE_HOME) for Hive access"
    FLUME_CLASSPATH="$FLUME_CLASSPATH:$HIVE_HOME/lib/*"
  fi
  if [ -d "${HCAT_HOME}/share/hcatalog" ]; then
    info "Including HCatalog libraries found via ($HCAT_HOME) for Hive access"
    FLUME_CLASSPATH="$FLUME_CLASSPATH:${HCAT_HOME}/share/hcatalog/*"
  fi
}

set_LD_LIBRARY_PATH(){
#Append the FLUME_JAVA_LIBRARY_PATH to whatever the user may have specified in
#flume-env.sh
  if [ -n "${FLUME_JAVA_LIBRARY_PATH}" ]; then
    export LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:${FLUME_JAVA_LIBRARY_PATH}"
  fi
}

display_help() {
  cat <<EOF
Usage: $0 <command> [options]...

commands:
  help                      display this help text
  agent                     run a Flume agent
  avro-client               run an avro Flume client
  version                   show Flume version info

global options:
  --conf-file,-f <conf>     use configs in <conf> directory
  --conf-uri, -u <conf>     use configs located at <conf>
  --conf-provider, -prov <provider-class> use the configuration provided by the provider class
  --conf-user, -user <user>    the user for accessing the configuration uri
  --conf-password, -pwd <password> the password for accessing the configuration uri
  --classpath,-C <cp>       append to the classpath
  --dryrun,-d               do not actually start Flume, just print the command
  --plugins-path <dirs>     colon-separated list of plugins.d directories. See the
                            plugins.d section in the user guide for more details.
                            Default: \$FLUME_HOME/plugins.d
  -Dproperty=value          sets a Java system property value
  -Xproperty=value          sets a Java -X option

agent options:
  --name,-n <name>          the name of this agent (required)
  --conf-file,-f <file>     specify a config file (required if -z, -u, and -prov are missing)
  --zkConnString,-z <str>   specify the ZooKeeper connection to use (required if -f missing)
  --zkBasePath,-p <path>    specify the base path in ZooKeeper for agent configs
  --no-reload-conf          do not reload config file if changed
  --help,-h                 display help text

avro-client options:
  --rpcProps,-P <file>   RPC client properties file with server connection params
  --host,-H <host>       hostname to which events will be sent
  --port,-p <port>       port of the avro source
  --dirname <dir>        directory to stream to avro source
  --filename,-F <file>   text file to stream to avro source (default: std input)
  --headerFile,-R <file> File containing event headers as key/value pairs on each new line
  --help,-h              display help text

  Either --rpcProps or both --host and --port must be specified.

Note that if <conf> directory is specified, then it is always included first
in the classpath.

EOF
}

run_flume() {
  local FLUME_APPLICATION_CLASS

  if [ "$#" -gt 0 ]; then
    FLUME_APPLICATION_CLASS=$1
    shift
  else
    error "Must specify flume application class" 1
  fi

  if [ ${CLEAN_FLAG} -ne 0 ]; then
    set -x
  fi
  $EXEC $JAVA_HOME/bin/java $JAVA_OPTS $FLUME_JAVA_OPTS "${arr_java_props[@]}" -cp "$FLUME_CLASSPATH" \
      -Djava.library.path=$FLUME_JAVA_LIBRARY_PATH "$FLUME_APPLICATION_CLASS" $*
}

################################
# main
################################

# set default params
FLUME_CLASSPATH=""
FLUME_JAVA_LIBRARY_PATH=""
JAVA_OPTS="-Xmx20m"
LD_LIBRARY_PATH=""

opt_conf=""
opt_classpath=""
opt_plugins_dirs=""
arr_java_props=()
arr_java_props_ct=0
opt_dryrun=""

mode=$1
shift

case "$mode" in
  help)
    display_help
    exit 0
    ;;
  agent)
    opt_agent=1
    ;;
  node)
    opt_agent=1
    warn "The \"node\" command is deprecated. Please use \"agent\" instead."
    ;;
  avro-client)
    opt_avro_client=1
    ;;
  tool)
    opt_tool=1
    ;;
  version)
   opt_version=1
   CLEAN_FLAG=0
   ;;
  *)
    error "Unknown or unspecified command '$mode'"
    echo
    display_help
    exit 1
    ;;
esac

args=""
while [ -n "$*" ] ; do
  arg=$1
  shift

  case "$arg" in
    --conf|-c)
      [ -n "$1" ] || error "Option --conf requires an argument" 1
      opt_conf=$1
      shift
      ;;
    --classpath|-C)
      [ -n "$1" ] || error "Option --classpath requires an argument" 1
      opt_classpath=$1
      shift
      ;;
    --dryrun|-d)
      opt_dryrun="1"
      ;;
    --plugins-path)
      opt_plugins_dirs=$1
      shift
      ;;
    -agentlib*)
      arr_java_props[arr_java_props_ct]=$arg
      ((++arr_java_props_ct))
      ;;
    -agentpath*)
      arr_java_props[arr_java_props_ct]=$arg
      ((++arr_java_props_ct))
      ;;
    -javaagent*)
      arr_java_props[arr_java_props_ct]=$arg
      ((++arr_java_props_ct))
      ;;
    -D*)
      arr_java_props[arr_java_props_ct]=$arg
      ((++arr_java_props_ct))
      ;;
    -X*)
      arr_java_props[arr_java_props_ct]=$arg
      ((++arr_java_props_ct))
      ;;
    *)
      args="$args $arg"
      ;;
  esac
done

# make opt_conf absolute
if [[ -n "$opt_conf" && -d "$opt_conf" ]]; then
  opt_conf=$(cd $opt_conf; pwd)
fi

# allow users to override the default env vars via conf/flume-env.sh
if [ -z "$opt_conf" ]; then
  warn "No configuration directory set! Use --conf <dir> to override."
elif [ -f "$opt_conf/flume-env.sh" ]; then
  info "Sourcing environment configuration script $opt_conf/flume-env.sh"
  source "$opt_conf/flume-env.sh"
fi

# prepend command-line classpath to env script classpath
if [ -n "${opt_classpath}" ]; then
  if [ -n "${FLUME_CLASSPATH}" ]; then
    FLUME_CLASSPATH="${opt_classpath}:${FLUME_CLASSPATH}"
  else
    FLUME_CLASSPATH="${opt_classpath}"
  fi
fi

if [ -z "${FLUME_HOME}" ]; then
  FLUME_HOME=$(cd $(dirname $0)/..; pwd)
fi

# prepend $FLUME_HOME/lib jars to the specified classpath (if any)
if [ -n "${FLUME_CLASSPATH}" ] ; then
  FLUME_CLASSPATH="${FLUME_HOME}/lib/*:$FLUME_CLASSPATH"
else
  FLUME_CLASSPATH="${FLUME_HOME}/lib/*"
fi

# load plugins.d directories
PLUGINS_DIRS=""
if [ -n "${opt_plugins_dirs}" ]; then
  PLUGINS_DIRS=$(sed -e 's/:/ /g' <<<${opt_plugins_dirs})
else
  PLUGINS_DIRS="${FLUME_HOME}/plugins.d"
fi

unset plugin_lib plugin_libext plugin_native
for PLUGINS_DIR in $PLUGINS_DIRS; do
  if [[ -d ${PLUGINS_DIR} ]]; then
    for plugin in ${PLUGINS_DIR}/*; do
      if [[ -d "$plugin/lib" ]]; then
        plugin_lib="${plugin_lib}${plugin_lib+:}${plugin}/lib/*"
      fi
      if [[ -d "$plugin/libext" ]]; then
        plugin_libext="${plugin_libext}${plugin_libext+:}${plugin}/libext/*"
      fi
      if [[ -d "$plugin/native" ]]; then
        plugin_native="${plugin_native}${plugin_native+:}${plugin}/native"
      fi
    done
  fi
done

if [[ -n "${plugin_lib}" ]]
then
  FLUME_CLASSPATH="${FLUME_CLASSPATH}:${plugin_lib}"
fi

if [[ -n "${plugin_libext}" ]]
then
  FLUME_CLASSPATH="${FLUME_CLASSPATH}:${plugin_libext}"
fi

if [[ -n "${plugin_native}" ]]
then
  if [[ -n "${FLUME_JAVA_LIBRARY_PATH}" ]]
  then
    FLUME_JAVA_LIBRARY_PATH="${FLUME_JAVA_LIBRARY_PATH}:${plugin_native}"
  else
    FLUME_JAVA_LIBRARY_PATH="${plugin_native}"
  fi
fi

# find java
if [ -z "${JAVA_HOME}" ] ; then
  warn "JAVA_HOME is not set!"
  # Try to use Bigtop to autodetect JAVA_HOME if it's available
  if [ -e /usr/libexec/bigtop-detect-javahome ] ; then
    . /usr/libexec/bigtop-detect-javahome
  elif [ -e /usr/lib/bigtop-utils/bigtop-detect-javahome ] ; then
    . /usr/lib/bigtop-utils/bigtop-detect-javahome
  fi

  # Using java from path if bigtop is not installed or couldn't find it
  if [ -z "${JAVA_HOME}" ] ; then
    JAVA_DEFAULT=$(type -p java)
    [ -n "$JAVA_DEFAULT" ] || error "Unable to find java executable. Is it in your PATH?" 1
    JAVA_HOME=$(cd $(dirname $JAVA_DEFAULT)/..; pwd)
  fi
fi

# look for hadoop libs
add_hadoop_paths
add_HBASE_paths
add_hive_paths

# prepend conf dir to classpath
if [ -n "$opt_conf" ]; then
  FLUME_CLASSPATH="$opt_conf:$FLUME_CLASSPATH"
fi

set_LD_LIBRARY_PATH
# allow dryrun
EXEC="exec"
if [ -n "${opt_dryrun}" ]; then
  warn "Dryrun mode enabled (will not actually initiate startup)"
  EXEC="echo"
fi

# finally, invoke the appropriate command
if [ -n "$opt_agent" ] ; then
  run_flume $FLUME_AGENT_CLASS $args
elif [ -n "$opt_avro_client" ] ; then
  run_flume $FLUME_AVRO_CLIENT_CLASS $args
elif [ -n "${opt_version}" ] ; then
  run_flume $FLUME_VERSION_CLASS $args
elif [ -n "${opt_tool}" ] ; then
  run_flume $FLUME_TOOLS_CLASS $args
else
  error "This message should never appear" 1
fi

exit 0