Mes premières dents (les incisives) sur Django...

[progfou.git] / ejabberd / auth-mysql.py

#!/usr/bin/env python
# -*- coding: UTF-8 -*-
# Depends: python-mysqldb
# install -o root -g ejabberd -m 0750 -c auth-mysql.py /etc/ejabberd/
# install -o ejabberd -g adm -m 0640 -c /dev/null /var/log/ejabberd/auth-mysql.log
# Test: printf '\000\031isuser:moussa.nombre:test' | ./auth-mysql.py | hd

import sys
import traceback
import struct
import time
import MySQLdb
import crypt

# TODO: transformer toute la partie gestion de BdD en une classe
# TODO: connexion persistente et gérer les coupures de service MySQL

_host = 'nss'
_user = 'jabber'
_passwd = 'password'
_db = 'mail'
_timeout = 2
_query = "SELECT * FROM auforg_virtual WHERE source=%s AND LENGTH(password)>1"
_log_filename = '/var/log/ejabberd/auth-mysql.log'
_cache_positive_ttl = (15*60)
_cache_negative_ttl = (1*60)

# ce qui suit est à usage interne, ne pas toucher...
_find_user_cache = {}
_authenticate_user_cache = {}

def find_user(user, host):
    """
    Returns (found, in_cache)
    where 'found' means the (user,host) was found in the database
    and 'in_cache' means it was found from the local cache
    """
    now = time.time()

    global _find_user_cache
    c = _find_user_cache.get((user,host))
    if c != None:
        if c['found'] and now < (c['time'] + _cache_positive_ttl):
            return (True, True)
        if not c['found'] and now < (c['time'] + _cache_negative_ttl):
            return (False, True)

    global _host, _user, _passwd, _db
    db = MySQLdb.connect(host=_host, user=_user, passwd=_passwd,
                         db=_db, connect_timeout=_timeout)
    cur = db.cursor(MySQLdb.cursors.DictCursor)
    nrows = cur.execute(_query, ('%s@%s' % (user,host), ))
    del cur, db
    found = (nrows > 0)
    _find_user_cache[(user,host)] = {'found': found, 'time': now}
    return (found, False)

def authenticate_user(user, host, password):
    """
    Returns (valid, in_cache)
    where 'valid' means the (user,host,password) was valid in the database
    and 'in_cache' means it was validated against the local cache
    """
    now = time.time()

    global _authenticate_user_cache
    c = _authenticate_user_cache.get((user,host,password))
    if c != None:
        if c['valid'] and now < (c['time'] + _cache_positive_ttl):
            return (True, True)
        if not c['valid'] and now < (c['time'] + _cache_negative_ttl):
            return (False, True)

    global _host, _user, _passwd, _db
    db = MySQLdb.connect(host=_host, user=_user, passwd=_passwd,
                         db=_db, connect_timeout=_timeout)
    cur = db.cursor(MySQLdb.cursors.DictCursor)
    nrows = cur.execute(_query, ('%s@%s' % (user,host), ))
    users = cur.fetchall()
    del cur, db
    valid = False
    if nrows > 0:
        for u in users:
            if crypt.crypt(password, u['password']) == u['password']:
                valid = True
                break
    _authenticate_user_cache[(user,host,password)] = {'valid': valid, 'time': now}
    return (valid, False)

def main():
    log_file = open(_log_filename, 'a')
    while True:
        try:
            nread = sys.stdin.read(2)
            if len(nread) == 0:
                time.sleep(0.25)
                continue
            now = time.strftime('%Y/%m/%d %H:%M:%S', time.localtime())
            if len(nread) < 2:
                log_file.write('%s bytes_read=%d\n' % (now, len(nread)))
                log_file.flush()
                continue
            size = struct.unpack('>h', nread)[0]
            data = sys.stdin.read(size)
            (operation, data) = data.split(':', 1)
            if operation == 'auth':
                (user, host, password) = data.split(':', 2)
                log_file.write('%s operation=%s user=%s host=%s\n'
                               % (now, operation, user, host))
                log_file.flush()
                (result, in_cache) = authenticate_user(user, host, password)
            elif operation == 'isuser':
                (user, host) = data.split(':', 1)
                log_file.write('%s operation=%s user=%s host=%s\n'
                               % (now, operation, user, host))
                log_file.flush()
                (result, in_cache) = find_user(user, host)
            elif operation == 'setpass':
                (user, host, password) = data.split(':', 2)
                log_file.write('%s operation=%s user=%s host=%s\n'
                               % (now, operation, user, host))
                log_file.flush()
                #result = set_user_password(user, host, password)
                result = False
                in_cache = False
            else:
                result = False
                in_cache = False
            if in_cache:
                in_cache = ' (from cache)'
            else:
                in_cache = ''
            log_file.write('%s => result=%s%s\n' % (now, result, in_cache))
            log_file.flush()
            sys.stdout.write(struct.pack('>hh', 2, result and 1 or 0))
            sys.stdout.flush()
        except:
            traceback.print_exc(file=log_file)
        #sys.exit()

if __name__ == '__main__':
    main()