#!/usr/bin/python3
# -*- coding: utf-8 -*-
import sys
import re
import sqlite3
import os.path
import ldap
import configparser
import time
import syslog


def get_domain():
    f = open('/etc/hostname')
    domain = f.read()
    f.close()
    return domain.rstrip()


def get_counter(now):
    db.execute("CREATE TABLE IF NOT EXISTS counter(time integer, cnt integer)")

    #delete old data
    cleanup_range(now)

    #return 1 hour sum counter value
    db.execute("SELECT sum(cnt) FROM counter")
    count = db.fetchone()[0]
    if count:
        return int(count)
    else:
        return 0


def cleanup_range(now):
    delta = now - 3600
    db.execute("DELETE FROM counter WHERE time<?", (delta,))
    con.commit()


def set_counter(now,mailnum):
    db.execute("INSERT INTO counter VALUES(?,?)", (now, mailnum))
    con.commit()


def get_limit():
    limit_value = 750
    try:
        # get the diskquota value from LDAP
        l = ldap.initialize(ldaphost)
        l.protocol_version = ldap.VERSION3
        l.simple_bind_s(username, password)
        searchScope = ldap.SCOPE_SUBTREE
        retrieveAttributes = ['diskquota']
        search_results = l.search_s(baseDN, searchScope, searchFilter, retrieveAttributes)
        quota = search_results[1][1].get('diskquota')[0]
        f.write("diskquota: "+str(quota)+"\n")
        if str(quota) == '80000':
            limit_value = 1500
        else:
            limit_value = 750
    except:
        pass

    f.write("limit: "+str(limit_value)+"\n")
    syslog.syslog(syslog.LOG_INFO, "limit value: "+str(limit_value))
    return limit_value


def check_throttle(data):

    sender = data.get('sender')
    if "@" in sender:
        sender_domain = sender.split('@', 1)[1]
    else:
        sender_domain = ''

    recipient = data.get('recipient')
    if "@" in recipient:
        recipient_domain = recipient.split('@', 1)[1]
    else:
        recipient_domain = ''

    mailnum = int(data.get('recipient_count'))

    # receive
    if sender_domain != domain and recipient_domain == domain:
        f.write("direction: receive domain: "+str(domain)+" senderdomain: "+str(sender_domain)+" receivedomain: "+str(recipient_domain))
        syslog.syslog(syslog.LOG_INFO, "direction: receive domain: "+str(domain)+" senderdomain: "+str(sender_domain)+" receivedomain: "+str(recipient_domain))
        return True
    # send
    else:
        if sender_domain != domain and recipient_domain == '':
            mailnum = 1
        now = int(time.time())
        counter_value = get_counter(now)
        limit_value = get_limit()
        f.write("direction: send\n")
        syslog.syslog(syslog.LOG_INFO, "direction: send counter: "+str(counter_value+mailnum)+" domain: "+str(domain)+" senderdomain: "+str(sender_domain)+" receivedomain: "+str(recipient_domain))
        #limit_value = 10 # for debug
        if counter_value+mailnum > limit_value:
            return False
        else:
            set_counter(now,mailnum)
            return True

if __name__ == '__main__':

    try:

        setting = ConfigParser.SafeConfigParser()
        setting.read("/var/www/ControlPanel/conf/setting.conf")
        dbpath = '/var/throttle/throttle.db'
        f = open("/var/throttle/debug.log","a")
        syslog.openlog('throttle', syslog.LOG_PID|syslog.LOG_PERROR, syslog.LOG_LOCAL0)

        # LDAP settings
        ldaphost = setting.get("ldap", "host")
        baseDN = setting.get("ldap", "dn_name")
        username = setting.get("ldap", "base_dn_name")
        password = setting.get("ldap", "pass")
        searchFilter = "(uid=*)"

        data = {}
        lineRx = re.compile(r'^\s*([^=\s]+)\s*=(.*)$')

        con = sqlite3.connect(dbpath)
        db = con.cursor()

        domain = get_domain()

        while 1:
            line = sys.stdin.readline() 
            if not line: break
            line = line.rstrip()

            #  end of entry
            if not line:

                result = check_throttle(data)
                
                if result:
                    sys.stdout.write('action=dunno\n\n')
                    sys.stdout.flush()
                else:
                    sys.stdout.write('action=451 mail rate limit exceeded -- try again later\n\n')
                    sys.stdout.flush()
                data = {}

                f.write("\n\n")
                f.flush()
                continue

            #  parse line
            m = lineRx.match(line)
            if not m: 
                continue

            #  save the string
            key = m.group(1)
            value = m.group(2)
            data[key] = value

            f.write(line+"\n")

        con.close()
        syslog.closelog()
        f.close()

    except Exception as e:
        syslog.syslog(syslog.LOG_INFO, "error: "+str(type(e))+" "+str(e))
        sys.stdout.write('action=dunno\n\n')
        sys.stdout.flush()