---

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

ENCRYPTED_FILE = "encrypted.avi"  
DECRYPTED_FILE = "decrypted.avi"  
SETTINGS_FILE = "setting.txt" 

def find_input_file():
    """Автоматически ищет файл object с любым расширением (.mp4, .png, .jpg и т.д.)"""
    supported_extensions = ['.mp4', '.avi', '.mov', '.mkv', '.png', '.jpg', '.jpeg', '.webp']
    for ext in supported_extensions:
        filename = f"object{ext}"
        if os.path.exists(filename):
            return filename
    return None

def generate_token(length=16):
    return secrets.token_hex(length)

def get_key_and_iv(token: str):
    key = hashlib.sha256(token.encode('utf-8')).digest()
    iv = hashlib.md5(token.encode('utf-8')).digest()
    return key, iv

def get_media_resolution(filepath: str):
    """Автоматическое получение разрешения (ширина x высота) для видео или картинки"""
    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_media(in_path: str, out_path: str, token: str):
    key, iv = get_key_and_iv(token)
    width, height = get_media_resolution(in_path)
    frame_size = width * height * 3  # RGB24 = 3 байта на пиксель

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

    # 2. Процесс записи обработанного потока в безпотерьный формат
    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}]...")

    # Для картинки ffmpeg с бесконечным лупом будет слать кадры вечно, 
    # поэтому для картинок ограничимся, например, 30 кадрами (1 секунда шума)
    # Если это исходно видео, он обработает все кадры видео.
    is_image = in_path.lower().endswith(('.png', '.jpg', '.jpeg', '.webp'))
    max_frames = 30 if is_image else 999999999

    while frame_index < max_frames:
        raw_frame = ff_read.stdout.read(frame_size)
        if not raw_frame:
            break

        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)
        ff_write.stdin.write(processed_frame)
        frame_index += 1

    # Принудительно завершаем чтение картинки, так как у нее нет конца
    ff_read.terminate()
    ff_write.stdin.close()
    ff_read.wait()
    ff_write.wait()
    print(f"[✓] Готово! Обработано кадров: {frame_index}")

def encrypt_mode():
    input_file = find_input_file()
    if not input_file:
        print(f"[!] Не найден файл 'object' с поддерживаемым расширением (.mp4, .png, .jpg и т.д.) в текущей папке!")
        return

    print(f"[+] Найден файл для шифрования: {input_file}")

    token = generate_token()
    with open(SETTINGS_FILE, "w", encoding="utf-8") as f:
        f.write(token)
    print(f"[+] Новый токен записан в {SETTINGS_FILE}: {token}")

    process_media(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

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

    process_media(ENCRYPTED_FILE, DECRYPTED_FILE, token)

if __name__ == "__main__":
    print("Выберите режим:")
    print("1 — Зашифровать (object.* -> 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("[!] Неверный выбор.")