【Python】自动化文件类检测与重命名脚本
脚本概述
这段Python脚本的核心功能是根据文件头字节和PE文件的特征自动判断文件类型，并为没有扩展名的文件添加适当的扩展名。它支持多种常见文件类型，包括压缩文件（如 .rar、.zip、.7z）、图像文件（如 .jpg、.png）、音频/视频文件（如 .mp3、.mp4）等。特别地，对于PE文件（可执行文件和动态链接库），脚本能够根据文件的特征自动判定其为 .exe 或 .dll 文件，并将文件重命名。
主要功能和步骤
文件类型检测：
脚本通过一个字典FILE_TYPES映射常见的文件头字节（magic bytes）到文件扩展名。通过读取文件的前几个字节，脚本能够匹配这些字节序列，进而确定文件的类型。例如：52617221代表.rar文件504B0304代表.zip文件FFD8FFE0代表.jpg文件
PE文件检测：
使用pefile库，脚本能够检测文件是否为PE文件（.exe或.dll）。PE文件有特定的文件头标识，通过检查文件的魔术字节和签名，脚本能够确认一个文件是否为有效的PE文件。自动重命名文件：
如果文件没有扩展名，脚本会根据文件头字节判断文件类型并自动为文件添加扩展名。对于PE文件，脚本会根据其特征（如是否为动态链接库）将文件重命名为.exe或.dll。扫描目录：
脚本会扫描指定的目录及其子目录，逐个处理文件。它检查每个文件是否已经有扩展名，如果没有，尝试根据文件的签名自动为其添加合适的扩展名。如果是PE文件，脚本会根据其特征将文件重命名为.exe或.dll。
代码实现
# coding=utf-8
import os
import pefile
import re
import struct
import logging
# 配置日志模块，用于记录信息、警告和错误
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
# 文件类型映射字典，键为文件头（magic bytes），值为文件扩展名
FILE_TYPES = {
"52617221": '.rar', # RAR文件
"504B0304": '.zip', # ZIP文件
"377ABCAF271C": '.7z', # 7Z文件
"25504446": '.pdf', # PDF文件
"89504E470D0A1A0A": '.png', # PNG文件
"FFD8FFE0": '.jpg', # JPEG文件
"47494638": '.gif', # GIF文件
"424D": '.bmp', # BMP文件
"494433": '.mp3', # MP3文件
"6674797069736F6D": '.mp4', # MP4文件
"4D5A": '.exe', # EXE文件
"4344303031": '.iso', # ISO文件
"7573746172": '.tar', # TAR文件
"1F8B08": '.gz', # GZ文件
"7801730D626260": '.dmg', # DMG文件
"664C6143": '.flac', # FLAC文件
"52494646": '.wav', # WAV文件
"52494646": '.avi', # AVI文件
"1A45DFA3": '.mkv', # MKV文件
"53514C69746520666F726D": '.sqlite', # SQLITE文件
"3C21444F4354595045": '.html', # HTML文件
"3C3F786D6C20": '.xml', # XML文件
"7B": '.json', # JSON文件（以 { 开头）
"5B": '.json', # JSON文件（以 [ 开头）
}
def bytes2hex(bytes_data):
"""将字节数据转换为十六进制字符串"""
return ''.join(f'{byte:02X}' for byte in bytes_data)
def detect_file_type(filename):
"""通过读取文件头字节，检测文件类型"""
try:
with open(filename, 'rb') as binfile:
# 遍历 FILE_TYPES 字典，检查文件头是否匹配
for magic, ext in FILE_TYPES.items():
num_of_bytes = len(magic) // 2
binfile.seek(0) # 重置文件指针到文件开头
header_bytes = binfile.read(num_of_bytes) # 读取与magic相同长度的字节
if bytes2hex(header_bytes) == magic: # 比较字节与签名
return ext # 返回匹配的文件扩展名
except Exception as e:
logging.error(f"Error reading file {filename}: {e}")
return '' # 未找到匹配的文件类型时返回空字符串
def is_pe_file(file_path):
"""检查文件是否为PE文件（即可执行文件，.exe 或 .dll）"""
try:
pe = pefile.PE(file_path) # 使用 pefile 库分析文件
pemagic = pe.DOS_HEADER.e_magic # 获取 DOS 头的魔术值
pesig = pe.NT_HEADERS.Signature # 获取 NT 头的签名
pe.close() # 关闭PE文件
return (pemagic == 0x5A4D) and (pesig == 0x4550) # 检查文件是否为有效的PE格式
except pefile.PEFormatError as e:
logging.warning(f"File {file_path} is not a valid PE file: {e}")
except Exception as e:
logging.error(f"Error checking PE file {file_path}: {e}")
return False # 不是有效的PE文件
def has_extension(file_path):
"""检查文件名是否包含扩展名"""
return bool(re.search(r"\.\w{2,4}$", file_path)) # 匹配以.开头，后面是2-4个字母或数字的扩展名
def rename_pe_file(file_path):
"""根据PE文件的特征重命名文件为.exe或.dll"""
try:
pe = pefile.PE(file_path) # 使用pefile库读取PE文件
characteristics = pe.FILE_HEADER.Characteristics # 获取PE文件的特征
pe.close() # 关闭PE文件
if characteristics & 0x2000: # IMAGE_FILE_DLL：如果文件为DLL
new_path = file_path + ".dll"
elif characteristics & 0x0102 or characteristics & 0x0022: # IMAGE_FILE_EXECUTABLE_IMAGE：如果文件为可执行文件
new_path = file_path + ".exe"
else:
logging.warning(f"Cannot determine type for {file_path}")
return # 无法确定类型时返回
os.rename(file_path, new_path) # 重命名文件
logging.info(f"Renamed {file_path} to {new_path}") # 记录日志
except Exception as e:
logging.error(f"Error renaming file {file_path}: {e}")
def scan_directory(directory):
"""扫描指定目录及其子目录，检测并处理文件"""
for root, _, files in os.walk(directory): # 遍历目录及其子目录
for file in files:
file_path = os.path.join(root, file) # 获取文件的完整路径
if is_pe_file(file_path): # 如果是PE文件
if not has_extension(file_path): # 如果没有扩展名
rename_pe_file(file_path) # 根据PE特征重命名文件
else:
logging.info(f"{file_path} already has an extension") # 如果已经有扩展名，记录日志
else:
logging.info(f"{file_path} is not a PE file") # 如果不是PE文件
if not has_extension(file_path): # 如果没有扩展名
ext = detect_file_type(file_path) # 检测文件类型
if ext: # 如果检测到文件类型
new_path = file_path + ext # 为文件添加扩展名
os.rename(file_path, new_path) # 重命名文件
logging.info(f"Renamed {file_path} to {new_path}") # 记录日志
if __name__ == '__main__':
target_directory = input("Enter the directory to scan: ").strip() # 提示用户输入要扫描的目录
if os.path.isdir(target_directory): # 检查目录是否有效
scan_directory(target_directory) # 开始扫描目录
else:
logging.error("Invalid directory path.") # 如果路径无效，记录错误脚本结构
文件类型检测函数：
通过detect_file_type()函数，脚本读取文件的前几个字节，并与FILE_TYPES中的已知签名进行比较。如果找到了匹配项，脚本将返回相应的文件扩展名。PE文件检测函数：
使用pefile.PE()，脚本可以打开PE文件并检查其文件头的魔术字节和签名。is_pe_file()函数通过检查这些特征来确认文件是否为PE文件。文件重命名函数：
rename_pe_file()函数会根据PE文件的特征（例如是否为DLL或可执行文件）自动为文件添加.exe或.dll扩展名。主目录扫描函数：
scan_directory()函数是脚本的核心，它遍历目标目录及其子目录，检查每个文件。如果文件没有扩展名，则调用相应的检测和重命名函数。
如何使用
安装依赖：
脚本使用了pefile库来处理PE文件，因此在运行脚本之前，您需要安装此库。可以通过以下命令安装：运行脚本：
执行脚本时，您只需要输入目标目录的路径，脚本会自动开始扫描目录，并处理所有文件：输出日志：
脚本使用logging模块记录文件处理的每一步，包括文件重命名和任何错误信息。通过查看日志，您可以了解每个文件的处理状态。
脚本的优势
自动化：通过自动检测文件类型并重命名，避免了手动检查和分类的繁琐工作。
支持多种文件类型：不仅支持常见的压缩文件和图像文件，还支持音视频文件以及PE文件的处理。
增强的PE文件处理：对PE文件的特征判断能够准确地为
.exe或.dll文件自动添加扩展名。可扩展性：脚本易于扩展，您可以根据需要添加更多的文件类型和处理逻辑。
改进和考虑
冲突处理：目前，脚本没有处理同名文件的冲突。如果目标目录中已经存在同名文件，脚本会直接覆盖它。为避免此类情况，您可以在脚本中添加文件名冲突检测。
错误处理：虽然脚本已经包含了一些错误处理（如文件无法读取），但对于不同操作系统的兼容性以及文件权限问题，仍然可以进一步改进错误处理机制。
结语
本脚本适用于需要批量处理文件的场景，特别是当文件扩展名丢失或错误时。通过自动化检测和重命名，可以节省大量的时间和精力，确保文件能够正确地分类和识别。
评论