209 lines
8.3 KiB
Python
209 lines
8.3 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
CLI интерфейс для стеганографических методов LSB и DCT.
|
||
|
||
Использование:
|
||
python main.py encode <method> <image> <output> [-m MESSAGE] [-f FILE]
|
||
python main.py decode <method> <image> [-o OUTPUT]
|
||
|
||
Примеры:
|
||
python main.py encode lsb image.png result.png -m "Secret message"
|
||
python main.py encode lsb image.png result.png -f secret.txt
|
||
python main.py decode lsb result.png
|
||
python main.py decode lsb result.png -o extracted.txt
|
||
"""
|
||
|
||
import argparse
|
||
import sys
|
||
import os
|
||
from core.lsb import encode_lsb, decode_lsb
|
||
from core.dct import encode_dct, decode_dct
|
||
from core.utils import calculate_capacity, calculate_psnr
|
||
|
||
|
||
def read_message_from_file(file_path: str) -> str:
|
||
"""Читает сообщение из текстового файла."""
|
||
with open(file_path, 'r', encoding='utf-8') as f:
|
||
return f.read()
|
||
|
||
|
||
def save_message_to_file(file_path: str, message: str) -> None:
|
||
"""Сохраняет извлеченное сообщение в файл."""
|
||
with open(file_path, 'w', encoding='utf-8') as f:
|
||
f.write(message)
|
||
print(f"Сообщение сохранено в: {file_path}")
|
||
|
||
|
||
def show_capacity(image_path: str) -> None:
|
||
"""Показывает максимальную вместимость изображения для LSB метода."""
|
||
try:
|
||
capacity = calculate_capacity(image_path)
|
||
print(f"Максимальная вместимость изображения: {capacity} байт")
|
||
print(f"Это примерно {capacity / 1024:.2f} КБ текста")
|
||
except Exception as e:
|
||
print(f"Ошибка при расчете вместимости: {e}")
|
||
|
||
|
||
def show_psnr(original_path: str, modified_path: str) -> None:
|
||
"""Показывает PSNR между двумя изображениями."""
|
||
try:
|
||
psnr = calculate_psnr(original_path, modified_path)
|
||
print(f"PSNR: {psnr:.2f} dB")
|
||
|
||
if psnr > 40:
|
||
print("Качество: Отличное (изменения практически незаметны)")
|
||
elif psnr > 30:
|
||
print("Качество: Хорошее")
|
||
elif psnr > 20:
|
||
print("Качество: Заметные искажения")
|
||
else:
|
||
print("Качество: Сильные искажения")
|
||
except Exception as e:
|
||
print(f"Ошибка при расчете PSNR: {e}")
|
||
|
||
|
||
def encode_command(args) -> None:
|
||
"""Обрабатывает команду encode."""
|
||
# Получаем сообщение
|
||
if args.message and args.file:
|
||
print("Ошибка: используйте либо -m, либо -f, но не оба")
|
||
sys.exit(1)
|
||
|
||
if args.message:
|
||
message = args.message
|
||
elif args.file:
|
||
try:
|
||
message = read_message_from_file(args.file)
|
||
except FileNotFoundError:
|
||
print(f"Ошибка: файл не найден - {args.file}")
|
||
sys.exit(1)
|
||
else:
|
||
print("Ошибка: укажите сообщение через -m или -f")
|
||
sys.exit(1)
|
||
|
||
# Проверяем вместимость (только для LSB)
|
||
if args.method == 'lsb':
|
||
try:
|
||
capacity = calculate_capacity(args.image)
|
||
msg_bytes = len(message.encode('utf-8'))
|
||
if msg_bytes > capacity:
|
||
print(f"Предупреждение: сообщение может не поместиться!")
|
||
print(f"Доступно: {capacity} байт, требуется: {msg_bytes} байт")
|
||
response = input("Продолжить? (y/n): ")
|
||
if response.lower() != 'y':
|
||
print("Отменено")
|
||
sys.exit(0)
|
||
except Exception as e:
|
||
print(f"Ошибка при проверке вместимости: {e}")
|
||
sys.exit(1)
|
||
|
||
# Выполняем кодирование
|
||
print(f"Кодирование методом {args.method.upper()}...")
|
||
|
||
if args.method == 'lsb':
|
||
success = encode_lsb(args.image, message, args.output)
|
||
elif args.method == 'dct':
|
||
success = encode_dct(args.image, message, args.output)
|
||
else:
|
||
print(f"Неизвестный метод: {args.method}")
|
||
success = False
|
||
|
||
if success:
|
||
print("Кодирование завершено успешно!")
|
||
if args.psnr:
|
||
show_psnr(args.image, args.output)
|
||
else:
|
||
print("Ошибка при кодировании")
|
||
sys.exit(1)
|
||
|
||
|
||
def decode_command(args) -> None:
|
||
"""Обрабатывает команду decode."""
|
||
print(f"Декодирование методом {args.method.upper()}...")
|
||
|
||
if args.method == 'lsb':
|
||
message = decode_lsb(args.image)
|
||
elif args.method == 'dct':
|
||
message = decode_dct(args.image)
|
||
else:
|
||
print(f"Неизвестный метод: {args.method}")
|
||
sys.exit(1)
|
||
|
||
if message:
|
||
if args.output:
|
||
save_message_to_file(args.output, message)
|
||
else:
|
||
print("\n" + "=" * 50)
|
||
print("Извлеченное сообщение:")
|
||
print("=" * 50)
|
||
print(message)
|
||
print("=" * 50)
|
||
else:
|
||
print("Не удалось извлечь сообщение или сообщение пустое")
|
||
sys.exit(1)
|
||
|
||
|
||
def main() -> None:
|
||
"""Главная функция: парсит аргументы и запускает команду."""
|
||
parser = argparse.ArgumentParser(
|
||
description="Steganography tool: hide data in images using LSB and DCT",
|
||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||
epilog="""
|
||
Примеры:
|
||
%(prog)s encode lsb image.png result.png -m "Secret message"
|
||
%(prog)s encode lsb image.png result.png -f secret.txt
|
||
%(prog)s decode lsb result.png
|
||
%(prog)s decode lsb result.png -o extracted.txt
|
||
%(prog)s capacity image.png
|
||
%(prog)s psnr original.png modified.png
|
||
"""
|
||
)
|
||
|
||
subparsers = parser.add_subparsers(dest='command', help='Команды')
|
||
|
||
# Команда encode
|
||
encode_parser = subparsers.add_parser('encode', help='Скрыть сообщение')
|
||
encode_parser.add_argument('method', choices=['lsb', 'dct'],
|
||
help='Метод стеганографии')
|
||
encode_parser.add_argument('image', help='Путь к исходному изображению')
|
||
encode_parser.add_argument('output', help='Путь для сохранения результата')
|
||
encode_parser.add_argument('-m', '--message', help='Текст сообщения')
|
||
encode_parser.add_argument('-f', '--file', help='Файл с сообщением')
|
||
encode_parser.add_argument('--psnr', action='store_true',
|
||
help='Показать PSNR после кодирования')
|
||
|
||
# Команда decode
|
||
decode_parser = subparsers.add_parser('decode', help='Извлечь сообщение')
|
||
decode_parser.add_argument('method', choices=['lsb', 'dct'],
|
||
help='Метод стеганографии')
|
||
decode_parser.add_argument('image', help='Путь к изображению')
|
||
decode_parser.add_argument('-o', '--output', help='Сохранить в файл')
|
||
|
||
# Команда capacity
|
||
capacity_parser = subparsers.add_parser('capacity',
|
||
help='Показать вместимость')
|
||
capacity_parser.add_argument('image', help='Путь к изображению')
|
||
|
||
# Команда psnr
|
||
psnr_parser = subparsers.add_parser('psnr',
|
||
help='Сравнить изображения')
|
||
psnr_parser.add_argument('original', help='Исходное изображение')
|
||
psnr_parser.add_argument('modified', help='Измененное изображение')
|
||
|
||
args = parser.parse_args()
|
||
|
||
if args.command == 'encode':
|
||
encode_command(args)
|
||
elif args.command == 'decode':
|
||
decode_command(args)
|
||
elif args.command == 'capacity':
|
||
show_capacity(args.image)
|
||
elif args.command == 'psnr':
|
||
show_psnr(args.original, args.modified)
|
||
else:
|
||
parser.print_help()
|
||
sys.exit(1)
|
||
|
||
|
||
if __name__ == '__main__':
|
||
main() |