From 8ce68479ae2bc8f9de6f23e3b910659cc26b441f Mon Sep 17 00:00:00 2001 From: Lovi <62809003+Lovi-0@users.noreply.github.com> Date: Thu, 13 Feb 2025 10:41:59 +0100 Subject: [PATCH] Fix signal handler mp4 --- .../Lib/Downloader/MP4/downloader.py | 82 ++++++++++--------- setup.py | 2 +- 2 files changed, 45 insertions(+), 39 deletions(-) diff --git a/StreamingCommunity/Lib/Downloader/MP4/downloader.py b/StreamingCommunity/Lib/Downloader/MP4/downloader.py index 6f64f34..86c921b 100644 --- a/StreamingCommunity/Lib/Downloader/MP4/downloader.py +++ b/StreamingCommunity/Lib/Downloader/MP4/downloader.py @@ -3,6 +3,7 @@ import os import re import sys +import time import signal import logging from functools import partial @@ -39,21 +40,38 @@ TELEGRAM_BOT = config_manager.get_bool('DEFAULT', 'telegram_bot') -def signal_handler(signum, frame, kill_handler): - """Signal handler for graceful interruption""" - kill_handler[0] = True - print("\nReceived interrupt signal. Completing current download...") +class InterruptHandler: + def __init__(self): + self.interrupt_count = 0 + self.last_interrupt_time = 0 + self.kill_download = False + self.force_quit = False +def signal_handler(signum, frame, interrupt_handler, original_handler): + """Enhanced signal handler for multiple interrupt scenarios""" + current_time = time.time() + + # Reset counter if more than 2 seconds have passed since last interrupt + if current_time - interrupt_handler.last_interrupt_time > 2: + interrupt_handler.interrupt_count = 0 + + interrupt_handler.interrupt_count += 1 + interrupt_handler.last_interrupt_time = current_time + + if interrupt_handler.interrupt_count == 1: + interrupt_handler.kill_download = True + console.print("\n[bold yellow]First interrupt received. Download will complete and save. Press Ctrl+C three times quickly to force quit.[/bold yellow]") + + elif interrupt_handler.interrupt_count >= 3: + interrupt_handler.force_quit = True + console.print("\n[bold red]Force quit activated. Saving partial download...[/bold red]") + signal.signal(signum, original_handler) def MP4_downloader(url: str, path: str, referer: str = None, headers_: dict = None): """ - Downloads an MP4 video from a given URL with robust error handling and SSL bypass. - - Parameters: - - url (str): The URL of the MP4 video to download. - - path (str): The local path where the downloaded MP4 file will be saved. - - referer (str, optional): The referer header value. - - headers_ (dict, optional): Custom headers for the request. + Downloads an MP4 video with enhanced interrupt handling. + - Single Ctrl+C: Completes download gracefully + - Triple Ctrl+C: Saves partial download and exits """ if TELEGRAM_BOT: bot = get_bot_instance() @@ -65,23 +83,19 @@ def MP4_downloader(url: str, path: str, referer: str = None, headers_: dict = No bot.send_message(f"Contenuto già scaricato!", None) return 400 - # Early return for link-only mode if GET_ONLY_LINK: return {'path': path, 'url': url} - # Validate URL if not (url.lower().startswith('http://') or url.lower().startswith('https://')): logging.error(f"Invalid URL: {url}") console.print(f"[bold red]Invalid URL: {url}[/bold red]") return None - # Prepare headers try: headers = {} if referer: headers['Referer'] = referer - # Use custom headers if provided, otherwise use default user agent if headers_: headers.update(headers_) else: @@ -93,17 +107,12 @@ def MP4_downloader(url: str, path: str, referer: str = None, headers_: dict = No return None temp_path = f"{path}.temp" - kill_handler = [False] # Using list for mutable state - original_handler = signal.signal(signal.SIGINT, partial(signal_handler, kill_handler=kill_handler)) + interrupt_handler = InterruptHandler() + original_handler = signal.signal(signal.SIGINT, partial(signal_handler, interrupt_handler=interrupt_handler, original_handler=signal.getsignal(signal.SIGINT))) try: - # Create a custom transport that bypasses SSL verification - transport = httpx.HTTPTransport( - verify=False, - http2=True - ) + transport = httpx.HTTPTransport(verify=False, http2=True) - # Download with streaming and progress tracking with httpx.Client(transport=transport, timeout=httpx.Timeout(60)) as client: with client.stream("GET", url, headers=headers, timeout=REQUEST_TIMEOUT) as response: response.raise_for_status() @@ -125,16 +134,16 @@ def MP4_downloader(url: str, path: str, referer: str = None, headers_: dict = No unit_scale=True, desc='Downloading', mininterval=0.05, - file=sys.stdout, # Using file=sys.stdout to force in-place updates because sys.stderr may not support carriage returns in this environment. + file=sys.stdout # Using file=sys.stdout to force in-place updates because sys.stderr may not support carriage returns in this environment. ) downloaded = 0 with open(temp_path, 'wb') as file, progress_bar as bar: try: for chunk in response.iter_bytes(chunk_size=1024): - if kill_handler[0]: - console.print("\n[bold yellow]Interrupting download...[/bold yellow]") - return None, True + if interrupt_handler.force_quit: + console.print("\n[bold red]Force quitting... Saving partial download.[/bold red]") + break if chunk: size = file.write(chunk) @@ -142,18 +151,15 @@ def MP4_downloader(url: str, path: str, referer: str = None, headers_: dict = No bar.update(size) except KeyboardInterrupt: - console.print("\n[bold red]Download interrupted by user.[/bold red]") - if os.path.exists(temp_path): - os.remove(temp_path) - return None, True + if not interrupt_handler.force_quit: + interrupt_handler.kill_download = True - # Rename temp file to final file if os.path.exists(temp_path): os.rename(temp_path, path) if os.path.exists(path): console.print(Panel( - f"[bold green]Download completed![/bold green]\n" + f"[bold green]Download completed{' (Partial)' if interrupt_handler.force_quit else ''}![/bold green]\n" f"[cyan]File size: [bold red]{internet_manager.format_file_size(os.path.getsize(path))}[/bold red]\n" f"[cyan]Duration: [bold]{print_duration_table(path, description=False, return_string=True)}[/bold]", title=f"{os.path.basename(path.replace('.mp4', ''))}", @@ -161,22 +167,22 @@ def MP4_downloader(url: str, path: str, referer: str = None, headers_: dict = No )) if TELEGRAM_BOT: - message = f"Download completato\nDimensione: {internet_manager.format_file_size(os.path.getsize(path))}\nDurata: {print_duration_table(path, description=False, return_string=True)}\nTitolo: {os.path.basename(path.replace('.mp4', ''))}" + message = f"Download completato{'(Parziale)' if interrupt_handler.force_quit else ''}\nDimensione: {internet_manager.format_file_size(os.path.getsize(path))}\nDurata: {print_duration_table(path, description=False, return_string=True)}\nTitolo: {os.path.basename(path.replace('.mp4', ''))}" clean_message = re.sub(r'\[[a-zA-Z]+\]', '', message) bot.send_message(clean_message, None) - return path, kill_handler[0] + return path, interrupt_handler.kill_download + else: console.print("[bold red]Download failed or file is empty.[/bold red]") - return None, kill_handler[0] + return None, interrupt_handler.kill_download except Exception as e: logging.error(f"Unexpected error: {e}") console.print(f"[bold red]Unexpected Error: {e}[/bold red]") if os.path.exists(temp_path): os.remove(temp_path) - return None, kill_handler[0] + return None, interrupt_handler.kill_download finally: - # Restore original signal handler signal.signal(signal.SIGINT, original_handler) \ No newline at end of file diff --git a/setup.py b/setup.py index a1f3cb1..798e04b 100644 --- a/setup.py +++ b/setup.py @@ -10,7 +10,7 @@ with open("requirements.txt", "r", encoding="utf-8-sig") as f: setup( name="StreamingCommunity", - version="2.5.7", + version="2.5.8", long_description=read_readme(), long_description_content_type="text/markdown", author="Lovi-0",