Python 通过模块函数调用扫描MD5值查找病毒文件示例

'''
Python 通过模块函数调用扫描MD5值查找病毒文件示例
f.py
'''
import os

def fib(n=1000):
    """Print a Fibonacci series up to n."""
    a,b = 0,1
    while a < n:
        print(a,end=',')
        a,b = b,a+b


def prime(x=100):
    """Print a prime number series up to x."""
    for n in range(2, x):
        for x in range(2, n):
            if n % x == 0:
                #print(n, 'equals', x, '*', n//x)
                break
        else:
        # loop fell through without finding a factor
             print(n, 'is a prime number')


def ask_ok(prompt, retries=4, reminder='Please try again!'):
    while True:
        ok = input(prompt)
        if ok in ('y', 'ye', 'yes'):
            return True
        if ok in ('n', 'no', 'nop', 'nope'):
            return False
        retries = retries - 1
        if retries < 0:
            raise ValueError('invalid user response')
        print(reminder)

def getmd5(file):
    import hashlib
    m = hashlib.md5()
    with open(file,'rb') as f:
        for line in f:
            m.update(line)
    md5code = m.hexdigest()
    print(md5code)
    return md5code


def getdisklist():
    import string    
    disklist = []
    d = string.ascii_uppercase
    #print(d)
    for w in d:        
        disk = w+':'        
        if os.path.isdir(disk):            
            disklist.append(disk)
    print(disklist)
    return disklist


def scan(disklist):
    #print(disklist)
    for disk in disklist:
        #print(disk)
        os.chdir(disk+'/')
        tree = os.walk('/')
        for dir in tree:
            for file in dir[2]:
                if '.jpg'in file or '.pdf'in file or '.com'in file or '.exe'in file or '.dll'in file:
                    myfile = disk+dir[0]+'/'+file
                    print(myfile)
                    mymd5code = getmd5(myfile)
                    with open ('md5.txt','a') as f:
                        f.write(myfile+'\n'+mymd5code+'\n')
                    print('md5: ',mymd5code)
                    if mymd5code == 'fef5c779d0b44382ef8f073ba0bbf7bb':
                        print('找到病毒文件：Office 2010 Toolkit')
                        return 

'''
Python 通过模块函数调用扫描MD5值查找病毒文件示例
scan.py
'''
import f

while True:
    print('''

    1、生成fib数列
    2、生成素数列表
    3、询问
    4、获取MD5
    5、获取磁盘列表
    6、扫描磁盘文件
    ''')
    w = input('请选择（1/2/3/4）：')
    if w == '1':
        f.fib(10000)
    elif w == '2':
        f.prime(1000)
    elif w == '3':
        f.ask_ok('OK?')
    elif w == '4':
        f.getmd5(r'./Office 2010 Toolkit.exe')
    elif w == '5':
        f.getdisklist()
    elif w == '6':
        disklist = f.getdisklist()
        f.scan(disklist)
    else:
        print('输入错误')