python ipset管理 增删白名单

为方便用ipset 来管理防火墙，写了下面Ipset类来对Ip进行管理

#!/usr/bin/env python
# coding: utf-8

import MySQLdb
import MySQLdb.cursors
import subprocess
import logging
import re,os
import xml.sax

class XMLHandler(xml.sax.ContentHandler):
    '''
    用来解析ipset数据
    '''
    def __init__(self):
        self.current_tag = ""
        self.current_set = ""
        self.mapping = {}

    def startElement(self, name, attributes):
        self.current_tag = name
        if name == 'ipset':
            self.current_set = attributes['name']
            self.mapping[self.current_set] = []

    def characters(self, data):
        if self.current_tag == 'elem' and data!='\n':
            self.mapping[self.current_set].append(data)

    def endElement(self, name):
        if name == 'ipset':
            self.current_set = ''

    def getData(self):
        return self.mapping

class IpsetPool(object):
    def __init__(self):
        self.msg = []
        self.config = '/etc/sysconfig/ipset'
        self.logger_file = '/data/logs/ipset.log'
        self.ipsets = ['manage','center','project']
        self.log = self.mylog()
        self.ipset_data = self.getIpsetData(XMLHandler)
        
    def sub_call(self,run_cmd, **kwargs):
        p = subprocess.Popen(
            run_cmd,
            shell=True,
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            **kwargs)
        outdata, errdata = p.communicate()
        retcode = p.wait()
        self.msg.append((False,errdata) if retcode != 0 else (True,outdata))
        return retcode, outdata, errdata

    @staticmethod
    def querydb(sql):
        host = "192.168.59.128"
        user = "dev"
        passwd = "123456"
        db = "gmweb_res"
        conn = MySQLdb.connect(
            host=host,
            user=user,
            passwd=passwd,
            db=db,
            charset="utf8",
            cursorclass=MySQLdb.cursors.DictCursor)
        cursor = conn.cursor()
        cursor.execute(sql)
        rs = cursor.fetchall()
        cursor.close()
        conn.commit()
        conn.close()
        return rs

    def getManageIps(self):
        sql = "select * from host where state not in (6) and `use` REGEXP ',2$|^2,|^2$|,2,';"
        return [x["ip1"] for x in self.querydb(sql)]

    def getProjectIps(self):
        from jgconf.models import projectConf
        return [i['saltIp'] for item in projectConf.objects.all() for i in item.getSaltServer()]

    def checkAddrIsIn(self, ip, setname):
        '''
        判断ip是否在某个set中
        '''
        if ip in self.ipset_data[setname]:
            return True
        else:
            return False

    def getIpsetData(self,xml_handler):
        '''
        获取机器上当前的ipset配置数据
        '''
        xh = xml_handler()
        xml.sax.parseString(self.sub_call('ipset list -o xml')[1], xh)
        return xh.getData()
    def createSet(self,setname):
        self.log.info('create {0} set'.format(setname))
        set_cmd = 'ipset create {0} hash:ip'.format(setname)
        return self.sub_call(set_cmd)

    def renderSetFile(self):
        '''
        重导配置
        '''
        self.sub_call('ipset save > {0}'.format(self.config))

    def createIpsets(self):
        for ipset in self.ipsets:
            self.createSet(ipset)

    def addIps2Set(self,setname,ips):
        if setname not in self.ipsets:
            self.log.error('invalid set name!')
            return False
        if not self.ipset_data.has_key(setname):
            self.createSet(setname)
        for ip in ips:
            if not self.checkAddrIsIn(ip,setname):
                self.log.info('add {0} {1}'.format(setname,ip))
                self.sub_call('ipset -A {0} {1}'.format(setname,ip))
        self.renderSetFile()
    def delIpsFromSet(self,setname,ips):
        self.log.info(ips)
        if setname in self.ipsets and self.ipset_data.has_key(setname):
            for ip in ips:
                if self.checkAddrIsIn(ip,setname):
                    self.log.info('delete {0} {1}'.format(setname,ip))
                    self.sub_call('ipset -D {0} {1}'.format(setname,ip))
        self.renderSetFile()
    def mylog(self):
        logger_dir = os.path.split(self.logger_file)[0]
        if not os.path.exists(logger_dir):
            os.makedirs(logger_dir)
        logger = logging.getLogger("reload")
        logger.setLevel(logging.DEBUG)
        # create file handler which logs even debug messages
        fh = logging.FileHandler(self.logger_file)
        fh.setLevel(logging.DEBUG)
        # create formatter and add it to the handlers
        formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
        fh.setFormatter(formatter)
        console = logging.StreamHandler()
        console.setLevel(logging.DEBUG)
        # add the handlers to the logger
        logger.addHandler(fh)
        logger.addHandler(console)
        return logger
    def reloadIpset(self):
        """
        从文件中加载最新集合
        """
        reloadlog = mylog()
        try:
            # 刷新清空当前规则
            sub_call("/etc/init.d/iptables stop")
            sub_call("/etc/init.d/ipset restart")
            sub_call("/etc/init.d/iptables start")
            reloadlog.info("reload成功")
        except Exception as e:
            reloadlog.info("ipset reload异常 %s" % e)
    def loadDefault(self):
        #self.addIps2Set('manage',self.getManageIps())
        self.addIps2Set('project',self.getProjectIps())

if __name__ == '__main__':
    p = IpsetPool()
    p.loadDefault()