mirror of
https://github.com/NohamR/OqeeRewind.git
synced 2026-08-26 10:12:08 +00:00
ae5035c6a8
- Switch bruteforce probes from GET to HEAD to avoid downloading full segment bodies (only the status code is needed) - Cap bruteforce concurrency with a semaphore (default 100) instead of issuing one task per tick, and cancel leftover tasks on early exit - Reuse a single aiohttp session / event loop for the manifest mode's AVC/HEVC/audio bruteforces to keep TCP connections alive - Hoist the per-request segment headers into a module-level constant - Use a set for the already-downloaded resume check in save_segments - Replace read()-into-memory copies with shutil.copyfileobj in merge_segments and decrypt
146 lines
4.6 KiB
Python
146 lines
4.6 KiB
Python
"""Utility functions for OqeeRewind, including verification, merging, and decryption."""
|
|
|
|
import os
|
|
import sys
|
|
import subprocess
|
|
import shutil
|
|
from utils.logging_config import logger
|
|
|
|
|
|
def verify_mp4ff():
|
|
"""Verify if mp4ff-decrypt is installed and available in PATH."""
|
|
if shutil.which("mp4ff-decrypt") is None:
|
|
logger.error("mp4ff-decrypt is not installed or not in PATH.")
|
|
logger.info("Please install it using:")
|
|
logger.info("go install github.com/Eyevinn/mp4ff/cmd/mp4ff-decrypt@latest")
|
|
sys.exit(1)
|
|
return True
|
|
|
|
|
|
def verify_cmd(path: str) -> bool:
|
|
"""Verify if the file provided at path is valid and exists, otherwise log error and exit."""
|
|
if not os.path.exists(path):
|
|
logger.error("File does not exist: %s", path)
|
|
sys.exit(1)
|
|
if not os.path.isfile(path):
|
|
logger.error("Path is not a file: %s", path)
|
|
sys.exit(1)
|
|
return True
|
|
|
|
|
|
def merge_segments(input_folder: str, track_id: str, output_file: str):
|
|
"""Merge downloaded segments into a single file using ffmpeg."""
|
|
segment_folder = os.path.join(input_folder, f"segments_{track_id}")
|
|
|
|
segment_files = sorted(
|
|
[f for f in os.listdir(segment_folder) if f.endswith(".m4s")],
|
|
key=lambda x: int(x.split(".")[0]),
|
|
)
|
|
with open(output_file, "wb") as outfile:
|
|
for fname in segment_files:
|
|
with open(f"{segment_folder}/{fname}", "rb") as infile:
|
|
shutil.copyfileobj(infile, outfile)
|
|
logger.info("Merged segments into %s", output_file)
|
|
|
|
|
|
def decrypt(segment_dir, init_path, output_file, key):
|
|
"""Decrypt segments in chunks of ~1GB to avoid loading entire file in memory.
|
|
|
|
Args:
|
|
segment_dir: Path to the directory containing .m4s segment files.
|
|
init_path: Path to the initialization file.
|
|
output_file: Path to the output decrypted file.
|
|
key: The decryption key in KID:KEY format.
|
|
|
|
Returns:
|
|
True if decryption succeeded, False otherwise.
|
|
"""
|
|
key = key.split(":")[1]
|
|
|
|
segment_files = sorted(
|
|
[f for f in os.listdir(segment_dir) if f.endswith(".m4s")],
|
|
key=lambda x: int(x.split(".")[0]),
|
|
)
|
|
|
|
if not segment_files:
|
|
logger.error("No segment files found in %s", segment_dir)
|
|
return False
|
|
|
|
logger.info(
|
|
"Decrypting %d segments from %s to %s",
|
|
len(segment_files),
|
|
segment_dir,
|
|
output_file,
|
|
)
|
|
|
|
chunk_num = 0
|
|
temp_files = []
|
|
TARGET_SIZE = 1 * 1024 * 1024 * 1024 # 1GB
|
|
|
|
i = 0
|
|
while i < len(segment_files):
|
|
chunk_files = []
|
|
chunk_size = 0
|
|
|
|
while i < len(segment_files) and chunk_size < TARGET_SIZE:
|
|
fname = segment_files[i]
|
|
fpath = os.path.join(segment_dir, fname)
|
|
fsize = os.path.getsize(fpath)
|
|
chunk_files.append(fname)
|
|
chunk_size += fsize
|
|
i += 1
|
|
|
|
logger.debug(
|
|
"Processing chunk %d: %d segments, %.2f MB",
|
|
chunk_num,
|
|
len(chunk_files),
|
|
chunk_size / (1024 * 1024),
|
|
)
|
|
|
|
chunk_merged = os.path.join(segment_dir, f"chunk_{chunk_num}_merged")
|
|
chunk_dec = os.path.join(segment_dir, f"chunk_{chunk_num}_dec")
|
|
temp_files.extend([chunk_merged, chunk_dec])
|
|
|
|
with open(chunk_merged, "wb") as outfile:
|
|
for fname in chunk_files:
|
|
with open(os.path.join(segment_dir, fname), "rb") as infile:
|
|
shutil.copyfileobj(infile, outfile)
|
|
|
|
logger.debug("Decrypting chunk %d", chunk_num)
|
|
result = subprocess.run(
|
|
[
|
|
"mp4ff-decrypt",
|
|
"-init",
|
|
init_path,
|
|
"-key",
|
|
key,
|
|
chunk_merged,
|
|
chunk_dec,
|
|
],
|
|
capture_output=True,
|
|
text=True,
|
|
check=False,
|
|
)
|
|
if result.returncode != 0:
|
|
logger.error("Decryption failed for chunk %d: %s", chunk_num, result.stderr)
|
|
for tf in temp_files:
|
|
if os.path.exists(tf):
|
|
os.remove(tf)
|
|
return False
|
|
|
|
chunk_num += 1
|
|
|
|
logger.debug("Concatenating %d decrypted chunks into %s", chunk_num, output_file)
|
|
with open(output_file, "wb") as outfile:
|
|
for c in range(chunk_num):
|
|
chunk_dec = os.path.join(segment_dir, f"chunk_{c}_dec")
|
|
with open(chunk_dec, "rb") as infile:
|
|
outfile.write(infile.read())
|
|
|
|
for tf in temp_files:
|
|
if os.path.exists(tf):
|
|
os.remove(tf)
|
|
|
|
logger.info("Successfully decrypted to %s", output_file)
|
|
return True
|