135 lines
4.5 KiB
Python
Executable File
135 lines
4.5 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
# PNG Metadata Editor
|
|
# Niles Rogoff 2016
|
|
import zlib, copy
|
|
class color_types(object):
|
|
GREYSCALE = 0
|
|
RGB = 2
|
|
PALETTE = 3
|
|
GREYSCALE_WITH_ALPHA = 4
|
|
RGB_WITH_ALPHA = 6
|
|
class PME(object):
|
|
def _assert(self, statement):
|
|
if not self._damaged:
|
|
assert(statement)
|
|
def _int(self, binary):
|
|
return int.from_bytes(binary, byteorder="big")
|
|
def _bytes(self, integer, length):
|
|
return integer.to_bytes(length, "big")
|
|
def __init__(self, filename=False, damaged=False):
|
|
self._damaged = damaged
|
|
self._init = False
|
|
self._magic_number = b'\x89PNG\r\n\x1a\n'
|
|
self.filename = filename
|
|
self.chunks = []
|
|
if filename:
|
|
f = open(self.filename, "rb")
|
|
self._assert(f.read(8) == self._magic_number)
|
|
while True:
|
|
length = f.read(4)
|
|
label = f.read(4)
|
|
data = f.read(self._int(length))
|
|
crc = f.read(4)
|
|
if not self._damaged:
|
|
self._verify_crc(label, data, crc)
|
|
self.chunks.append([length, label, data, crc])
|
|
if label == b"IEND":
|
|
break
|
|
self._assert(self.chunks[0][1] == b"IHDR")
|
|
self.recalculate_properties()
|
|
else:
|
|
self.chunks = [
|
|
[b"\0\0\0\0", b"IHDR", b"\0" * 13, b""],
|
|
[b"\0\0\0\0", b"IDAT", b"", b""],
|
|
[b"\0\0\0\0", b"IEND", b"", b""]
|
|
]
|
|
self.width = self.height = 0 # the user can decide
|
|
self.bit_depth = 8 # a sane default
|
|
self.color_type = color_types.RGB_WITH_ALPHA
|
|
self.compression_method = 0
|
|
self.filter_method = 0 # The only one as far as I know
|
|
self.interlace_method = 0
|
|
self._init = True
|
|
self.recalculate_IHDR()
|
|
for i in range(len(self.chunks)):
|
|
self.recalculate_crc(i)
|
|
self.recalculate_length(i)
|
|
# The alternative to this is to have an useless property accessor for all seven properties. Trust me, this is better.
|
|
def __setattr__(self, name, value):
|
|
super(PME, self).__setattr__(name, value)
|
|
if name in ["width", "height", "bit_depth", "color_type", "compression_method", "filter_method", "interlace_method"]:
|
|
self.recalculate_IHDR()
|
|
def _calculate_crc(self, label, binary):
|
|
calculated = zlib.crc32(label)
|
|
calculated = zlib.crc32(binary, calculated) & 0xFFFFFFFF
|
|
return calculated
|
|
def _verify_crc(self, label, binary, crc):
|
|
assert(self._calculate_crc(label, binary) == self._int(crc))
|
|
def recalculate_properties(self):
|
|
self.width = self._int(self.chunks[0][2][0:4])
|
|
self.height = self._int(self.chunks[0][2][4:8])
|
|
self.bit_depth = self._int([self.chunks[0][2][8]])
|
|
self.color_type = self._int([self.chunks[0][2][9]])
|
|
self.compression_method = self._int([self.chunks[0][2][10]])
|
|
self.filter_method = self._int([self.chunks[0][2][11]])
|
|
self.interlace_method = self._int([self.chunks[0][2][12]])
|
|
def recalculate_IHDR(self):
|
|
if not self._init: return
|
|
final = self._bytes(self.width, 4)
|
|
final += self._bytes(self.height, 4)
|
|
final += self._bytes(self.bit_depth, 1)
|
|
final += self._bytes(self.color_type, 1)
|
|
final += self._bytes(self.compression_method, 1)
|
|
final += self._bytes(self.filter_method, 1)
|
|
final += self._bytes(self.interlace_method, 1)
|
|
self.chunks[0][2] = final
|
|
self.recalculate_crc(0)
|
|
self.recalculate_length(0)
|
|
def _index(self, index):
|
|
if type(index) == list:
|
|
return self.chunks.index(index)
|
|
if type(index) == str:
|
|
index = bytes(index)
|
|
if type(index) == bytes: # PLEASE FOR THE LOVE OF GOD DO NOT DO THIS
|
|
index = [x[1] for x in self.chunks].index(index) # index
|
|
return index
|
|
def recalculate_crc(self, index):
|
|
index = self._index(index)
|
|
self.chunks[index][3] = self._bytes(self._calculate_crc(self.chunks[index][1], self.chunks[index][2]), 4)
|
|
def recalculate_length(self, index):
|
|
index = self._index(index)
|
|
self.chunks[index][0] = self._bytes(len(self.chunks[index][2]), 4)
|
|
def save(self,filename=False):
|
|
if not filename:
|
|
filename = self.filename
|
|
f = open(filename, "wb")
|
|
f.write(self._magic_number)
|
|
for chunk in self.chunks:
|
|
for field in chunk:
|
|
f.write(field)
|
|
decompress = zlib.decompress
|
|
compress = zlib.compress
|
|
|
|
def get_concatenated_idat_data(self):
|
|
data = b""
|
|
for chunk in self.chunks:
|
|
if chunk[1] == b'IDAT':
|
|
data += chunk[2]
|
|
return data
|
|
def write_raw_idat_data(self, data):
|
|
i = 0
|
|
for chunk in copy.deepcopy(self.chunks):
|
|
if chunk[1] == b'IDAT':
|
|
i += 1
|
|
if i > 1:
|
|
del self.chunks[self.chunks.index(chunk)]
|
|
for index in range(len(self.chunks)):
|
|
if self.chunks[index][1] == b'IDAT':
|
|
self.chunks[index][2] = data
|
|
self.recalculate_crc(index)
|
|
self.recalculate_length(index)
|
|
return True
|
|
return False
|
|
decompress = zlib.decompress
|
|
compress = zlib.compress
|