#!/usr/bin/env python3

import atexit
import logging
import struct
import sys
import time

import bcrypt
import psycopg2


# Database connection details. The credentials here need access to the
# Mastodon database, which "ejabberd" is unlikely to have on your system
# by default. You shoud grant SELECT privileges to ejabberd on the
# "accounts" and "users" tables, to play it safe, or include the
# Mastodon DB user credentials here (don't).
db_host = "localhost"
db_port = 5432
db_user = "ejabberd"
db_pass = ""
db_name = "mastodon"

# This is the query that pulls the password hash for the given user.
# Mastodon doesn't store the domain for local accounts in the database,
# so we ignore the host component and try to match username where the
# domain is NULL.
db_query_getpass = """
    SELECT users.encrypted_password AS password
    FROM accounts
        INNER JOIN users
        ON accounts.id = users.account_id
    WHERE
        lower(accounts.username) = %(user)s
        AND accounts.domain IS NULL
    """

########################################################################
# Setup
########################################################################

sys.stderr = open('/var/log/ejabberd/extauth_err.log', 'a')
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s %(levelname)s %(message)s',
    filename='/var/log/ejabberd/extauth.log',
    filemode='a',
)

try:
    # Connect to the DB, set autocommit and readonly otherwise
    # postgresql seems to have a tendency to keep things "idle in
    # transaction" and table locks eventually grind Mastodon to a halt.
    # We don't make any changes anyway.
    database = psycopg2.connect(
        host=db_host,
        user=db_user,
        password=db_pass,
        database=db_name,
        port=db_port,
    )
    database.set_session(readonly=True, autocommit=True)
    logging.debug(database.get_dsn_parameters())
except:
    logging.error("Unable to initialize database, check settings!")
    time.sleep(10)
    sys.exit(1)


@atexit.register
def close_db():
    cursor.close()
    database.close()


logging.info('auth-mastodon script started, waiting for ejabberd requests')


class EjabberdInputError(Exception):

    def __init__(self, value):
        self.value = value

    def __str__(self):
        return repr(self.value)


########################################################################
# Declarations
########################################################################

def ejabberd_in():
    logging.debug("trying to read 2 bytes from ejabberd:")

    input_length = sys.stdin.buffer.read(2)

    if len(input_length) is not 2:
        logging.debug("ejabberd sent us wrong things!")
        raise EjabberdInputError('Wrong input from ejabberd!')

    logging.debug('got 2 bytes via stdin: %s' % input_length)

    (size,) = struct.unpack('>h', input_length)
    logging.debug('size of data: %i' % size)

    income = sys.stdin.read(size)
    logging.debug("incoming data: %s" % income)

    return income


def ejabberd_out(bool):
    logging.debug("Ejabberd gets: %s" % bool)

    token = genanswer(bool)

    logging.debug(
        "sent bytes: %#x %#x %#x %#x"
        % (token[0], token[1], token[2], token[3]))

    sys.stdout.buffer.write(token)
    sys.stdout.buffer.flush()


def genanswer(bool):
    answer = 0
    if bool:
        answer = 1
    token = struct.pack('>hh', 2, answer)
    return token


def get_password(user, host):
    # Right now we ignore the host component, as Mastodon doesn't store
    # it for local accounts. It may be required one day, so the code to
    # handle passing it to the query is left in for now.
    cursor = database.cursor()
    cursor.execute(db_query_getpass, {"user": user.lower(), "host": host})
    data = cursor.fetchone()
    cursor.close()
    return data[0] if data != None else None


def isuser(user, host):
    return get_password(user, host) != None


def auth(user, host, password):
    db_password = get_password(user, host)
    if db_password == None:
        logging.debug("Wrong username: %s@%s" % (user, host))
        return False
    else:
        if bcrypt.checkpw(password.encode('utf8'), db_password.encode('utf8')):
            logging.debug("Validated %s against hash %s" % (user, db_password))
            return True
        else:
            logging.debug("Wrong password for user: %s@%s" % (user, host))
            return False


########################################################################
# Main Loop
########################################################################

exitcode = 0

while True:
    logging.debug("start of infinite loop")

    try:
        ejab_request = ejabberd_in().split(':', 3)
    except EOFError:
        break
    except Exception as e:
        logging.exception("Exception occured while reading stdin")
        raise

    op_result = False
    try:
        # Only 'auth' and 'isuser' implemented, placeholders left to
        # maybe expose other functions later but for now let's not even
        # think about modifying the Mastodon DB
        if ejab_request[0] == "auth":
            op_result = auth(ejab_request[1], ejab_request[2], ejab_request[3])
        elif ejab_request[0] == "isuser":
            op_result = isuser(ejab_request[1], ejab_request[2])
        elif ejab_request[0] == "setpass":
            op_result = False
        elif ejab_request[0] == "tryregister":
            op_result = False
        elif ejab_request[0] == "removeuser":
            op_result = False
        elif ejab_request[0] == "removeuser3":
            op_result = False
    except Exception:
        logging.exception("Exception occured")

    ejabberd_out(op_result)
    logging.info("successful" if op_result else "unsuccessful")

logging.debug("end of infinite loop")
logging.info('extauth script terminating')
database.close()
sys.exit(exitcode)