#!/bin/bash

METASPLOIT_BASEDIR=/usr/share/metasploit-framework

DB_CONF=${METASPLOIT_BASEDIR}/config/database.yml
DB_NAME=msf
DB_USER=msf
DB_PORT="${PGPORT:-5432}"
PG_SERVICE=postgresql

## Bash Colour output
RED="\033[01;31m"
GREEN="\033[01;32m"
YELLOW="\033[01;33m"
RESET="\033[00m"

if echo "$DB_PORT" | grep -q '[^0-9]' || [ $DB_PORT -lt 1 ] || [ $DB_PORT -gt 65535 ]; then
  echo -e "${RED}[-]${RESET} Error: invalid port number, should be between 1 and 65535" 1>&2
  exit 1
fi

pw_gen() {
  openssl rand -base64 32
}

pg_cmd() {
  local command
  local command_args
  command="$1"
  command_args="${@:2}"
  shift
  su - postgres -c "$command -p $DB_PORT $command_args"
}

db_exists() {
  if pg_cmd psql -lqt | cut -d \| -f 1 | grep -qw $1; then
    return 0
  fi
  return 1
}

user_exists() {
  if echo "SELECT usename FROM pg_user;" | pg_cmd psql -qt postgres | grep -qw $1; then
    return 0
  fi
  return 1
}

start_db() {
  if ! systemctl status ${PG_SERVICE} >/dev/null; then
    echo -e "${GREEN}[+]${RESET} Starting database"
    systemctl start ${PG_SERVICE}
  else
    echo -e "${YELLOW}[i]${RESET} Database already started"
  fi
}

stop_db() {
  if systemctl status ${PG_SERVICE} >/dev/null; then
    echo -e "${GREEN}[+]${RESET} Stopping database"
    systemctl stop ${PG_SERVICE}
  else
    echo -e "${YELLOW}[i]${RESET} Database already stopped"
  fi
}

init_db() {
  start_db

  if [ -e ${DB_CONF} ]; then
    echo -e "${YELLOW}[i]${RESET} The database appears to be already configured, skipping initialization"
    return
  fi

  DB_PASS=$(pw_gen)
  if user_exists ${DB_USER}; then
    echo -e "${GREEN}[+]${RESET} Resetting password of database user '${DB_USER}'"
    printf "ALTER ROLE ${DB_USER} WITH PASSWORD '$DB_PASS';\n" | pg_cmd psql postgres >/dev/null
  else
    echo -e "${GREEN}[+]${RESET} Creating database user '${DB_USER}'"
    TMPFILE=$(mktemp) || (echo -e "${RED}[-]${RESET} Error: Couldn't create temp file" && exit 1)
    printf "%s\n%s\n" "${DB_PASS}" "${DB_PASS}" | pg_cmd createuser -S -D -R -P ${DB_USER} >/dev/null 2>"${TMPFILE}"
    grep -v "^Enter password for new role: $\|^Enter it again: $" "${TMPFILE}"
    rm -f "${TMPFILE}"
  fi

  if ! db_exists ${DB_NAME}; then
    echo -e "${GREEN}[+]${RESET} Creating databases '${DB_NAME}'"
    pg_cmd createdb ${DB_NAME} -O ${DB_USER} -T template0 -E UTF-8
  fi

  if ! db_exists ${DB_NAME}_test; then
    echo -e "${GREEN}[+]${RESET} Creating databases '${DB_NAME}_test'"
    pg_cmd createdb ${DB_NAME}_test -O ${DB_USER} -T template0 -E UTF-8
  fi

  echo -e "${GREEN}[+]${RESET} Creating configuration file '${DB_CONF}'"
  cat > ${DB_CONF} <<-EOF
development:
  adapter: postgresql
  database: ${DB_NAME}
  username: ${DB_USER}
  password: ${DB_PASS}
  host: localhost
  port: $DB_PORT
  pool: 5
  timeout: 5

production:
  adapter: postgresql
  database: ${DB_NAME}
  username: ${DB_USER}
  password: ${DB_PASS}
  host: localhost
  port: $DB_PORT
  pool: 5
  timeout: 5

test:
  adapter: postgresql
  database: ${DB_NAME}_test
  username: ${DB_USER}
  password: ${DB_PASS}
  host: localhost
  port: $DB_PORT
  pool: 5
  timeout: 5
EOF

  echo -e "${GREEN}[+]${RESET} Creating initial database schema"
  cd ${METASPLOIT_BASEDIR}/
  bundle exec rake db:migrate >/dev/null
}

