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

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

try:
  #登录错误尝试次数
  whiteip='123.207.175.190,'
  failedtimes = 3
  ips="";
  logfile = r'/var/log/secure'
  denyfile = r'/etc/hosts.deny'
  months_31 = ['Jan','Mar','May','Jul','Aug','Oct','Dec']
  months_30 = ['Apr','Jun','Sep','Nov']
  month_28or29 = 'Feb'
  months = {
          'Jan':1,'Feb':2,'Mar':3,'Apr':4,'May':5,'Jun':6,
          'Jul':7,'Aug':8,'Sep':9,'Oct':10,'Nov':11,'Dec':12
         }
  month_days = {}
  for mon in months_31:
    month_days[mon] = 31
  for mon in months_30:
    month_days[mon] = 30
  if date.isocalendar(date.today())[0] % 4 == 0:
    month_days[month_28or29] = 29
  else:
    month_days[month_28or29] = 28

  def search_source():
    t = date.today()
    month = t.strftime('%b')
    day = t.strftime('%d')
    pat = re.compile('.+sshd.+Failed password.+ (\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
      if line.split()[0] == month and (int(day) - int(line.split()[1])) < 7 and (int(day) - int(line.split()[1])) >= 0:
        if re.search(pat,line):
          lines.append(line)
        elif (months[month] - months[line.split()[0]]) == 1 or (months[month] - months[line.split()[0]]) == -11:
          if (int(day) + month_days[line.split()[0]] - int(line.split()[1])) < 7 and 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

  def deny_ips(count):
    f = open(denyfile,'w')
    for ip in count:
      if count[ip] >= failedtimes:
        word = 'ALL: %s #failed %d times in a week.\n' % (ip,count[ip])
        f.write(word)
    f.close()

  #设置ipest防火墙
  def deny_ipset(count):
    global ips;
    for ip in count:
      if count[ip] >= failedtimes and ip not in whiteip:
        #os.system("ipset -exist add ssh_black "+ip);
        rs=subprocess.getoutput("sudo ipset list ssh_black");
        if ip not in rs:
          rs=subprocess.getoutput("sudo ipset add ssh_black "+ip);
          if rs=="":
            ips+='{"'+ip+'":'+str(count[ip])+'},';
	 
  def main():
    lines = search_source()
    #print(lines);
    count = count_ips(lines)
    #print(count);
    deny_ipset(count)
    #save forever
    #os.system("service ipset save");
    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");