Files
CourseWork/main.py

220 lines
8.5 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

#!/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)
def gui_main():
"""Запуск GUI"""
from gui.main_window import main as gui_start
gui_start()
if __name__ == '__main__':
if len(sys.argv) > 1 and sys.argv[1] == '--gui':
gui_main()
else:
main()
if __name__ == '__main__':
main()