2023-10-02 17:25:40 +00:00
|
|
|
|
import binascii
|
|
|
|
|
import json
|
|
|
|
|
import os
|
|
|
|
|
import random
|
|
|
|
|
import re
|
|
|
|
|
import sys
|
|
|
|
|
import time
|
|
|
|
|
from PIL import Image
|
2023-12-11 17:05:48 +00:00
|
|
|
|
from type.PNG import PNG
|
|
|
|
|
# from type.JPG import JPG
|
|
|
|
|
# from type.GIF import GIF
|
2023-10-02 17:25:40 +00:00
|
|
|
|
|
2023-12-11 17:05:48 +00:00
|
|
|
|
types = [PNG] # 暂时先只处理PNG,其他的都没写好呢[PNG, JPG, GIF]
|
2023-10-02 17:25:40 +00:00
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def check_image_type(filename: str, data: bytearray):
|
|
|
|
|
file_type1 = None
|
2023-12-11 17:05:48 +00:00
|
|
|
|
for model in types:
|
|
|
|
|
if re.match(f'.*\.{model.extension}$', filename.lower()):
|
|
|
|
|
file_type1 = model
|
|
|
|
|
break
|
2023-10-02 17:25:40 +00:00
|
|
|
|
|
|
|
|
|
file_type2 = None
|
|
|
|
|
weight = 0
|
2023-12-11 17:05:48 +00:00
|
|
|
|
for model in types:
|
|
|
|
|
weight2 = model.test(data)
|
|
|
|
|
if weight2 > weight:
|
|
|
|
|
file_type2 = model
|
|
|
|
|
weight = weight2
|
2023-10-02 17:25:40 +00:00
|
|
|
|
|
|
|
|
|
if file_type1 != None:
|
|
|
|
|
if file_type1 == file_type2:
|
2023-12-11 17:05:48 +00:00
|
|
|
|
print(f'看起来文件格式是{file_type1.extension}')
|
2023-10-02 17:25:40 +00:00
|
|
|
|
else:
|
2023-12-11 17:05:48 +00:00
|
|
|
|
print(f'文件名中格式是{file_type1.extension},但是解析文件头格式发现{weight}%是{file_type2.extension}')
|
2023-10-02 17:25:40 +00:00
|
|
|
|
else:
|
|
|
|
|
if file_type2 != None:
|
|
|
|
|
file_type1 = file_type2
|
2023-12-11 17:05:48 +00:00
|
|
|
|
print(f'检测文件头认为文件格式{weight}%是{file_type2.extension}')
|
2023-10-02 17:25:40 +00:00
|
|
|
|
else:
|
|
|
|
|
print('未检测到任何图片格式特征,该文件可能不是图片文件')
|
|
|
|
|
|
2023-12-11 17:05:48 +00:00
|
|
|
|
for i in range(len(types)):
|
|
|
|
|
print(f' {i}、{types[i].extension}')
|
|
|
|
|
index = input('要指定文件格式吗?或者直接回车。请输入文件格式序号:')
|
2023-10-02 17:25:40 +00:00
|
|
|
|
if index.isdigit():
|
|
|
|
|
index = int(index)
|
2023-12-11 17:05:48 +00:00
|
|
|
|
file_type1 = types[index]
|
2023-10-02 17:25:40 +00:00
|
|
|
|
|
|
|
|
|
return file_type1
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
# 总函数
|
|
|
|
|
def main(filename: str):
|
|
|
|
|
# 读取图片数据
|
|
|
|
|
data = None
|
|
|
|
|
with open(filename, 'rb') as f:
|
|
|
|
|
data = bytearray(f.read())
|
|
|
|
|
|
|
|
|
|
# 确定文件格式
|
2023-12-11 17:05:48 +00:00
|
|
|
|
model = check_image_type(filename, data)
|
2023-10-02 17:25:40 +00:00
|
|
|
|
|
|
|
|
|
# 检测并修复
|
|
|
|
|
changed = False
|
|
|
|
|
if model != None:
|
|
|
|
|
info = model.get_info(data)
|
|
|
|
|
# print(json.dumps(info))
|
|
|
|
|
changed1, data, tail = model.check_tail(data)
|
|
|
|
|
changed2, data, head = model.repair_sig(data)
|
|
|
|
|
changed3, data, width, height = model.repair_size(data)
|
2023-12-11 17:05:48 +00:00
|
|
|
|
changed4, data, data_append = model.check_data(data)
|
|
|
|
|
checked_channel = False # 是否检测过RGB通道,如果没检测过,则下面还会尝试检测一遍
|
|
|
|
|
try:
|
|
|
|
|
# 尝试分析通道信息
|
|
|
|
|
changed5, data_r, data_g, data_b = model.check_channels(filename, width, height)
|
|
|
|
|
checked_channel = True
|
|
|
|
|
except Exception as e:
|
|
|
|
|
pass
|
|
|
|
|
#print('[x] 尝试打开图片失败,失败原因:' + str(e))
|
|
|
|
|
|
2023-10-02 17:25:40 +00:00
|
|
|
|
changed = changed1 or changed2 or changed3
|
|
|
|
|
else:
|
|
|
|
|
print('[*] 未进行任何处理')
|
|
|
|
|
|
|
|
|
|
if changed:
|
|
|
|
|
# 创建结果目录
|
|
|
|
|
basename = os.path.basename(filename)
|
|
|
|
|
dirname = filename.replace('.', '_') + str(int(time.time()))
|
|
|
|
|
dirpath = os.path.join(os.path.dirname(os.path.abspath(filename)), dirname)
|
|
|
|
|
os.mkdir(dirpath)
|
|
|
|
|
|
|
|
|
|
# 保存解析结果
|
|
|
|
|
imagepath = os.path.join(dirpath, basename)
|
|
|
|
|
f = open(imagepath, 'wb')
|
|
|
|
|
f.write(data)
|
|
|
|
|
f.close()
|
|
|
|
|
if changed1 and tail:
|
|
|
|
|
f = open(os.path.join(dirpath, 'tail'), 'wb')
|
|
|
|
|
f.write(tail)
|
|
|
|
|
f.close()
|
|
|
|
|
if changed2 and head:
|
|
|
|
|
f = open(os.path.join(dirpath, 'head'), 'wb')
|
|
|
|
|
f.write(head)
|
|
|
|
|
f.close()
|
2023-12-11 17:05:48 +00:00
|
|
|
|
for i in range(len(data_append)):
|
|
|
|
|
f = open(os.path.join(dirpath, 'data_append_' + str(i)), 'wb')
|
|
|
|
|
f.write(data_append[i])
|
|
|
|
|
f.close()
|
2023-10-02 17:25:40 +00:00
|
|
|
|
print(f'[*] 分析结果已保存到 {dirpath}')
|
2023-12-11 17:05:48 +00:00
|
|
|
|
if not checked_channel:
|
|
|
|
|
# 尝试分析通道信息
|
|
|
|
|
changed5, data_r, data_g, data_b = model.check_channels(imagepath, width, height)
|
2023-10-02 17:25:40 +00:00
|
|
|
|
try:
|
2023-12-11 17:05:48 +00:00
|
|
|
|
# 尝试打开图片看一看
|
2023-10-02 17:25:40 +00:00
|
|
|
|
im = Image.open(imagepath)
|
|
|
|
|
im.show()
|
|
|
|
|
except Exception as e:
|
|
|
|
|
print('[x] 尝试打开图片失败,失败原因:' + str(e))
|
|
|
|
|
else:
|
|
|
|
|
print('[*] 未进行任何处理')
|
|
|
|
|
|
|
|
|
|
print('[*] 分析完成!')
|
|
|
|
|
|
|
|
|
|
# 开始执行
|
2023-12-11 17:05:48 +00:00
|
|
|
|
if len(sys.argv) > 1:
|
|
|
|
|
filename = sys.argv[1]
|
|
|
|
|
|
2023-10-02 17:25:40 +00:00
|
|
|
|
main(filename)
|