/*
 * Copyright (C) 2006 Manuel Mausz (manuel@mausz.at)
 * Origin code copyright (c) mjd@digitaleveryware.com 2003
 *  (http://www.digitaleveryware.com/projects/greylisting/)
 *
 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License
 * as published by the Free Software Foundation; either
 * version 2 of the License, or (at your option) any later
 * version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software Foundation,
 * Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdarg.h>
#include <errno.h>
#include <mysql.h>

#define SQLCMDSIZE     2048
#define RET_NOTFOUND   0
#define RET_ACCEPT     1
#define RET_REJECT     2
#define RET_TEMPREJECT 3
#define CMD_TEMPREJECT "E451 temporary failure (#4.3.0)\n"
#define CMD_REJECT     "E553 sorry, your envelope sender has been denied (#5.7.1)\n"
#define LOGLEVEL_FATAL 1
#define LOGLEVEL_ERROR 2
#define LOGLEVEL_WARN  3
#define LOGLEVEL_INFO  4
#define LOGLEVEL_DEBUG 5
#define MAXCONFIGLINESIZE 1024 // maybe change this to dynamic allocation sometime
#if MYSQL_VERSION_ID >= 50003
# define QUERYSIZE 500
#else
# define QUERYSIZE 700
#endif

static char *configfile = "control/greylisting";
static char *mysql_host = NULL;
static char *mysql_user = NULL;
static char *mysql_pass = NULL;
static char *mysql_db   = NULL;
unsigned int mysql_port = 3306;
unsigned int block_expire  = 55;
unsigned int record_expire = 500;
unsigned int record_expire_good = 36;
static char *relay_ip;
static char *mail_from;
static char *rcpt_to;
static int loglevel = LOGLEVEL_WARN;

void gllog(unsigned int level,  char* format, ...)
{
  va_list args;
  if (level > loglevel)
    return;
  va_start(args, format);
  vfprintf(stderr, format, args);
  va_end(args);
}

int load_config(void)
{
  char *tmp, *delim, *atpos;
  FILE *config;
  char buf[MAXCONFIGLINESIZE];
  int i;
  unsigned int userlen;

  /* first check for logging var */
  tmp = getenv("GLLOGLEVEL");
  if (tmp)
    loglevel = atoi(tmp);

  /* check if greylisting is enabled */
  if (!getenv("GREYLISTING") || getenv("RELAYCLIENT"))
  {
    gllog(LOGLEVEL_DEBUG, "greylisting: greylisting is not enabled\n");
    return 0;
  }

  /* basic environment variables needed */
  relay_ip  = getenv("TCPREMOTEIP");
  mail_from = getenv("SMTPMAILFROM");
  rcpt_to   = getenv("SMTPRCPTTO");
  if (!relay_ip || !mail_from || !rcpt_to)
  {
    gllog(LOGLEVEL_FATAL, "greylisting: one of the following envvars is undefined: TCPREMOTEIP, SMTPMAILFROM, SMTPRCPTTO\n");
    return 0;
  }

  /* check for BATV ("prvs=X=u@d.t" minimum) */
  if (strlen(mail_from) > 11
      && mail_from[0] == 'p' && mail_from[1] == 'r' && mail_from[2] == 'v'
      && mail_from[3] == 's' && mail_from[4] == '=')
  {
    /* BATV: prvs=HASH=user@domain.tld */
    if ((delim = strchr(mail_from + 5, '=')))
      mail_from = delim + 1;
    /* BATV: prvs=user/HASH@domain.tld */
    else if ((delim = strchr(mail_from + 5, '/')) && (atpos = strchr(delim, '@')))
    {
      userlen = delim - mail_from - 5;
      memmove(atpos - userlen, mail_from + 5, userlen);
      mail_from = atpos - userlen;
    }
  }

  /* avoid buffer overflows (max. query is ~410 chars long) */
  if (strlen(relay_ip) + strlen(mail_from) + strlen(rcpt_to) > SQLCMDSIZE - QUERYSIZE)
  {
    gllog(LOGLEVEL_FATAL, "greylisting: buffer overflow protection occurs\n");
    return 0;
  }

  /* fetch config file path */
  tmp = getenv("GLCONFIGFILE");
  if (tmp)
    configfile = tmp;

  /* fetch config file content */
  gllog(LOGLEVEL_DEBUG, "greylisting: configfile=%s\n", configfile);
  config = fopen(configfile, "r");
  if (!config)
    gllog(LOGLEVEL_DEBUG, "greylisting: configfile error: %s\n", strerror(errno));
  else
  {
    while((tmp = fgets(buf, sizeof(buf), config)))
    {
      if (buf[0] == '#' || buf[0] == ';')
        continue;
      for(i = 0; i < strlen(buf) && buf[i] != '\r' && buf[i] != '\n'; i++);
      buf[i] = 0;
      if (strstr(tmp, "mysql_host=") == tmp)
      {
        free(mysql_host);
        mysql_host = strdup(tmp + strlen("mysql_host="));
      }
      else if (strstr(tmp, "mysql_port=") == tmp)
        mysql_port = atoi(tmp + strlen("mysql_port="));
      else if (strstr(tmp, "mysql_user=") == tmp)
      {
        free(mysql_user);
        mysql_user = strdup(tmp + strlen("mysql_user="));
      }
      else if (strstr(tmp, "mysql_pass=") == tmp)
      {
        free(mysql_pass);
        mysql_pass = strdup(tmp + strlen("mysql_pass="));
      }
      else if (strstr(tmp, "mysql_db=") == tmp)
      {
        free(mysql_db);
        mysql_db = strdup(tmp + strlen("mysql_db="));
      }
      else if (strstr(tmp, "block_expire=") == tmp)
        block_expire = atoi(tmp + strlen("block_expire="));
      else if (strstr(tmp, "record_expire=") == tmp)
        record_expire = atoi(tmp + strlen("record_expire="));
      else if (strstr(tmp, "record_expire_good=") == tmp)
        record_expire_good = atoi(tmp + strlen("record_expire_good="));
      else if (strstr(tmp, "loglevel=") == tmp && !getenv("GLLOGLEVEL"))
        loglevel = atoi(tmp + strlen("loglevel="));
    }
    fclose(config);
  }

  /* environment variables */
  tmp = getenv("GLMYSQLHOST");
  if (tmp)
  {
    free(mysql_host);
    mysql_host = strdup(tmp);
  }

  tmp = getenv("GLMYSQLPORT");
  if (tmp)
    mysql_port = atoi(tmp);

  tmp = getenv("GLMYSQLUSER");
  if (tmp)
  {
    free(mysql_user);
    mysql_user = strdup(tmp);
  }

  tmp = getenv("GLMYSQLPASS");
  if (tmp)
  {
    free(mysql_pass);
    mysql_pass = strdup(tmp);
  }

  tmp = getenv("GLMYSQLDB");
  if (tmp)
  {
    free(mysql_db);
    mysql_db = strdup(tmp);
  }

  tmp = getenv("GLBLOCKEXPIRE");
  if (tmp)
    block_expire = atoi(tmp);

  tmp = getenv("GLRECORDEXPIRE");
  if (tmp)
    record_expire = atoi(tmp);

  tmp = getenv("GLRECORDEXPIREGOOD");
  if (tmp)
    record_expire_good = atoi(tmp);

  /* logging */
  gllog(LOGLEVEL_DEBUG, "greylisting: mysql: host=%s, port=%d, user=%s, pass=******\n", mysql_host, mysql_port, mysql_user);
  gllog(LOGLEVEL_DEBUG, "greylisting: block_expire=%d, record_expire=%d, record_expire_good=%d\n", block_expire, record_expire, record_expire_good);
  return 1;
}

