Spaces:

Yassine
/

Stego

Sleeping

Stego / stc.py

YassineYousfi

even lower payload

46444c7 over 2 years ago

No virus

5.04 kB

	#!/usr/bin/env python3

	import sys
	import os.path
	import math
	import random
	import struct
	import hashlib
	from PIL import Image
	from ctypes import *
	from Crypto.Cipher import AES
	from Crypto.Random import get_random_bytes
	from Crypto.Util.Padding import pad, unpad

	def prepare_message(text, password):

	content_data = text.encode('utf-8')

	# Prepare a header with basic data about the message
	content_ver=struct.pack("B", 1) # version 1
	content_len=struct.pack("!I", len(content_data))
	content=content_ver+content_len+content_data

	# encrypt
	enc = encrypt(content, password)

	array=[]
	for b in enc:
	for i in range(8):
	array.append((b >> i) & 1)
	return array


	# {{{ encrypt()
	def encrypt(plain_text, password):

	salt = get_random_bytes(AES.block_size)

	# use the Scrypt KDF to get a private key from the password
	private_key = hashlib.scrypt(
	password.encode(), salt=salt, n=2**14, r=8, p=1, dklen=32)

	cipher = AES.new(private_key, AES.MODE_CBC)
	cipher_text = cipher.encrypt(pad(plain_text, AES.block_size))
	enc = salt+cipher.iv+cipher_text

	return enc
	# }}}

	# {{{ decrypt()
	def decrypt(cipher_text, password):

	salt = cipher_text[:AES.block_size]
	iv = cipher_text[AES.block_size:AES.block_size*2]
	cipher_text = cipher_text[AES.block_size*2:]

	# Fix padding
	mxlen = len(cipher_text)-(len(cipher_text)%AES.block_size)
	cipher_text = cipher_text[:mxlen]

	private_key = hashlib.scrypt(
	password.encode(), salt=salt, n=2**14, r=8, p=1, dklen=32)

	cipher = AES.new(private_key, AES.MODE_CBC, iv=iv)
	decrypted = cipher.decrypt(cipher_text)
	#decrypted = unpad(decrypted, AES.block_size)

	return decrypted
	# }}}




	def embed(input_img_path, cost_matrix, msg_file_path, password, output_img_path, payload=0.10):

	me = os.path.abspath(os.path.dirname(__file__))
	lib = cdll.LoadLibrary(os.path.join(me, "lib", "stc.so"))

	# Prepare cover image
	im=Image.open(input_img_path)
	if im.mode in ['L']:
	width, height = im.size
	if im.mode in ['RGB', 'RGBA', 'RGBX']:
	pass
	I = im.load()
	cover = (c_int(widthheight))()
	idx=0
	for j in range(height):
	for i in range(width):
	cover[idx] = I[i, j]
	idx += 1


	# Prepare costs
	INF = 2**31-1
	costs = (c_float(widthheight*3))()
	idx=0
	for j in range(height):
	for i in range(width):
	if cover[idx]==0:
	costs[3*idx+0] = INF
	costs[3*idx+1] = 0
	costs[3*idx+2] = cost_matrix[j, i]
	elif cover[idx]==255:
	costs[3*idx+0] = cost_matrix[j, i]
	costs[3*idx+1] = 0
	costs[3*idx+2] = INF
	else:
	costs[3*idx+0] = cost_matrix[j, i]
	costs[3*idx+1] = 0
	costs[3*idx+2] = cost_matrix[j, i]
	idx += 1

	# Prepare message
	msg_bits = prepare_message(msg_file_path, password)
	if len(msg_bits)>widthheightpayload:
	print("Message too long")
	sys.exit(0)
	m = int(widthheightpayload)
	message = (c_ubyte*m)()
	for i in range(m):
	if i<len(msg_bits):
	message[i] = msg_bits[i]
	else:
	# This doesn't look optimal
	message[i] = 0
	# Hide message
	stego = (c_int(widthheight))()
	a = lib.stc_hide(width*height, cover, costs, m, message, stego)

	# Save output message
	idx=0
	for j in range(height):
	for i in range(width):
	im.putpixel((i, j), stego[idx])
	idx += 1
	im.save(output_img_path)
	im.close()



	def extract(stego_img_path, password, output_msg_path, payload=0.10):

	me = os.path.abspath(os.path.dirname(__file__))
	lib = cdll.LoadLibrary(os.path.join(me, "lib", "stc.so"))

	# Prepare stego image
	im=Image.open(stego_img_path)
	if im.mode in ['L']:
	width, height = im.size
	if im.mode in ['RGB', 'RGBA', 'RGBX']:
	pass
	I = im.load()
	stego = (c_int(widthheight))()
	idx=0
	for j in range(height):
	for i in range(width):
	stego[idx] = I[i, j]
	idx += 1

	# Extract the message
	n = width*height;
	m = int(n*payload)
	extracted_message = (c_ubyte*m)()
	s = lib.stc_unhide(n, stego, m, extracted_message)

	# Save the message
	enc = bytearray()
	idx=0
	bitidx=0
	bitval=0
	for b in extracted_message:
	if bitidx==8:
	enc.append(bitval)
	bitidx=0
	bitval=0
	bitval \|= b<<bitidx
	bitidx+=1
	if bitidx==8:
	enc.append(bitval)

	# decrypt
	cleartext = decrypt(enc, password)

	# Extract the header and the message
	content_ver=struct.unpack_from("B", cleartext, 0)
	content_len=struct.unpack_from("!I", cleartext, 1)
	content=cleartext[5:content_len[0]+5]

	f = open(output_msg_path, 'w')
	f.write(content.decode())
	f.close()