[Python網絡編程]淺析守護進程後臺任務的設計與實現

    在做基於B/S應用中,經常有需要後臺運行任務的需求,最簡單比如發送郵件。在一些如防火牆,WAF等項目中,前臺只是爲了展示內容與各種參數配置,後臺守護進程纔是重頭戲。所以在防火牆配置頁面中可能會經常看到調用cgi,但真正做事的一般並不是cgi,比如說執行關機命令,他們的邏輯如下:


   (ps:上圖所說的前臺界面包含通常web開發中的後端,不然也沒有socket一說)


    爲什麼要這麼設計

你可能疑惑爲什麼要這麼設計,我覺得理由如下:
首先有一點說明,像防火牆等基本上都運行在類Linux平臺上
    1.安全問題  cgi一般也就擁有www權限,但執行關鍵等命令需要root,所以需要讓後臺守護進程去幹
    2.一般類似防火牆的後臺守護進程是C/C++寫的,在消息格式上很方便處理,如填充一個消息結構體發送出去,後臺進程只需要強制轉換爲定義的結構體,就輕鬆獲得傳遞的參數值。

那可不可以去掉中間的cig模塊,直接發送消息給後臺守護進程呢?
我覺得是可以的,本文的重點也是實現這個方案。

如何實現

由於最近一直在windows下,所以我們的守護進程是運行在windows下的,但其實windows並沒有守護進程的概念,相對應的是服務的概念。這裏需要安裝pywin32包。
class MgrService(win32serviceutil.ServiceFramework): 
    """
    Usage: 'python topmgr.py install|remove|start|stop|restart'
    """
    #服務名
    _svc_name_ = "Mgr"
    #服務顯示名稱
    _svc_display_name_ = "Daemon Mgr"
    #服務描述
    _svc_description_ = "Daemon Mgr"

    def __init__(self, args): 
        win32serviceutil.ServiceFramework.__init__(self, args) 
        self.hWaitStop = win32event.CreateEvent(None, 0, 0, None)

    def SvcDoRun(self):
        self.ReportServiceStatus(win32service.SERVICE_START_PENDING)
        INFO("mgr startting...")
        self.ReportServiceStatus(win32service.SERVICE_RUNNING)
        self.start()
        # 等待服務被停止
        INFO("mgr waitting...")
        win32event.WaitForSingleObject(self.hWaitStop, win32event.INFINITE)
        INFO("mgr end")
        
    def SvcStop(self): 
        self.ReportServiceStatus(win32service.SERVICE_STOP_PENDING)
        INFO("mgr stopping...")
        self.stop()
        INFO("mgr stopped")
        # 設置事件
        win32event.SetEvent(self.hWaitStop)
        self.ReportServiceStatus(win32service.SERVICE_STOPPED)

    def start(self): pass

    def stop(self): pass
很簡單,這樣就實現了windows中的服務,也就是說脫離終端,運行於後臺。INFO等函數只是簡單的記錄作用,可直接忽略。
我們要實現自己的後臺程序,只需要繼承MgrService,並提供start,stop方法就可以了。

由於我們是通過socket來傳遞消息的,所以在start方法中要監聽端口,等待連接,處理連接,這個大家都很擅長。在這裏我選擇了
 單線程,基於協程,底層使用libev(libevent)--- gevent這個高性能網絡庫。對gevent有興趣的童鞋可以看看深度分析gevent運行流程
class Engine(MgrService):
    rbufsize = -1
    wbufsize = 0

    def start(self):
        INFO('wait connection')
        self.server = StreamServer((HOST, PORT), self.msg_handle)
        self.server.serve_forever()

    def msg_handle(self,socket,address):
        try:
            rfile = socket.makefile('rb', self.rbufsize)
            wfile = socket.makefile('wb', self.wbufsize)
            headers = Message(rfile).dict

            INFO('get a connection from:%s,headers:%s' % (str(address), headers))

            if 'module' in headers and headers['module'] in MODULES:
                MODULES[headers['module']].handle(wfile, headers)
        except Exception:
            ERROR('msg_handle exception,please check')

    def stop(self):
        if hasattr(self, server):
            self.server.stop()
當有新連接到來,由msg_handle處理,首先讀取發送來的消息,消息格式使用了最簡單的http的格式,即(鍵名:鍵值)的格式,你要問我爲什麼採用這個格式,哈哈,格式簡單,python有現成的庫解析。

考慮到後期模塊可能很多,所以我們的處理流程自動根據消息的模塊參數,調用對應模塊的handle方法。
上面代碼的那個MODULES是個全局變量,當你添加一個模塊的時候需要註冊到MODULES中,我提供了module_register方法。
MODULES = {           # module: handle module class
}

def module_register(module_name, handle_class):
    if module_name in MODULES:
        WARN('duplicate module_name:' + module_name)
    else:
        MODULES[module_name] = handle_class

