#!/bin/bash

METASPLOIT_BASEDIR=/opt/metasploit
PG_BASEDIR=/var/lib/postgres

DB_CONF=${METASPLOIT_BASEDIR}/config/database.yml
DB_NAME=msf
DB_USER=msf
DB_PORT="${2:-5432}"
PG_SERVICE=postgresql
PG_CONF=${PG_BASEDIR}/data/postgresql.conf
METASPLOIT_PROFILE=/etc/profile.d/metasploit.sh

## Bash Colour output
if tput setaf 0 &>/dev/null; then
    RESET="$(tput sgr0)"
    BOLD="$(tput bold)"
    GREEN="${BOLD}$(tput setaf 2)"
    RED="${BOLD}$(tput setaf 1)"
    YELLOW="${BOLD}$(tput setaf 3)"
else
    RESET="\e[0m"
    BOLD="\e[1m"
    GREEN="${BOLD}\e[32m"
    RED="${BOLD}\e[31m"
    YELLOW="${BOLD}\e[33m"
fi
readonly RESET BOLD GREEN RED YELLOW

# print error and exit
err() {
    echo "${RED}[-]${RESET} ERROR: ${@}"
}

# print warning
warn() {
    echo "${YELLOW}[!]${RESET} WARNING: ${@}"
}

# print message
msg() {
    echo "${GREEN}[+]${RESET} ${@}"
}

pw_gen() {
    openssl rand -base64 32
}

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

db_exists() {
    if pg_cmd "psql -p ${DB_PORT} -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 -p ${DB_PORT} -qt postgres" | grep -qw $1; then
        return 0
    fi
    return 1
}

check_pg() {
    if [ -e ${PG_CONF} ]; then
        sed -i "/^port/d" ${PG_CONF}
    else
        if [[ "$(stat -c "%U:%G" /var/lib/postgres/)" != "postgres:postgres" ]]; then
            chown -R postgres:postgres /var/lib/postgres/
        fi
        pg_cmd "initdb --locale en_US.UTF-8 -D '$PG_BASEDIR/data'" >/dev/null
    fi
    pg_cmd "echo 'port = ${DB_PORT}' >> '${PG_CONF}'"
}

check_port() {
    if ! [[ "${DB_PORT}" =~ ^[0-9]+$ ]] || ((${DB_PORT} < 1 || ${DB_PORT} > 65535)); then
        return 0
    fi
    return 1
}

start_db() {
    check_pg

    if ! systemctl status ${PG_SERVICE} >/dev/null; then
        msg "Starting database"
        systemctl start ${PG_SERVICE}
    else
        warn "Database already started"
    fi
}

stop_db() {
    if systemctl status ${PG_SERVICE} >/dev/null; then
        msg "Stopping database"
        systemctl stop ${PG_SERVICE}
    else
        warn "Database already stopped"
    fi
}

init_db() {
    start_db

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

    DB_PASS=$(pw_gen)
    if user_exists ${DB_USER}; then
        msg "Resetting password of database user '${DB_USER}'"
        printf "ALTER ROLE ${DB_USER} WITH PASSWORD '$DB_PASS';\n" | pg_cmd psql -p ${DB_PORT} postgres >/dev/null
    else
        msg "Creating database user '${DB_USER}'"
        TMPFILE=$(mktemp) || (err "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} -p ${DB_PORT} >/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
        msg "Creating databases '${DB_NAME}'"
        pg_cmd createdb ${DB_NAME} -O ${DB_USER} -T template0 -E UTF-8 -p ${DB_PORT}
    fi

    if ! db_exists ${DB_NAME}_test; then
        msg "Creating databases '${DB_NAME}_test'"
        pg_cmd createdb ${DB_NAME}_test -O ${DB_USER} -T template0 -E UTF-8 -p ${DB_PORT}
    fi

    msg "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
    chmod 644 ${DB_CONF}

    echo "export MSF_DATABASE_CONFIG=${DB_CONF}" >${METASPLOIT_PROFILE}
    chmod 644 ${METASPLOIT_PROFILE}
}

delete_db() {
    start_db

    if db_exists ${DB_NAME}; then
        msg "Dropping databases '${DB_NAME}'"
        pg_cmd dropdb ${DB_NAME} -p ${DB_PORT}
    fi

    if db_exists ${DB_NAME}_test; then
        msg "Dropping databases '${DB_NAME}_test'"
        pg_cmd dropdb ${DB_NAME}_test -p ${DB_PORT}
    fi

    if user_exists ${DB_USER}; then
        msg "Dropping database user '${DB_USER}'"
        pg_cmd dropuser ${DB_USER} -p ${DB_PORT}
    fi

    sed -i "/^port/d" ${PG_CONF}

    msg "Deleting configuration file ${DB_CONF}"
    rm -f ${DB_CONF} ${METASPLOIT_PROFILE}

    stop_db
}

reinit_db() {
    delete_db
    init_db
}

status_db() {
    ## Check if the port is free
    if lsof -Pi :${DB_PORT} -sTCP:LISTEN -t >/dev/null; then
        echo ""
        ## Show network services
        lsof -Pi :${DB_PORT} -sTCP:LISTEN
        echo ""
        ## Show running processes
        ps -f $(lsof -Pi :${DB_PORT} -sTCP:LISTEN -t)
        echo ""
    else
        warn "No network service running"
    fi

    if [ -e ${DB_CONF} ]; then
        msg "Detected configuration file (${DB_CONF})"
    else
        warn "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 "You can use an specific port number for the"
    echo "postgresql connection after the operation."
    echo "Example: ${PROG} init 5433"
    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 [ "$#" -lt 1 ]; then
    usage
fi

if check_port; then
    err "${DB_PORT} is not a valid port number"
    exit 1
fi

if [ $(id -u) -ne 0 ]; then
    err "$0 must be ${RED}run as root${RESET}"
    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 ;;
*)
    err "unrecognized action '${RED}${1}${RESET}'"
    usage
    ;;
esac