delete_db() {
  start_db

  if db_exists ${DB_NAME}; then
    echo -e "${GREEN}[+]${RESET} Dropping databases '${DB_NAME}'"
    pg_cmd dropdb ${DB_NAME}
  fi

  if db_exists ${DB_NAME}_test; then
    echo -e "${GREEN}[+]${RESET} Dropping databases '${DB_NAME}_test'"
    pg_cmd dropdb ${DB_NAME}_test
  fi

  if user_exists ${DB_USER}; then
    echo -e "${GREEN}[+]${RESET} Dropping database user '${DB_USER}'"
    pg_cmd dropuser ${DB_USER}
  fi

  echo -e "${GREEN}[+]${RESET} Deleting configuration file ${DB_CONF}"
  rm -f ${DB_CONF}

  stop_db
}

reinit_db() {
  delete_db
  init_db
}

status_db() {
  systemctl --no-pager -l status ${PG_SERVICE}

  ## Check if the port is free
  if lsof -Pi :${DB_PORT} -sTCP:LISTEN -t >/dev/null ; then
    echo ""
    ## Show network servies
    lsof -Pi :${DB_PORT} -sTCP:LISTEN
    echo ""
    ## Show running processes
    ps -f $( lsof -Pi :${DB_PORT} -sTCP:LISTEN -t )
    echo ""
  else
    echo -e "${YELLOW}[i]${RESET} No network service running"
  fi

  if [ -e ${DB_CONF} ]; then
    echo -e "${GREEN}[+]${RESET} Detected configuration file (${DB_CONF})"
  else
    echo -e "${YELLOW}[i]${RESET} No configuration file found"
  fi
}

run_db() {
  ## Is there a config file?
  if [ ! -e ${DB_CONF} ]; then
    ## No, so lets create one (first time run!)
    init_db
  else
   ## There is, so just start up
    start_db
  fi

  ## Run metasploit-framework's main console
  exec msfconsole
}

usage() {
  PROG="$(basename $0)"
  echo
  echo "Manage the metasploit framework database"
  echo
  echo "You can use an specific port number for the"
  echo "PostgreSQL connection setting the PGPORT variable"
  echo "in the current shell."
  echo
  echo "Example: PGPORT=5433 ${PROG} init"
  echo
  echo "  ${PROG} init     # start and initialize the database"
  echo "  ${PROG} reinit   # delete and reinitialize the database"
  echo "  ${PROG} delete   # delete database and stop using it"
  echo "  ${PROG} start    # start the database"
  echo "  ${PROG} stop     # stop the database"
  echo "  ${PROG} status   # check service status"
  echo "  ${PROG} run      # start the database and run msfconsole"
  echo
  exit 0
}

if [ "$#" -ne 1 ]; then
  usage
fi

if [ $(id -u) -ne 0 ]; then
  echo -e "${RED}[-]${RESET} Error: $0 must be ${RED}run as root${RESET}" 1>&2
  exit 1
fi

case $1 in
  init) init_db ;;
  reinit) reinit_db ;;
  delete) delete_db ;;
  start) start_db ;;
  stop) stop_db ;;
  status) status_db ;;
  run) run_db ;;
  *) echo -e "${RED}[-]${RESET} Error: unrecognized action '${RED}${1}${RESET}'" 1>&2; usage ;;
esac

