python脚本 限制 外部访问 linux服务器端口

注意：该脚本会清空linux防火墙的filter表的规则和用户自定义链路

脚本的效果是将端口限制为仅服务器内部访问 

可以提供ip地址白名单

具体脚本：

#!/usr/bin/python3
import argparse, subprocess, sys, re


def popen(cmd):
    global result
    try:
        result = subprocess.run(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, encoding='utf-8', check=True)
    except Exception as e:
        raise Exception(e.stderr)
    file_count = result.stdout.strip()  # 假设grep -c 的输出只有一行且没有额外的空格或换行符
    return file_count


def assignment():
    parser = argparse.ArgumentParser(description='manual to this script')
    parser.add_argument("--ip_white_list", type=str, default="")
    parser.add_argument("--port_list", type=str, default="")
    args = parser.parse_args()

    ip_white_list = args.ip_white_list
    port_list = args.port_list

    if not port_list :
        raise Exception("--port_list不能为空")

    ip_white_list = f'{ip_white_list},'.split(",")
    port_list = f'{port_list},'.split(",")

    ip_white_list = set(filter(lambda x:x!='',ip_white_list))
    port_list = set(filter(lambda x:x!='',port_list))

    for port in port_list:
        if re.search(r"[^0-9]", port):
            raise Exception("--port_list 不是合法的端口")
        pass

    if ip_white_list :
        for ip_white in ip_white_list:
            if re.search(r"[^0-9.]", ip_white):
                raise Exception("--ip_white_list 不是合法的ip地址")
            pass

    return ip_white_list, port_list


def clear_iptables_rule():
    popen("iptables -F ")
    popen("iptables -X ")
    pass


def ip_white(ip_white_list, port_list):
    popen("iptables -t filter -P INPUT ACCEPT")
    popen("iptables -t filter -P OUTPUT ACCEPT")
    popen("iptables -A INPUT -p tcp -s 127.0.0.1 -j ACCEPT")
    popen("iptables -A INPUT -p tcp -s $(ip r g  1|cut -d ' ' -f 7) -j ACCEPT")
    if ip_white_list:
        for ip in ip_white_list:
            for port in port_list:
                popen(f"iptables -A INPUT -p tcp -s {ip} --dport {port} -j ACCEPT")


def port_disables(port_list):
    if port_list:
        for port in port_list:
            popen(f"iptables -A INPUT -p tcp --dport {port} -j REJECT ")
    pass


def clear_script():
    popen(f"echo '' > {sys.argv[0]}")
    pass


if __name__ == "__main__":
    ip_white_list, port_list = assignment()
    clear_iptables_rule()
    ip_white(ip_white_list, port_list)
    port_disables(port_list)
    clear_script()
    pass

执行脚本示例选项

./ta  --port_list 8081,8080 --ip_white_list 192.168.0.1,120.0.0.2