void cleanup()
{
  free(mysql_host);
  free(mysql_user);
  free(mysql_pass);
}

int mysql_query_wrapper(MYSQL *mysql, char *query)
{
  int result;

  result = mysql_query(mysql, query);
  gllog(LOGLEVEL_DEBUG, "greylisting: mysql: %s - ret=%d\n", query, result);
  return result;
}

/* check if relay_ip or rcpt_to is white-/blacklisted */
int check_listed(MYSQL *mysql)
{
  MYSQL_RES *res;
  MYSQL_ROW row;
  char query[SQLCMDSIZE];
  int found = RET_NOTFOUND;
  char *rcpt_to_esc = NULL;
  char *domain_esc = NULL;
  char *domain = NULL;

  domain = strrchr(rcpt_to, '@');
  /* fallback to full rcpt_to if there's no domain */
  domain = (domain) ? domain + 1 : rcpt_to;
  domain_esc = malloc(strlen(domain)*2 + 1);
  mysql_real_escape_string(mysql, domain_esc, domain, strlen(domain));
  rcpt_to_esc = malloc(strlen(rcpt_to)*2 + 1);
  mysql_real_escape_string(mysql, rcpt_to_esc, rcpt_to, strlen(rcpt_to));
  sprintf(query, "SET @uipaddr = inet_aton('%s'), @udomain = '%s', @urcpt_to = '%s'; ", relay_ip, domain_esc, rcpt_to_esc);
  if (mysql_query_wrapper(mysql, query))
  {
    gllog(LOGLEVEL_ERROR, "greylisting: mysql: %s\n", mysql_error(mysql));
    return 0;
  }
  free(domain_esc);
  free(rcpt_to_esc);

#if MYSQL_VERSION_ID >= 50003
  sprintf(query,
    "SELECT `id`, `block_expires` >= UTC_TIMESTAMP(), `block_expires` < UTC_TIMESTAMP() "
    "FROM `greylisting_lists` "
    "WHERE `record_expires` > UTC_TIMESTAMP() "
    "AND ( "
      "( "
        "`rcpt_to` IS NULL "
        "AND `ipaddr_start` <= @uipaddr "
        "AND @uipaddr <= `ipaddr_end` "
      ") OR ( "
        "`ipaddr` IS NULL "
        "AND ( "
          "`rcpt_to` = @udomain "
          "OR `rcpt_to` = @urcpt_to "
        ") "
      ") "
    ") "
    "ORDER BY (`ipaddr_end` - `ipaddr_start`) ASC "
    "LIMIT 1");
#else
  sprintf(query,
    "SELECT `id`, `block_expires` >= UTC_TIMESTAMP(), `block_expires` < UTC_TIMESTAMP() "
    "FROM `greylisting_lists` "
    "WHERE `record_expires` > UTC_TIMESTAMP() "
    "AND ( "
      "( "
        "`rcpt_to` IS NULL "
        "AND (@base := IF(INSTR(`ipaddr`, '.'), 32, 128)) "
        "AND IF( "
          "INSTR(`ipaddr`, '/'), "
          "(@ipaddr_start := inet_aton(substring_index(`ipaddr`, '/', 1))) "
          "AND (@ipaddr_count := POW(2, @base - substring_index(`ipaddr`, '/', -1))) "
          "AND (@ipaddr_end := @ipaddr_start + @ipaddr_count - 1), "
          "(@ipaddr_start := inet_aton(`ipaddr`)) "
          "AND (@ipaddr_end := @ipaddr_start) "
        ") "
        "AND @ipaddr_start <= @uipaddr "
        "AND @uipaddr <= @ipaddr_end "
      ") OR ( "
        "`ipaddr` IS NULL "
        "AND ( "
          "`rcpt_to` = @udomain "
          "OR `rcpt_to` = @urcpt_to "
        ") "
      ") "
    ") "
    "ORDER BY (@ipaddr_end - @ipaddr_start) ASC "
    "LIMIT 1");
#endif

  if (mysql_query_wrapper(mysql, query) ||
      !(res = mysql_store_result(mysql)))
  {
    gllog(LOGLEVEL_ERROR, "greylisting: mysql: %s\n", mysql_error(mysql));
    return 0;
  }

  if ((row = mysql_fetch_row(res)))
  {
    if (atoi(row[1]))
    {
      found = RET_REJECT;
      gllog(LOGLEVEL_INFO, "greylisting: %s/%s is blacklisted (id=%s) - rejecting\n", relay_ip, domain, row[0]);
    }
    else if (atoi(row[2]))
    {
      found = RET_ACCEPT;
      gllog(LOGLEVEL_INFO, "greylisting: %s/%s is whitelisted (id=%s) - accepting\n", relay_ip, domain, row[0]);
    }
  }

  mysql_free_result(res);
  return found;
}

