Don't be a punk, punk
import binascii | |
import struct | |
class Punk(object): | |
_END_CHUNK_TYPE = 'IEND' | |
_PUNK_CHUNK_TYPE = 'puNk' | |
_MAX_BYTES = 2147483647 | |
_chunks = dict() | |
def __init__(self): | |
self._mode = None | |
self._file = None | |
self._output = None | |
self._bytes_to_hide = None | |
self._bytes_read = 0 | |
def decode(self, input_file, output_file): | |
self._mode = 'decode' | |
self._file = open(input_file, 'rb+') | |
self._output = open(output_file, 'wb+') | |
# First move cursor past the signature | |
self._read_bytes(8) | |
# Start reading chunks | |
self._read_next_chunk() | |
def encode(self, input_file, bytes_to_hide): | |
self._mode = 'encode' | |
self._file = open(input_file, 'rb+') | |
self._bytes_to_hide = bytes_to_hide | |
# First move cursor past the signature | |
self._read_bytes(8) | |
# Start reading chunks | |
self._read_next_chunk() | |
def _read_bytes_as_hex(self, position): | |
return self._read_bytes(position).encode('hex') | |
def _read_bytes_as_ascii(self, position): | |
return self._read_bytes(position).encode('ascii') | |
def _read_bytes_as_int(self, position): | |
return int(self._read_bytes_as_hex(position), 16) | |
def _read_bytes(self, byte_count): | |
self._bytes_read += byte_count | |
return self._file.read(byte_count) | |
def _rewind_bytes(self, byte_count): | |
self._bytes_read -= byte_count | |
self._file.seek(self._bytes_read) | |
def _inject_punk_chunk(self): | |
# Move back 8 bytes. | |
self._rewind_bytes(8) | |
chunk_size = len(self._bytes_to_hide) | |
print 'Hiding', (chunk_size / 1024), 'kB (', chunk_size, 'bytes)' | |
# Create a temporary byte array for the CRC check. | |
tmp_bytes = bytearray() | |
# First write the chunk type | |
tmp_bytes.extend(bytearray(self._PUNK_CHUNK_TYPE)) | |
# Now write the bytes of whatever we're trying to hide | |
tmp_bytes.extend(self._bytes_to_hide) | |
print 'Injecting punk chunk' | |
# Write the chunk size | |
self._file.write(bytearray(struct.pack('!i', chunk_size))) | |
# And the type | |
self._file.write(bytearray(self._PUNK_CHUNK_TYPE)) | |
self._file.write(self._bytes_to_hide) | |
crc = binascii.crc32(tmp_bytes) | |
self._file.write(bytearray(struct.pack('!i', crc))) | |
# Write the end chunk. Start with the size. | |
self._file.write(bytearray(struct.pack('!i', 0))) | |
# Then the chunk type. | |
self._file.write(bytearray(self._END_CHUNK_TYPE)) | |
print 'Punk chunk injected' | |
def _read_next_chunk(self): | |
chunk_size = self._read_bytes_as_int(4) | |
print 'Chunk size:', chunk_size | |
chunk_type = self._read_bytes_as_ascii(4) | |
print 'Chunk type:', chunk_type | |
if self._mode == 'encode' and chunk_type == self._END_CHUNK_TYPE: | |
self._inject_punk_chunk() | |
print 'Reached EOF' | |
self._file.close() | |
return | |
content = self._read_bytes(chunk_size) | |
crc = self._read_bytes_as_hex(4) | |
print 'CRC:', crc | |
if self._mode == 'decode' and chunk_type == self._PUNK_CHUNK_TYPE: | |
print "Found a punk chunk", len(content), "bytes. Writing to file" | |
self._output.write(bytearray(content)) | |
self._output.close() | |
self._file.close() | |
return | |
self._read_next_chunk() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
This comment has been minimized.
I think you forgot to add CRC for the IEND chunk in the output file.