#!/usr/bin/python

import bcrypt
import sys
import psycopg2
import logging
import struct
from struct import *

db_name = "my_db"
db_user = "my_id"
db_pass = "my_pass"
db_host = "localhost"
db_table = "my_table"
db_username_field = "name"
db_password_field = "pass"
domain_suffix = "@example.net"

sys.stderr = open('/tmp/pg_auth_err.log', 'w')
logging.basicConfig(level=logging.DEBUG,
                    format='%(asctime)s %(levelname)s %(message)s',
                    filename='/tmp/pg_auth_extauth.log',
                    filemode='w')

logging.info('extauth script started, waiting for ejabberd requests')


def hash_password(password, salt=None):
    if not salt:
        salt = bcrypt.gensalt()
    return bcrypt.hashpw(password, salt)


def get_connection():
    return psycopg2.connect(
        "dbname='%s' user='%s' host='%s' password='%s'" % (
            db_name, db_user, db_host, db_pass
            ))


def ejabberd_in():
    logging.debug("trying to read 2 bytes from ejabberd:")
    try:
        input_length = sys.stdin.read(2)
    except IOError:
        logging.debug("ioerror")
        if len(input_length) != 2:
            logging.debug("ejabberd sent us wrong things!")
            raise Exception('Wrong input from ejabberd!')

    logging.debug('got 2 bytes via stdin: %s' % input_length)
    (size, ) = unpack('>h', input_length)
    logging.debug('size of data: %i' % size)
    income = sys.stdin.read(size).split(':')
    logging.debug("incoming data: %s" % income)
    return income


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


def ejabberd_out(bool):
    logging.debug("Ejabberd gets: %s" % bool)
    token = genanswer(bool)
    logging.debug("sent bytes: %#x %#x %#x %#x" % (
            ord(token[0]), ord(token[1]), ord(token[2]), ord(token[3])))
    sys.stdout.write(token)
    sys.stdout.flush()


def db_entry(in_user):
    con = get_connection()
    cur = con.cursor()
    cur.execute("SELECT %s,%s FROM %s WHERE %s ='%s'" % (
            db_username_field, db_password_field, db_table,
            db_username_field, in_user))
    return cur.fetchone()


def isuser(in_user, in_host):
    data = db_entry(in_user)
    out = False
    if data == None or data[0] == None:
        out = False
        logging.debug("Wrong username: %s" % (in_user))
    name = "%s@%s" % (in_user, in_host)
    try:
        resultname = "%s%s" % (data[0], domain_suffix)
    except:
        return False

    if name == resultname:
        out = True
    return out


def tryregister(username, domain, password):
    if isuser(username, domain):
        return False

    query = "INSERT INTO USERS (username, password) values ('%s', '%s');" % (
        username, hash_password(password))

    con = get_connection()
    cur = con.cursor()
    cur.execute(query)
    try:
        con.commit()
        return True
    except:
        return False


def setpass(username, domain, password):
    if not isuser(username, domain):
        logging.info("Trying to change password for invalid user %s@%s" % (
                username, domain))
        return False

    query = "UPDATE USERS set password='%s' where username='%s';" % (
        hash_password(password), username)

    con = get_connection()
    cur = con.cursor()
    cur.execute(query)
    try:
        con.commit()
        return True
    except:
        import traceback
        traceback.print_exc()
        sys.stderr.write(traceback.format_exc())
        sys.stderr.flush()
        return False


def removeuser(username):
    query = "DELETE FROM USERS where username='%s';" % (
        username)

    con = get_connection()
    cur = con.cursor()
    cur.execute(query)
    try:
        con.commit()
        return True
    except:
        return False


def auth(in_user, in_host, password):
    data = db_entry(in_user)
    out = False
    if data == None:
        logging.debug("Wrong username: %s" % (in_user))
        return False

    name = "%s@%s" % (in_user, in_host)
    resultname = "%s%s" % (data[0], domain_suffix)
    if name == resultname:
        if hash_password(password, data[1]) == data[1]:
            out = True
        else:
            logging.debug("Wrong password for user: %s" % (in_user))
            out = False
    else:
        out = False

    return out


def log_result(op, in_user, bool):
    if bool:
        logging.info("%s successful for %s\n" % (op, in_user))
    else:
        logging.info("%s unsuccessful for %s\n" % (op, in_user))


if __name__ == '__main__':
    while True:
        logging.debug("start of processing loop")
        try:
            ejab_request = ejabberd_in()
        except Exception, inst:
            logging.info("Exception occured: %s", inst)
            continue

        logging.debug('operation: %s' % (ejab_request[0]))
        op_result = False
        if ejab_request[0] == "auth":
            op_result = auth(ejab_request[1], ejab_request[2], ejab_request[3])
            ejabberd_out(op_result)
            log_result(ejab_request[0], ejab_request[1], op_result)
        elif ejab_request[0] == "isuser":
            op_result = isuser(ejab_request[1], ejab_request[2])
            ejabberd_out(op_result)
            log_result(ejab_request[0], ejab_request[1], op_result)
        elif ejab_request[0] == "setpass":
            op_result = setpass(ejab_request[1],
                                ejab_request[2], ejab_request[3])
            ejabberd_out(op_result)
            log_result(ejab_request[0], ejab_request[1], op_result)
        elif ejab_request[0] == 'tryregister':
            op_result = tryregister(ejab_request[1], ejab_request[2],
                                    ejab_request[3])
            ejabberd_out(op_result)
            log_result(ejab_request[0], ejab_request[1], op_result)
        elif ejab_request[0] == 'removeuser':
            op_result = removeuser(ejab_request[1])
            ejabberd_out(op_result)
            log_result(ejab_request[0], ejab_request[1], op_result)
        elif ejab_request[0] == 'removeuser3':
            op_result = removeuser(ejab_request[1])
            ejabberd_out(op_result)
            log_result(ejab_request[0], ejab_request[1], op_result)

    logging.debug("end of infinite loop")
    logging.info('extauth script terminating')