117 lines
2.7 KiB
Python
117 lines
2.7 KiB
Python
#!/usr/bin/env python
|
|
#
|
|
# Copyright 2013 Google, Inc. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
#
|
|
# Google Author(s): Behdad Esfahbod
|
|
#
|
|
|
|
import struct
|
|
import sys
|
|
from io import BytesIO
|
|
|
|
|
|
try:
|
|
basestring # py2
|
|
except NameError:
|
|
basestring = str # py3
|
|
|
|
|
|
class PNG:
|
|
|
|
signature = bytearray ((137,80,78,71,13,10,26,10))
|
|
|
|
def __init__ (self, f):
|
|
|
|
if isinstance(f, basestring):
|
|
f = open (f, 'rb')
|
|
|
|
self.f = f
|
|
self.IHDR = None
|
|
|
|
def tell (self):
|
|
return self.f.tell ()
|
|
|
|
def seek (self, pos):
|
|
self.f.seek (pos)
|
|
|
|
def stream (self):
|
|
return self.f
|
|
|
|
def data (self):
|
|
self.seek (0)
|
|
return bytearray (self.f.read ())
|
|
|
|
class BadSignature (Exception): pass
|
|
class BadChunk (Exception): pass
|
|
|
|
def read_signature (self):
|
|
header = bytearray (self.f.read (8))
|
|
if header != PNG.signature:
|
|
raise PNG.BadSignature
|
|
return PNG.signature
|
|
|
|
def read_chunk (self):
|
|
buf = self.f.read (4)
|
|
length = struct.unpack (">I", buf)[0]
|
|
chunk_type = self.f.read (4)
|
|
chunk_data = self.f.read (length)
|
|
if len (chunk_data) != length:
|
|
raise PNG.BadChunk
|
|
crc = self.f.read (4)
|
|
if len (crc) != 4:
|
|
raise PNG.BadChunk
|
|
return (chunk_type, chunk_data, crc)
|
|
|
|
def read_IHDR (self):
|
|
(chunk_type, chunk_data, crc) = self.read_chunk ()
|
|
if chunk_type != b"IHDR":
|
|
raise PNG.BadChunk
|
|
# Width: 4 bytes
|
|
# Height: 4 bytes
|
|
# Bit depth: 1 byte
|
|
# Color type: 1 byte
|
|
# Compression method: 1 byte
|
|
# Filter method: 1 byte
|
|
# Interlace method: 1 byte
|
|
return struct.unpack (">IIBBBBB", chunk_data)
|
|
|
|
def read_header (self):
|
|
self.read_signature ()
|
|
self.IHDR = self.read_IHDR ()
|
|
return self.IHDR
|
|
|
|
def get_size (self):
|
|
if not self.IHDR:
|
|
pos = self.tell ()
|
|
self.seek (0)
|
|
self.read_header ()
|
|
self.seek (pos)
|
|
return self.IHDR[0:2]
|
|
|
|
def filter_chunks (self, chunks):
|
|
self.seek (0);
|
|
out = BytesIO ()
|
|
out.write (self.read_signature ())
|
|
while True:
|
|
chunk_type, chunk_data, crc = self.read_chunk ()
|
|
if chunk_type in chunks:
|
|
out.write (struct.pack (">I", len (chunk_data)))
|
|
out.write (chunk_type)
|
|
out.write (chunk_data)
|
|
out.write (crc)
|
|
if chunk_type == b"IEND":
|
|
break
|
|
return PNG (out)
|