到這裏一切都很自然,但貌似只假設模塊有handle方法,自己寫一個模塊還是很費事,你需要自己去想怎麼調用,最有返回什麼格式的數據,這都是一件頭疼的事情,所以最好提供一個基類模塊。
class Module(object):
    SECRE_KEY = "YI-LUO-KEHAN"
    MODULE_NAME = "BASE_MODULE"
    PREFIX = "do_"  # method prefix

    def __init__(self, wfile, headers):
        self.wfile = wfile
        self.headers = headers

    def __getattr__(self, name):
        try:
            return self.headers[name]
        except Exception:
            ERROR("%s has no attr:%s,please check" %(self.MODULE_NAME, name))            

    @classmethod
    def handle(cls, wfile, headers):
        module_obj = cls(wfile, headers)
        module_obj.schedule_default()

    def verify(self):
        if hmac.new(self.SECRE_KEY, self.MODULE_NAME).hexdigest() == self.signature:
            return True
        else:
            WARN("client verify failed,signature:%s" % str(self.signature))

    def schedule_default(self):
        err_code = 0
        if self.verify() and self.action:
            func_name = self.PREFIX + self.action
            try:
                getattr(self, func_name)()
            except AttributeError:
                err_code = 1
                ERROR("%s has no method:%s" %(self.MODULE_NAME, func_name))
            except Exception:
                err_code = 2
                ERROR("module:%s,method:%s,exception" % (self.MODULE_NAME, func_name))              
        else:
            err_code = 3

        if err_code:
            self.send_error({'err_code':err_code})

    def send_success(self, msg=''):
        data = {'success':True,'msg':msg}
        self.wfile.write(json.dumps(data))

    def send_error(self, msg=''):
        data = {'success':False,'msg':msg}
        self.wfile.write(json.dumps(data))

在基類模塊中我們提供了默認的處理流程,即根據消息中action,調用do_action方法,並提供了一個簡單但很有效的認證方法,通過消息的signature字段,可能有些簡陋,但沒關係,你可以定義自己的認證方法。

下面該寫我們自己的模塊了,
TASK = {}  # task_id: pid
class ScanModule(Module):
    MODULE_NAME = "SCAN_MODULE"

    def do_start(self):
        self.send_success('start ok')
        DEBUG('------------task start------------')
        task_ids = [int(task_id) for task_id in self.task_ids.split(',') if int(task_id) not in TASK]

        for task_id in task_ids:
            try:
                cmd = 'python scan.py -t %s' % task_id
                DEBUG(cmd)
                self.sub = Popen(cmd, shell=True, cwd=CWD)
                pid = int(self.sub.pid)
                TASK[task_id] = pid
                INFO('%s start a new task,task_id:%s,pid:%s' %(self.MODULE_NAME, task_id, pid))
            except Exception:
                ERROR('%s start a new task,task_id:%s failed' % (self.MODULE_NAME, task_id))

    def do_stop(self):
        self.send_success('stop ok')
        DEBUG('------------task stop------------')
        task_ids = [int(task_id) for task_id in self.task_ids.split(',') if int(task_id) in TASK]

        for task_id in task_ids:
            pid = TASK.pop(task_id)
            try:
                INFO('%s stop a new task,task_id:%s,pid:%s' %(self.MODULE_NAME, task_id, pid))
                call(['taskkill', '/F', '/T', '/PID', str(pid)])
            except Exception:
                ERROR('%s taskkill a task failed,task_id:%s,pid:%s' %(self.MODULE_NAME, task_id, pid))


module_register(ScanModule.MODULE_NAME, ScanModule)
上面實現了一個簡單的掃描模塊,支持兩個action,start,stop。
start很簡單,調用gevent的subprocess.Popen運行子進程,並記錄pid,stop則使用taskkill直接殺掉該進程。
這裏有兩點需要注意:
    1.不要用原生的subprocess模塊,因爲原生的subprocess是阻塞的,這可能導致主處理邏輯也阻塞,不能服務更多的請求
最後別忘了調用module_register註冊相應模塊。
    2.方法一開始最好就返回結果,因爲前臺很可能在等待返回。所以說as soon as possible

下面提供一個客戶端用於測試,client.py
#!/usr/bin/env python
#-*-encoding:UTF-8-*-

import hmac
import gevent
from gevent import monkey
monkey.patch_socket()

addr = ('localhost', 6667)


def send_request(module_name,request_headers):
    SECRE_KEY = "YI-LUO-KEHAN"
    socket = gevent.socket.socket()
    socket.connect(addr)
    request_headers['module'] = module_name
    request_headers['signature'] = hmac.new(SECRE_KEY, module_name).hexdigest()
    h = ["%s:%s" %(k, v) for k,v in request_headers.iteritems()]
    h.append('\n')
    request = '\n'.join(h)
    socket.send(request)
    print socket.recv(8192)
    socket.close()

if __name__ =="__main__":
    import sys
    if sys.argv[1] == 'start':
        send_request('SCAN_MODULE',{'action':'start','task_ids':'1'})
    else:
        send_request('SCAN_MODULE',{'action':'stop','task_ids':'1'})

    
    

我們來簡單的測試一下:
注意:由於要註冊到服務,cmd需要管理員權限
至於start中調用的scan.py隨便寫一個就可以

截圖如下,我們看到成功!!!


本文代碼已放到github,https://github.com/Skycrab/pymgr
感興趣的童鞋可以參考,請大家多提意見。




發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章