123

# -*- coding: utf-8 -*-
"""steg_ANAL.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1jeb954KZS2bvn9AFzsW_rfEGUI8H0Mlv
"""

#@title
import numpy as np
import matplotlib.image as img
import matplotlib.pyplot as plt
from PIL import Image
import random
import string

test_message = "Among other public buildings in a certain town, which for many reasons it will be prudent to refrain from mentioning, and to which I will assign no fictitious name, there is one anciently common to most towns, great or small: to wit, a workhouse; and in this workhouse was born; on a day and date which I need not trouble myself to repeat, inasmuch as it can be of no possible consequence to the reader, in this stage of the business at all events; the item of mortality whose name is prefixed to the head of this chapter."

### Для теста на переполнение
# Генератор случайного текста
def random_text_generator(size):
    return ''.join(random.choices(string.ascii_letters, k=size))

# Разница между двумя строками
def str_diff(s1, s2):
    return sum(x != y for (x, y) in zip(s1, s2))

def str_to_bin(str):
    return '{0:b}'.format(str)

def bin_to_int(bin):
    return int("".join(bin), 2)

# Функция для шифрования сообщения методом LSB
# img - конейнер (изображение)
# message - сообщение-строка (обычная строка)
# Первые max_bits бита в зашифрованном изображении - ключ (бинарное представление длинны сообщения (в битах))

def lsb_encode(img, message, max_bits=22):
    encode_img = np.copy(img)

    # Ascii в бинарный формат
    encoded_message = str(bin(int.from_bytes(message.encode(), 'big')))[2:]

#     print(f"Длинна сообщения: {len(encoded_message)}")

    # Создание ключа
    key = str_to_bin(len(encoded_message))

    # Дополнение ключа до key_size битов
    for i in range(max_bits - len(key)):
        key = "0" + key

#     print(f"Ключ: {key}")

    # Процент, на сколько контейнер заполнен (в процентах)
    im_size = encode_img.shape[0] * encode_img.shape[1]
    cont_fullness = 100 * abs(len(key) + len(encoded_message)) / im_size

#     print(f"Заполненность контейнера(для последнего бита Red): {cont_fullness}")

    counter = 0
    x = 0
    y = 0

    # Записть ключа и сообщения в изображение ( с переполнением )
    while(counter < len(encoded_message) + max_bits):
        r                  = encode_img[x,y][0]
        binary_r           = list(str_to_bin(r))
        if (counter < max_bits):
            binary_r[-1]   = key[counter]
        else:
            binary_r[-1]   = encoded_message[counter - max_bits]

        encode_img[x,y][0] = bin_to_int(binary_r)
        counter += 1

        if(y == encode_img.shape[1] - 1):
            if(x == encode_img.shape[0] - 1):
                x = 0
                y = 0
            else:
                x +=1
                y  =0
        else:
            y +=1

    return (encode_img, cont_fullness)

def lsb_decode(decode_img, max_bits=22, isShuffle=False, shuffle_key=None):
    message = ""
    key = ""

    # Получить ключ
    for y in range(0, max_bits):
        r         = decode_img[0, y][0]
        binary_r  = list(str_to_bin(r))
        key      += binary_r[-1]

    key = bin_to_int(key)

    x = 0
    y = 0
    counter = 0

    while(counter < key + max_bits):
        r = decode_img[x,y][0]
        binary_r = list(str_to_bin(r))
        if (counter >= max_bits):
            message += binary_r[-1]
        counter += 1

        if(y == decode_img.shape[1] - 1):
            if(x == decode_img.shape[0] - 1):
                x = 0
                y = 0
            else:
                x +=1
                y  =0
        else:
            y +=1

    result = ''
    message = '0' + message

    for i in range (0, len(message), 8):
        bin_oct = message[i: i+8]
        bin_oct = int(bin_oct, 2)
        try:
            symbol = bin_oct.to_bytes((bin_oct.bit_length() + 7) // 8, 'big').decode()
            result += symbol
        except UnicodeDecodeError:
            result += ''

    return result

# Читаем изображение
image = img.imread("C:/lsb/cat300x300.jpg")

plt.figure(figsize = (80,50))
plt.imshow(image)

output = lsb_encode(image, test_message)

plt.figure(figsize = (80,50))
plt.imshow(output[0])

lsb_decode(output[0])

# Метрики (для R-канала)
maxD = 0
nmse = 0
diff_sum_2 = 0
diff_inp_2 = 0
max_inp    = 0
max_inp_diff    = 0
diff_2     = 0

for x in range(0, image.shape[0]):
    for y in range(0, image.shape[1]):
        r_inp = int(image[x,y][0])
        r_out = int(output[0][x,y][0])

        diff_abs  = abs(r_out - r_inp)

        diff_sum_2 += (r_inp-r_out)**2
        diff_inp_2 += (r_inp)**2

        diff_2         += (r_inp - r_out)**2
        max_inp_diff   = max(max_inp_diff, r_inp**2)

        maxD  = max(diff_abs, maxD)

print("maxD: ", maxD)
print("nmse: ", diff_sum_2 / diff_inp_2)
print("snr: ", 1 / (diff_sum_2 / diff_inp_2))
print("psnr: ", max_inp_diff / diff_2)

### Тест на ошибку при переполнении
test_counts         = 30
test_letter_range   = (1000, 30000)
results_fullness    = []
results_errors      = []

for i in range(test_counts):
    letter_size       = (i + 1) * int((test_letter_range[1]) / test_counts)
    test_mess         = random_text_generator(letter_size)
    test_output       = lsb_encode(image, test_mess)
    test_out_im       = test_output[0]
    test_fullness     = test_output[1]
    output_mess       = lsb_decode(test_out_im)[0:letter_size]

    diff              = str_diff(test_mess, output_mess)

    # Ошибка (0...1)
    error             = ((100 * diff) / letter_size)

    results_fullness.append(test_fullness)
    results_errors.append(error)

    print('\n')

    print((test_fullness, letter_size, len(output_mess), diff, error))

plt.figure(figsize = (15, 10))
plt.xlabel("Заполнение контейнера (%)")
plt.ylabel("Ошибка (%)")
plt.plot(results_fullness,results_errors)