---

import os
import sys
import secrets
import hashlib
import subprocess
from Crypto.Cipher import AES
from Crypto.Util import Counter

# Файлы по умолчанию
INPUT_FILE = "object.mp4"      # Исходный файл (может быть .mp4, .png, .jpg, .avi и т.д.)
ENCRYPTED_FILE = "encrypted.avi"  # Файл с шифрованным видео/шумом
DECRYPTED_FILE = "decrypted.avi"  # Восстановленный файл
SETTINGS_FILE = "setting.txt" # Файл для хранения токена

def generate_token(length=16):
    """Генерация безопасного рандомного токена"""
    return secrets.token_hex(length)

def get_key_and_iv(token: str):
    """Получение 256-битного ключа и 128-битного IV из текстового токена"""
    key = hashlib.sha256(token.encode('utf-8')).digest()
    iv = hashlib.md5(token.encode('utf-8')).digest()
    return key, iv

def get_video_resolution(filepath: str):
    """Автоматическое получение разрешения видео (ширина x высота) через ffprobe"""
    cmd = [
        'ffprobe', '-v', 'error',
        '-select_streams', 'v:0',
        '-show_entries', 'stream=width,height',
        '-of', 'csv=s=x:p=0',
        filepath
    ]
    try:
        output = subprocess.check_output(cmd).decode('utf-8').strip()
        width, height = map(int, output.split('x'))
        return width, height
    except Exception as e:
        print(f"[!] Ошибка определения разрешения файла {filepath}: {e}")
        sys.exit(1)

def process_video(in_path: str, out_path: str, token: str):
    """Шифрование / Дешифрование потока через FFmpeg + AES-CTR"""
    key, iv = get_key_and_iv(token)
    width, height = get_video_resolution(in_path)
    frame_size = width * height * 3  # RGB24 = 3 байта на пиксель

    # 1. Процесс чтения исходного видео в формате сырых пикселей (rawvideo RGB24)
    ff_read = subprocess.Popen([
        'ffmpeg', '-y', '-i', in_path,
        '-f', 'image2pipe', '-pix_fmt', 'rgb24', '-vcodec', 'rawvideo', '-'
    ], stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, bufsize=10**8)

    # 2. Процесс записи обработанного потока пикселей в безпотерьный формат (FFV1)
    ff_write = subprocess.Popen([
        'ffmpeg', '-y',
        '-f', 'rawvideo', '-pix_fmt', 'rgb24', '-s', f'{width}x{height}', '-r', '30',
        '-i', '-',
        '-c:v', 'ffv1', out_path
    ], stdin=subprocess.PIPE, stderr=subprocess.DEVNULL)

    frame_index = 0
    print(f"Обработка [{in_path} -> {out_path}]...")

    while True:
        raw_frame = ff_read.stdout.read(frame_size)
        if not raw_frame:
            break

        # Настраиваем счетчик AES-CTR для текущего кадра
        ctr_int = int.from_bytes(iv, 'big') + frame_index
        ctr = Counter.new(128, initial_value=ctr_int)
        cipher = AES.new(key, AES.MODE_CTR, counter=ctr)

        # Шифруем / расшифровываем байты кадра
        processed_frame = cipher.encrypt(raw_frame)

        # Записываем результат в выходной процесс FFmpeg
        ff_write.stdin.write(processed_frame)
        frame_index += 1

    ff_read.stdout.close()
    ff_write.stdin.close()
    ff_read.wait()
    ff_write.wait()
    print(f"[✓] Готово! Обработано кадров: {frame_index}")

def encrypt_mode():
    if not os.path.exists(INPUT_FILE):
        print(f"[!] Файл {INPUT_FILE} не найден!")
        return

    # 1. Генерируем токен и сохраняем в setting.txt
    token = generate_token()
    with open(SETTINGS_FILE, "w", encoding="utf-8") as f:
        f.write(token)
    print(f"[+] Новый токен сгенерирован и записан в {SETTINGS_FILE}: {token}")

    # 2. Запускаем шифрование
    process_video(INPUT_FILE, ENCRYPTED_FILE, token)

def decrypt_mode():
    if not os.path.exists(ENCRYPTED_FILE):
        print(f"[!] Зашифрованный файл {ENCRYPTED_FILE} не найден!")
        return

    if not os.path.exists(SETTINGS_FILE):
        print(f"[!] Файл настроек {SETTINGS_FILE} не найден!")
        return

    # 1. Считываем токен
    with open(SETTINGS_FILE, "r", encoding="utf-8") as f:
        token = f.read().strip()
    print(f"[+] Использован токен из {SETTINGS_FILE}: {token}")

    # 2. Запускаем дешифрование
    process_video(ENCRYPTED_FILE, DECRYPTED_FILE, token)

if __name__ == "__main__":
    print("Выберите режим:")
    print("1 — Зашифровать (object.mp4 -> encrypted.avi + создать setting.txt)")
    print("2 — Расшифровать (encrypted.avi -> decrypted.avi по токену из setting.txt)")
    
    choice = input("Введите цифру (1 или 2): ").strip()
    
    if choice == "1":
        encrypt_mode()
    elif choice == "2":
        decrypt_mode()
    else:
        print("[!] Неверный выбор.")