main.py

import json
import os
import subprocess
import time
from typing import Optional

from pynput import keyboard
from selenium import webdriver
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.chrome.service import Service
from selenium.webdriver.common.desired_capabilities import DesiredCapabilities
from webdriver_manager.chrome import ChromeDriverManager


class M3U8Downloader:
    """
    A class responsible for detecting and downloading M3U8 video streams from a connected Chrome browser.
    """

    SAVE_DIR = "downloads"

    def __init__(self) -> None:
        """
        Initializes the M3U8Downloader by setting up the browser driver and creating the save directory.
        """
        os.makedirs(self.SAVE_DIR, exist_ok=True)
        self.driver = self._setup_browser()
        print("[+] Browser connected. Waiting for F4 command...")

    @staticmethod
    def _setup_browser() -> webdriver.Chrome:
        """
        Sets up and returns a Chrome WebDriver instance with necessary options.

        :return: Configured Chrome WebDriver instance.
        """
        chrome_options = Options()
        chrome_options.debugger_address = "localhost:9222"
        chrome_options.set_capability('goog:loggingPrefs', {'performance': 'ALL'})
        capabilities = DesiredCapabilities.CHROME
        service = Service(ChromeDriverManager().install())
        return webdriver.Chrome(service=service, options=chrome_options)

    def get_m3u8_url(self) -> Optional[str]:
        """
        Extracts an M3U8 URL from browser logs if available.

        :return: M3U8 URL if found, otherwise None.
        """
        logs = self.driver.get_log("performance")
        for entry in logs:
            log_json = json.loads(entry["message"])
            try:
                url = log_json["message"]["params"]["request"]["url"]
                if ".m3u8" in url:
                    print(f"[+] M3U8 stream found: {url}")
                    return url
            except (KeyError, TypeError):
                continue
        return None

    def download_video(self, m3u8_url: str) -> None:
        """
        Downloads the video from the given M3U8 URL using FFmpeg.

        :param m3u8_url: URL of the M3U8 stream to download.
        """
        timestamp = int(time.time())
        output_filename = os.path.join(self.SAVE_DIR, f"video_{timestamp}.mp4")

        command = [
            "ffmpeg", "-i", m3u8_url, "-c", "copy", "-bsf:a", "aac_adtstoasc", output_filename
        ]
        print(f"[+] Downloading video: {output_filename}")
        subprocess.run(command, check=True)
        print(f"[+] Video saved: {output_filename}")

    def process(self) -> None:
        """
        Initiates the process of detecting and downloading an M3U8 stream.
        """
        m3u8_url = self.get_m3u8_url()
        if m3u8_url:
            self.download_video(m3u8_url)
        else:
            print("[-] No M3U8 link found. Try again.")


def on_press(key: keyboard.Key) -> None:
    """
    Callback function triggered when a key is pressed.
    Initiates the M3U8 detection and download process when F4 is pressed.

    :param key: The key that was pressed.
    """
    if key == keyboard.Key.f4:
        print("[+] F4 pressed. Searching for M3U8 link...")
        downloader.process()


if __name__ == "__main__":
    downloader = M3U8Downloader()
    with keyboard.Listener(on_press=on_press) as listener:
        listener.join()