-
Notifications
You must be signed in to change notification settings - Fork 26
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
1 parent
f1317e3
commit 38aeacb
Showing
1 changed file
with
135 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,135 @@ | ||
import math | ||
|
||
import numpy as np | ||
import itertools as it | ||
|
||
from constants.constants import * | ||
from utils.bit_storing import store_bits, retrieve_bits | ||
from utils.conversor import conv, bin2Arr, arr2int | ||
from utils.ctxt import * | ||
|
||
# Inyect a file and a filename into an array | ||
def inject_file(arr:np.ndarray, file:bytes, filename:bytes, store_random:bool) -> np.ndarray: | ||
#------------------------------------------------------- | ||
print(ctxt("\nPreparing...\n", Fore.MAGENTA)) | ||
#------------------------------------------------------- | ||
# Convert file and filename to binary | ||
file_bin = conv(int(file.hex(), 16), 2) | ||
filename_bin = conv(int(filename.hex(), 16), 2) | ||
#------------------------------------------------------- | ||
# Flatten the array | ||
arr_flat = arr.flatten() | ||
#------------------------------------------------------- | ||
# Get the maximum number of bits that can be modified in each channel and its base and mask | ||
channel_bits = np.iinfo(arr.dtype).bits | ||
max_mod_bits_base = 2**channel_bits | ||
max_mod_bits_mask = max_mod_bits_base | ||
#------------------------------------------------------- | ||
# Maximum size the file can have | ||
max_file_size = arr.size - MAX_FN_SIZE_BIN - 1 | ||
|
||
# Length of max_file_size in base 2 | ||
max_file_size_len = len(conv(max_file_size, 2)) | ||
|
||
# Available pixels that can be used to store the file | ||
available_channels = max_file_size - max_file_size_len | ||
#------------------------------------------------------- | ||
# Times the file_len_b2 is bigger than the available channels (rounded up) | ||
file_len_b2 = len(file_bin) | ||
times = math.ceil(file_len_b2 / available_channels) | ||
|
||
# If times > channel_bits | ||
if times > channel_bits: | ||
raise Exception(f"Base file is too small to store the file") | ||
|
||
# times must be between 1 and max_mod_bits (maximum modification of bits you can do to a channel) | ||
times = times if times <= channel_bits else channel_bits | ||
|
||
mod_bits = times # The number of bits that will be modified | ||
base = 2**mod_bits # The base that will be used to store the file | ||
base_mask = max_mod_bits_base - base # The mask that will be used to store the file | ||
bin_mask = max_mod_bits_base - 2 | ||
#------------------------------------------------------- | ||
# Convert the filename to base 2 and fill it left with 0s | ||
filename_conv = bin2Arr(filename_bin.zfill(MAX_FN_SIZE_BIN), 2) | ||
|
||
# Convert the file to base {base} | ||
file_conv = bin2Arr(file_bin, base) | ||
|
||
# Get the file size and convert it to base 2 | ||
file_size = len(file_conv) | ||
file_size_conv = bin2Arr(conv(file_size, 2).zfill(max_file_size_len),2) # image resolution in pixels | ||
#------------------------------------------------------- | ||
print(f"Modified bits per channel: {ctxt(mod_bits, Fore.YELLOW)}") | ||
print(f"Base file modification: {ctxt(f'{round(base/(2**channel_bits)*100, 2)}%', Fore.YELLOW)}") | ||
#------------------------------------------------------- | ||
# (1) Store mod_bits at idx 0 | ||
arr_flat[0] = store_bits(arr_flat[0], mod_bits, max_mod_bits_mask) | ||
|
||
# Generate random values | ||
print(ctxt("\nGenerating random values...", Fore.MAGENTA)) | ||
end_idx = 1 + len(file_size_conv) + len(filename_conv) + len(file_conv) | ||
if store_random: rands = np.random.randint(low = 0, high=base, size=arr_flat.size - end_idx) | ||
|
||
# Use zip to inject information in channels, using base_mask | ||
print(ctxt("\nInjecting information...", Fore.MAGENTA)) | ||
arr_flat[1:] = [ | ||
store_bits(channel, val, base_mask) for channel, val in zip( | ||
arr_flat[1:], | ||
np.concatenate((file_size_conv, filename_conv, file_conv, rands)), | ||
) | ||
] | ||
#------------------------------------------------------- | ||
# Return flat_img to original shape if needed | ||
if arr.shape != arr_flat.shape: | ||
print(ctxt("\nReshaping...", Fore.MAGENTA)) | ||
return arr_flat.reshape(arr.shape) | ||
|
||
else: # If the array was already flat | ||
return arr_flat | ||
|
||
|
||
|
||
# Extract the file from the array | ||
def extract_file(mod_arr:np.ndarray) -> dict[bytes, bytes]: | ||
#------------------------------------------------------- | ||
print(ctxt("\nPreparing...", Fore.MAGENTA)) | ||
#------------------------------------------------------- | ||
# Get the maximum number of bits that can be modified in each channel and its base | ||
channel_bits = np.iinfo(mod_arr.dtype).bits | ||
max_mod_bits_base = 2**channel_bits | ||
#------------------------------------------------------- | ||
# Flatten the array | ||
mod_arr_flat = mod_arr.flatten() | ||
#------------------------------------------------------- | ||
# Maximum size the file can have | ||
max_file_size = len(mod_arr_flat) - MAX_FN_SIZE_BIN - 1 | ||
# Length of the maximum size in binary | ||
max_file_size_len = len(conv(max_file_size, 2)) | ||
#------------------------------------------------------- | ||
# Modified bits per channel and Base | ||
mod_bits = retrieve_bits(mod_arr_flat[0], max_mod_bits_base) | ||
base = 2**mod_bits | ||
idx = 1 | ||
#------------------------------------------------------- | ||
# Get the file size | ||
file_size_arr = [retrieve_bits(x, 2) for x in mod_arr_flat[idx:idx+max_file_size_len]] | ||
file_size = arr2int(file_size_arr, 2) | ||
idx += max_file_size_len | ||
|
||
# Get the filename | ||
print(ctxt("\nRetrieving filename...", Fore.MAGENTA)) | ||
filename_arr = [retrieve_bits(x, 2) for x in mod_arr_flat[idx:idx+MAX_FN_SIZE_BIN]] | ||
filename_hex = conv(arr2int(filename_arr, 2), 16) | ||
idx += MAX_FN_SIZE_BIN | ||
|
||
# Get the file | ||
print(ctxt("\nRetrieving input file...", Fore.MAGENTA)) | ||
file_arr = [retrieve_bits(x, base) for x in mod_arr_flat[idx:idx+file_size]] | ||
file_hex = conv(arr2int(file_arr, base), 16) | ||
#------------------------------------------------------- | ||
# Return the filename and the input file in bytes | ||
return { | ||
"filename": bytes.fromhex(filename_hex), | ||
"file": bytes.fromhex(file_hex) | ||
} |