int check_greylisted(MYSQL *mysql)
{
  MYSQL_RES *res;
  MYSQL_ROW row;
  char query[SQLCMDSIZE];
  char *mail_from_esc = NULL;
  char *rcpt_to_esc   = NULL;
  char *relay_ip_sub  = NULL;
  char *ptr;
  char ipdelimeter;
  int ret = RET_NOTFOUND;

  mail_from_esc = malloc(strlen(mail_from)*2 + 1);
  rcpt_to_esc   = malloc(strlen(rcpt_to)*2 + 1);
  mysql_real_escape_string(mysql, mail_from_esc, mail_from, strlen(mail_from));
  mysql_real_escape_string(mysql, rcpt_to_esc, rcpt_to, strlen(rcpt_to));

  /*
   * 0 ... query matches anything in the same /24 subnet
   * 1 ... query does an exact ip match
   */
  if (0)
  {
    sprintf(query,
      "SELECT `id`, `block_expires` < UTC_TIMESTAMP() "
      "FROM `greylisting_data` "
      "WHERE `record_expires` > UTC_TIMESTAMP() "
        "AND `relay_ip` = '%s' "
        "AND `mail_from` = '%s' "
        "AND `rcpt_to` = '%s' "
      "LIMIT 1",
      relay_ip,
      mail_from_esc,
      rcpt_to_esc);
  }
  else
  {
    /* strip off the last octet */
    ipdelimeter = '.';
    if (!strchr(relay_ip, '.'))
      ipdelimeter = ':';
    relay_ip_sub = strdup(relay_ip);
    ptr = strrchr(relay_ip_sub, ipdelimeter);
    if (ptr)
      *ptr = '\0';

    sprintf(query,
      "SELECT `id`, `block_expires` < UTC_TIMESTAMP() "
      "FROM `greylisting_data` "
      "WHERE `record_expires` > UTC_TIMESTAMP() "
        "AND `relay_ip` LIKE '%s%c%%' "
        "AND `mail_from` = '%s' "
        "AND `rcpt_to` = '%s' "
      "LIMIT 1",
      relay_ip_sub,
      ipdelimeter,
      mail_from_esc,
      rcpt_to_esc);
    free(relay_ip_sub);
  }

  if (mysql_query_wrapper(mysql, query) ||
      !(res = mysql_store_result(mysql)))
  {
    gllog(LOGLEVEL_ERROR, "greylisting: mysql: %s\n", mysql_error(mysql));
    return RET_NOTFOUND;
  }

  if ((row = mysql_fetch_row(res)))
  {
    if (atoi(row[1]))
    {
      sprintf(query,
        "UPDATE `greylisting_data` "
        "SET `record_expires` = UTC_TIMESTAMP() + INTERVAL %u DAY, `passed_count` = `passed_count` + 1 "
        "WHERE `id` = '%s'",
        record_expire_good, row[0]);
      ret = RET_ACCEPT;
      gllog(LOGLEVEL_INFO, "greylisting: %s (%s -> %s) exists (id=%s) - accepting\n", relay_ip, mail_from, rcpt_to, row[0]);
    }
    else
    {
      sprintf(query,
        "UPDATE `greylisting_data` "
        "SET `blocked_count` = `blocked_count` + 1 "
        "WHERE `id` = '%s'",
        row[0]);
      ret = RET_TEMPREJECT;
      gllog(LOGLEVEL_INFO, "greylisting: %s (%s -> %s) is blocked (id=%s) - temp. rejecting\n", relay_ip, mail_from, rcpt_to, row[0]);
    }
  }
  else
  {
    sprintf(query,
      "INSERT INTO `greylisting_data` "
      "VALUES (0, '%s', '%s', '%s', UTC_TIMESTAMP() + INTERVAL %u MINUTE, UTC_TIMESTAMP() + INTERVAL %u MINUTE, 1, 0, 0, UTC_TIMESTAMP(), UTC_TIMESTAMP())",
      relay_ip, mail_from_esc, rcpt_to_esc, block_expire, record_expire);
    ret = RET_TEMPREJECT;
    gllog(LOGLEVEL_INFO, "greylisting: %s (%s -> %s) doesn't exist. - temp. rejecting\n", relay_ip, mail_from, rcpt_to);
  }
  mysql_free_result(res);

  if (mysql_query_wrapper(mysql, query))
  {
    gllog(LOGLEVEL_ERROR, "greylisting: mysql: %s\n", mysql_error(mysql));
    return RET_NOTFOUND;
  }

  free(mail_from_esc);
  free(rcpt_to_esc);
  return ret;
}

