#!/data/setupfiles/python3/bin/python3

import re
from datetime import date
import os
import subprocess
import argparse

try:
  #登录错误尝试次数
  failedtimes = 3
  #日志路径
  logfile = r'/data/setupfiles/mysql/mysql-error.log'
  #ipset黑名单名称
  blackname='mysql_black'
  ips="";
  def search_source():
    t = date.today()
    month = t.strftime('%b')
    day = t.strftime('%d')
    pat = re.compile('.+Access denied for user.+''(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})''.+')
    lines = []
    f = open(logfile,'r')
    for line in f:
      if len(line) < 10:
        continue
      #print(line)
      if re.search(pat,line):
        lines.append(line)
    return lines

  def count_ips(lines):
    count = {}
    if len(lines) == 0:
      return count;
    pat = re.compile('''(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})''')
    for line in lines:
      ip = re.findall(pat,line)[0]
      if ip in count:
        count[ip] += 1
      else:
        count[ip] = 1
    return count


  #设置ipset防火墙
  def deny_ipset(count):
    global ips;
    for ip in count:
      if count[ip] >= failedtimes:
        rs=subprocess.getoutput("sudo ipset list "+blackname+"");
        if ip not in rs:
          rs=subprocess.getoutput("sudo ipset add "+blackname+" "+ip);
          if rs=="":
            ips+='{"'+ip+'":'+str(count[ip])+'},';
	 
  def main():
    lines = search_source()
    #print(lines);
    count = count_ips(lines)
    #print(count);
    deny_ipset(count)
    if ips=="":
      print("no");
    else:
      print(ips)
  if  __name__=="__main__":
    #声明参数
    parser = argparse.ArgumentParser(description='检查登录失败次数\n命令：python secure.py --failedtimes=6');
    parser.add_argument('--failedtimes', type=int, default=3, help='登录失败次数,默认值:6')
    args = parser.parse_args();
    failedtimes=int(args.failedtimes);
    #
    main()
except Exception as e:
  print("no");