【Pyhon】获取文件MIME类型，根据文件类型自定义文件后缀

场景

下载样本，都是MD5命名的无后缀文件，需要自己手动查询然后修改文件后缀。

根据文件类型自定义后缀可以很方便地根据后缀判断用什么工具分析。

使用说明

libmagic

地址：https://pypi.org/project/python-magic-bin/0.4.14/

根据系统版本选择安装程序

https://pypi.org/project/python-magic-bin/0.4.14/#files

如果是windows版本，下载python_magic_bin-0.4.14-py2.py3-none-win32.whl

安装

pip install python_magic_bin-0.4.14-py2.py3-none-win32.whl

运行结果

- 自定义后缀前

D:\test>01Malware_identification_extension.py
[!!! need extension define !!!]libmagic.dll     ==>     PE32 executable (DLL) (GUI) Intel 80386, for MS Windows
[!!! need extension define !!!]magic.mgc        ==>     magic binary file for file(1) cmd (version 14) (little endian)
file count: 2

- 自定义后缀后

D:\test>01Malware_identification_extension.py
libmagic.dll    ==>     PE32 executable (DLL) (GUI) Intel 80386, for MS Windows
[!!! need extension define !!!]magic.mgc        ==>     magic binary file for file(1) cmd (version 14) (little endian)
file count: 2

其中need extension define是未指定后缀格式所以显示的提示，可以通过修改config.ini来自定义要改名的后缀。然后自动改文件名。

代码实现

#!/usr/bin/python
# -*- coding: UTF-8 -*-

# 第三方库 https://pypi.org/project/python-magic-bin/0.4.14/
import os
import magic
import sys
reload(sys)
sys.setdefaultencoding('utf8')   

# 识别格式的字典
identification_dict = {}

# 读取文件函数
def read_file(file_path):
    if not os.path.exists(file_path):
        print 'Please confirm correct filepath !'
        sys.exit(0)
    else:
        with open(file_path, 'r') as source:
            for line in source:
                line = line.rstrip('\r\n').rstrip('\n')
                key =   line.split(':')[0]
                value = line.split(':')[1]
                identification_dict[key] = value

                #ip_list.append(line.rstrip('\r\n').rstrip('\n'))

# 后缀增加
def ident_extension(var):
    # 自定义增加后缀方法
    # '文件描述':'.后缀',
    return identification_dict.get(var,False) #'error'为默认返回值，可自设置

# 增加文件后缀名
def change_filename(filepath,filemime):
    # 如果原来就有后缀，就不修改了
    # portion = os.path.splitext(filepath)
    # if len(portion[1])>0:
    #     return False

    # 如果后缀在数据库中就改名
    if ident_extension(filemime) != False :
        portion = filepath + ident_extension(filemime)
        os.rename(filepath,portion)
        return True
    else:
        return False

if __name__ == '__main__':
    # 读取配置文件存放进字典
    read_file(os.getcwd()+'//config.ini')
    # 存放的路径
    path = os.getcwd()+'//sample'
    # 判断目录是否存在
    if os.path.exists(path) == False:
          print u'usage:需要把无后缀名的样本放入sample目录内，与脚本放置在一起，config.ini是根据文件描述自定义后缀的配置文件'
          print u'C:.'
          print u'└─identification_extension.py'
          print u'│'
          print u'└─sample'
    # 遍历目录内文件名
    nCout = 0                                                       # 文件计数
    #path_text = open('path.txt', 'w')                               # 将文件名写入文本
    for fpathe, dirs, fs in os.walk(path):                          # 遍历路径
        for f in fs:                                                # 遍历文件名
            nCout = nCout + 1                                       # 计数器增加
            # 识别单个后缀函数
            mimetype = magic.from_file(os.path.join(fpathe, f))
            if change_filename(os.path.join(fpathe, f),mimetype):       # 修 改文件后缀
                print f + '\t==>\t' + mimetype
            else:
                print  '[!!! need extension define !!!]' + f + '\t==>\t' + mimetype
    print 'file count: {0}'.format(nCout)                           # 打印计数器

效果