import os
import shutil as sh
import pyclamd as cav
from pyinotify import WatchManager, Notifier, ThreadedNotifier, EventsCodes, ProcessEvent
 
#ประกาศ socket อันนี้จะแตกต่างไปตาม OS ในกรณีนี้เป็นของ Ubuntu 7.10
cav.init_unix_socket('/var/run/clamav/clamd.ctl')
 
wm = WatchManager()
#ตรวจสอบเฉพาะการสร้างและการแก้ไขไฟล์ ถ้าเปิดมาอ่านเฉยๆ ก็ปล่อยไป
mask = EventsCodes.IN_CREATE | EventsCodes.IN_MODIFY 
 
class PAfterAccess(ProcessEvent):
    def process_IN_CREATE(self, event):
        fname = os.path.join(event.path, event.name)
        self.do_scan(fname)
 
    def process_IN_MODIFY(self, event):
        fname = os.path.join(event.path, event.name)
        self.do_scan(fname)
 
    def do_scan(self, fname):
        res = None
        try:
            res = cav.scan_file(fname)
            print "Scanned %s result is %s" % (fname,str(res))
        except cav.ScanError:
            #ScanError มักพบในกรณีที่ไฟล์นั้นหายไปแล้ว เช่นกรณีไฟล์ Swap
            print "Can't scan file %s" % fname
        #ฟังก์ชั่น scan_file จะคืนค่าเป็น None ถ้าไม่เจอไวรัส
        #ในกรณีนี้ถ้าเจอไวรัสก็ให้ย้ายไฟล์ไปเก็บที่อื่น
        if res != None:
            hname = fname[1:].replace("/","-")
            sh.move(fname,"/home2/quarantine/"+hname)
 
 
notifier = Notifier(wm, PAfterAccess())
#ระบุตำแหน่งที่ต้องการแสกน ควรตั้ง auto_add=True ไว้ด้วยเผื่อการสร้าง folder ใหม่
wdd = wm.add_watch('/home', mask, rec=True,auto_add=True)
 
#ที่เหลือยกมาจาก Tutorial
while True:  # loop forever
    try:
        # process the queue of events as explained above
        notifier.process_events()
        if notifier.check_events():
            # read notified events and enqeue them
            notifier.read_events()
        # you can do some tasks here...
    except KeyboardInterrupt:
        # destroy the inotify's instance on this interrupt (stop monitoring)
        notifier.stop()
        break