#!/usr/bin/env python3

import subprocess
import sys
import time
import shutil
import os


def getAllLink(cmds):
    cmd_ret = subprocess.Popen(cmds, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    while cmd_ret.poll() == None:
        time.sleep(1)
    if cmd_ret.poll() != 0:
        raise AttributeError("Failed to execute cmds %s,reason:\n%s" % (cmds, "\n".join(cmd_ret.stderr)))

    line = str(cmd_ret.stdout.readline(), encoding="utf-8")
    while True:
        if line:
            yield line
        else:
            break
        line = str(cmd_ret.stdout.readline(), encoding="utf-8")

def statisticalData(paras):
    ret_dict = {}
    for ele in paras:
        tmp_data = ele.split()
        tmp_ip = tmp_data[4].split(":")[0]
        tmp_rq = int(tmp_data[1])
        tmp_sq = int(tmp_data[2])
        tmp_state = tmp_data[5]

        ret_dict.setdefault(tmp_ip, {})
        ret_dict[tmp_ip].setdefault("recvQ", 0)
        ret_dict[tmp_ip].setdefault("sendQ", 0)
        ret_dict[tmp_ip].setdefault(tmp_state, 0)

        ret_dict[tmp_ip]["recvQ"] += tmp_rq
        ret_dict[tmp_ip]["sendQ"] += tmp_sq
        ret_dict[tmp_ip][tmp_state] += 1

    return ret_dict

def judgeData(paras):
    ret_list = []
    for k, v in paras.items():
        recvQueue_NUM = v.get('recvQ', 0)
        sendQueue_NUM = v.get('sendQ', 0)

        established_NUM = v.get("ESTABLISHED", 0)
        closed_NUM = v.get("CLOSED", 0)

        # closing_NUM = v.get("CLOSING", 0)
        # closewait_NUM = v.get("CLOSE_WAIT", 0)
        # unknown_NUM = v.get("UNKNOWN", 0)
        # synsent_NUM = v.get("UNKNOWN", 0)
        # synrecv_NUM = v.get("UNKNOWN", 0)

        warn_num = 0
        for k1, v1 in v.items():
            if k1 not in ["recvQ", "sendQ", "ESTABLISHED", "CLOSED"]:
                warn_num += v1
        ok_num = established_NUM + closed_NUM

        if ok_num > 100:
            ret_list.append(k)
        elif ok_num == 0 and warn_num > 100:
            ret_list.append(k)
        elif warn_num > 300:
            ret_list.append(k)

    return ret_list

def recordsData(paras, rfile, wfile):
    denyLevel = [1, 3, 7, 30, 365, 999*365]
    with open(rfile, "r") as rfobj, open(wfile, "w") as wfobj:
        #rline = str(rfobj.readline(), encoding="utf-8")
        rline = rfobj.readline()
        while rline:
            tline = rline.split()
            nowTime = int(time.time())
            if tline[0] in paras:
                tline[1] += 1
                tline[2] = nowTime
                tline[3] = nowTime + denyLevel[int(tline[1])]*3600*24
                paras.remove(tline[0])
            tline = " ".join(tline)
            tline = "%s\n" % tline
            wfobj.write(tline)
            #rline = str(rfobj.readline(), encoding="utf-8")
            rline = rfobj.readline()
        for ele in paras:
            nowTime = int(time.time())
            line = "%s 0 %d %d\n" %(ele, nowTime, nowTime + 3600*24)
            #line = bytes(line, encoding="utf-8")
            wfobj.write(line)

    shutil.move(wfile, rfile)
    return True

def rewriteIpsFile(rfile, tmpipsfile):
    prev_text = """# Generated by iptables-save v1.4.7 on Thu May 31 20:23:03 2018
*filter
:INPUT DROP [0:0]
:FORWARD DROP [0:0]
:OUTPUT ACCEPT [60:4296]
-A INPUT -s 192.168.0.0/16 -j ACCEPT
-A INPUT -s 127.0.0.1/32 -j ACCEPT
-A INPUT -s 103.212.33.132/32 -j ACCEPT
-A INPUT -s 221.120.162.0/24 -j ACCEPT
-A INPUT -m conntrack --ctstate RELATED,ESTABLISHED -j ACCEPT
"""
    suffix_text = """-A INPUT -p icmp -j REJECT --reject-with icmp-host-prohibited
-A INPUT -p tcp -m tcp --dport 80 -j ACCEPT
-A INPUT -p tcp -m tcp --dport 443 -j ACCEPT
-A INPUT -p udp -m udp --dport 80 -j ACCEPT
-A INPUT -p udp -m udp --dport 443 -j ACCEPT
COMMIT
# Completed on Thu May 31 20:23:03 2018
"""
    with open(rfile, "r") as rfobj, open(tmpipsfile, "w") as wfobj:
        wfobj.write(prev_text)
        #rline = str(rfobj.readline(), encoding="utf-8")
        rline = rfobj.readline()
        while rline:
            rline = rline.split()
            nowTime = int(time.time())
            if nowTime < int(rline[3]):
                line = "-A INPUT -s %s/32 -j DROP\n" % rline[0]
                wfobj.write(line)
            rline = rfobj.readline()
        wfobj.write(suffix_text)
    return True

def restartIps(tmpfile, actualfile):
    fmode = os.stat(actualfile).st_mode
    fuid = os.stat(actualfile).st_uid
    fgid = os.stat(actualfile).st_gid

    shutil.move(tmpfile, actualfile)
    os.chmod(actualfile, fmode)
    os.chown(actualfile, uid=fuid, gid=fgid)
    cmd = "/etc/init.d/iptables restart"
    cmd_ret = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    while cmd_ret.poll() == None:
        time.sleep(1)
    ret_code = cmd_ret.poll()
    if ret_code == 0:
        return True
    return False

def main():
    cmd = "netstat -ntlupa |grep \"202.181.24.107:80\""
    rfile = "/root/ips.txt"
    tmp_wfile = "/tmp/tmp_ips.txt"
    ips_file = "/etc/sysconfig/iptables"
    tmp_ipsfile = "/tmp/tmp_ips"
    try:
        getret = getAllLink(cmd)
        stsret = statisticalData(getret)
        if not len(stsret):
            return True
        judret = judgeData(stsret)
        recret = recordsData(judret, rfile, tmp_wfile)
        if not recret:
            raise IOError("Failed to write data in %s" % tmp_wfile)
        rwret = rewriteIpsFile(rfile, tmp_ipsfile)
        if not rwret:
            raise IOError("Failed to rewrite file %s" % tmp_ipsfile)
        resrvret = restartIps(tmp_ipsfile, ips_file)
        if not resrvret:
            raise IOError("Failed to restart service iptables.")
    except AttributeError:
        return True
    except :
        return False

    return True


if __name__ == "__main__":
    while True:
        RET = main()
        if not RET:
            break
        time.sleep(5)

    sys.exit(0) if RET else sys.exit(1)