int main()
{
  int ret = 1;
  int greylisted = 0;
  MYSQL *mysql = NULL;

  /* load config */
  if (ret && !load_config())
    ret = 0;

  /* connect to mysql */
  if (ret)
  {
    mysql_library_init(-1, NULL, NULL);
    mysql = mysql_init(NULL);
    if (!mysql_real_connect(mysql, mysql_host, mysql_user, mysql_pass, mysql_db, mysql_port, NULL, 0))
    {
      gllog(LOGLEVEL_FATAL, "greylisting: mysql: %s\n", mysql_error(mysql));
      ret = 0;
    }
  }

  /* greylisting checks */
  if (ret && !greylisted)
  {
    greylisted = check_listed(mysql);
    if (greylisted == RET_NOTFOUND)
      greylisted = check_greylisted(mysql);
  }

  /* print smtp error code */
  if (ret)
  {
    switch(greylisted)
    {
      case RET_REJECT:
        printf(CMD_REJECT);
        break;
      case RET_TEMPREJECT:
        printf(CMD_TEMPREJECT);
        break;
    }
  }

  /* cleanup stuff */
  gllog(LOGLEVEL_DEBUG, "greylisting: exiting\n");
  if (mysql)
    mysql_close(mysql);
  mysql_library_end();
  cleanup();
  return !ret;
}