#!/usr/bin/python2

from __future__ import print_function

import sys, os, time

#import GeoIP
#gi = GeoIP.new(GeoIP.GEOIP_STANDARD)
#def country_code_by_addr(addr):
#    gi = GeoIP.new(GeoIP.GEOIP_STANDARD)

import maxminddb
mmdb_reader = maxminddb.Reader("/usr/share/GeoIP/GeoLite2-Country.mmdb")
def country_code_by_addr(addr):
    ret = mmdb_reader.get(addr)
    if ret and ("country" in ret):
      return ret["country"]["iso_code"]
    return ""

PROTECTED_COUNTRIES = [
  "SK", "CZ", "HU", "PL",
  "AU", "BE", "HR", "CY", "DK", "EE", "FI", "FR", "DE", "GR", "IE",
  "IT", "LV", "LT", "LU", "MT", "PT", "SI", "ES", "SE", "UK"
  # spammer EU
  #"BG", "NL", "RO"
]

class watcher:
  def __init__(self, filename):
      self.filename = filename
      self.reopen()
      #self.f.seek(0, 2) # seek to end of file
  def fstat(self):
      return os.stat(self.filename).st_ino
  def reopen(self):
      self.f = open(self.filename, "r")
      self.fd = self.f.fileno()
      self.stat = self.fstat()
  def readline(self):
      row = self.f.readline()
      if row:
        return row
      time.sleep(1)
      if self.fstat()!=self.stat:
        self.reopen()
      return ""
  def process(self):
      row = self.readline()
      if not row:
        return
      elif "SASL LOGIN authentication failed:" in row:
        ip = row.split("[", 2)[-1].split("]", 1)[0]
      else:
        return
      country = country_code_by_addr(ip)
      yield ip, country

class ipset_target:
  cmd = "ipset"
  def __init__(self, setname="blocker"):
      self.setname = setname
      print("Loading")
      self.load()
  def parse_row(self, row):
      a = row.split(" ")
      if ' comment "' in row:
        aa = row.split('"')
        country, date = row.split('"')[1].split(" ")
        return a[0], country, date
      # fallback to non country mode
      return a[0], "XX", "now"
  def load(self):
      self.data = {}
      s = os.popen("%s list %s" % (self.cmd, self.setname))
      # search for Members: string
      for i in range(100):
        row = s.readline().strip()
        if not row:
          break
        if row == "Members:":
          #print("members found")
          break
      # load current settings
      for i in range(1000):
        row = s.readline().strip()
        if not row:
          break
        ip, country, date = self.parse_row(row)
        print("IP loaded:", ip, country, date)
        self.data[ip] = dict(country=country, date=date)
  def add(self, ip, country=""):
      if ip in self.data:
        return # already added
      date = time.strftime("%Y%m%d-%H%M%S")
      print("fail_blocker: add %s comment '%s %s'" % (ip, country, date))
      os.popen(
        "%s add blocker %s comment '%s %s'"
        % (self.cmd, ip, country, date)
      ).read()
      self.data[ip] = dict(country=country, date=date)

class sudo_ipset_target(ipset_target):
  cmd = "sudo ipset"

class debug_target:
  def __init__(self, setname="blocker"):
      self.setname = setname
      self.data = {}
      print("Debug target")
  def add(self, ip, country=""):
      if ip in self.data:
        return # already added
      date = time.strftime("%Y%m%d-%H%M%S")
      print(
        "%s # comment %s %s"
        % (ip, country, date)
      )
      self.data[ip] = dict(country=country, date=date)

if len(sys.argv)>1:
  worker = watcher(sys.argv[1])
else:
  worker = watcher("/var/log/maillog")

if "debug" in sys.argv:
  target = debug_target()
else:
  target = ipset_target()
  print("loaded")

while True:
  for ip, country in worker.process():
    if country not in PROTECTED_COUNTRIES:
      target.add(ip, country)