-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathlibpme.py
executable file
·134 lines (133 loc) · 4.46 KB
/
libpme.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
#!/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