def embed_message(self, message):
self._open_image()
self.pix_array = self.image.load()
# the message packet will look like <indicator><msg_len><msg>
# ensure that there is enough room in the picture to store the message
# Indicator will be max of 2 bytes, length = 4, and the message * 8
full_msg_len = 2 + 4 + (len(message) * 8)
if full_msg_len > (self.image.size[0] * self.image.size[1]):
raise simplePicStegoError.Error("Message is too large")
# set indicator
message_to_embed = struct.pack(">H", simplePicStegoDefines.indicator)
# append length
message_to_embed += struct.pack(">I", len(message))
# append message as byte string
for c in message:
message_to_embed += struct.pack(">B", ord(c))
next_pixel_iter = self._get_next_pixel()
for byte in message_to_embed:
for bit_shift in range(8):
new_bit = ((byte << bit_shift) & 0x80) >> 7
pixel = next_pixel_iter.__next__()
# Going to use the red pixel for now. Can configure later
pixel = ((pixel[0] & 0xfe | new_bit), pixel[1], pixel[2])
self.pix_array[self.current_height, self.current_width] = pixel
self.image.save(self.filename.split(".")[0] + "_new.png")
print("Finished encoding")
def _get_next_pixel(self):
image_width, image_height = self.image.size
for i in range(image_height):
for j in range(image_width):
self.current_height = i
self.current_width = j
yield self.pix_array[i, j]
else:
raise simplePicStegoError.Error("All out of picture")
def init_program():
parser = argparse.ArgumentParser(description="An app that embeds strings into images")
# parser.add_argument("--version", action="version", version="%(prog)s %s" % version)
parser.add_argument("-e", action="store", dest="encode_file", help="The file name to store the string",
default=False)
parser.add_argument("-m", action="store", dest="message", help="The message to store. Combine with -e",
default=None)
parser.add_argument("-d", action="store", dest="decode_file", help="The file to extract the message")
results = parser.parse_args()
if (results.encode_file and results.decode_file) or (not results.emcode_file and not results.decode_file):
raise UnknownFunctionError("Must either encode or decode a file")
elif results.encode_file: # create object to encode message into file and perform operation
if results.encode_file.split(".")[1] != "png":
raise simplePicStegoError.Error("Can only support png file right now")
simplePicStegoEmbed.PicEmbed(results.encode_file, results.message).embed_message()
elif results.decode_file: # create object to attempt to find a message within an image file
if results.decode_file.split(".")[1] != "png":
raise simplePicStegoError.Error("Can only support png file right now")
message = simplePicStegoReveal.SimplePicStegoReveal(results.decode_file).reveal()
print(message)
def _search_for_indicator(self):
"""
Searches image for indicator. Raises except of none found.
Should be run first when extracting message
:return: None
"""
# Indicator will be 2 bytes
indicator = self._get_byte_from_pix()
try:
while True:
indicator = (indicator << 8) & 0xffff
indicator = indicator | self._get_byte_from_pix()
if indicator == simplePicStegoDefines.indicator:
return
except simplePicStegoError.Error:
raise simplePicStegoError.Error("Could not find indicator")
def _open_image(self):
if not os.path.isfile(self.filename):
raise simplePicStegoError.Error("Could not find file %s." % self.filename)
self.image = Image.open(self.filename)