Initial commit

This commit is contained in:
Niles Rogoff 2016-11-30 18:54:08 -05:00 committed by Niles Rogoff
commit 557f81ff7c
No known key found for this signature in database
GPG Key ID: B78B908F23430F80
1 changed files with 127 additions and 0 deletions

127
libpme.py Normal file
View File

@ -0,0 +1,127 @@
#!/usr/bin/env python3
# PNG Metadata Editor
# Niles Rogoff 2016
import struct, zlib, copy
class PME(object):
class color_types(object):
GREYSCALE = 0
RGB = 2
PALETTE = 3
GREYSCALE_WITH_ALPHA = 4
RGB_WITH_ALPHA = 6
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):
self._init = False
self._magic_number = b'\x89PNG\r\n\x1a\n'
self.filename = filename
self.chunks = []
if filename:
f = open(self.filename, "rb")
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)
self._verify_crc(label, data, crc)
self.chunks.append([length, label, data, crc])
if label == b"IEND":
break
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 = self.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