#!/bin/bash

METASPLOIT_BASEDIR=/usr/share/metasploit-framework

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

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

pw_gen() {
  openssl rand -base64 32
}

pg_cmd() {
  su - postgres -c "$*"
}

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

usage() {
  PROG="$(basename $0)"
  echo
  echo "Manage the metasploit framework database"
  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

