mediafire.py

#!/usr/bin/env python3
import asyncio
import hashlib
import json
from re import findall
from typing import Tuple

import httpx
from fastapi import Body, FastAPI
from fastapi.responses import JSONResponse
from gazpacho import Soup
from gazpacho.utils import HTTPError
from pydantic import AnyUrl, BaseModel, Field

from sse_starlette import EventSourceResponse, ServerSentEvent
from starlette.exceptions import HTTPException

app = FastAPI()


class bcolors:
    HEADER = "\033[95m"
    OKBLUE = "\033[94m"
    OKCYAN = "\033[96m"
    OKGREEN = "\033[92m"
    WARNING = "\033[93m"
    FAIL = "\033[91m"
    ENDC = "\033[0m"
    BOLD = "\033[1m"
    UNDERLINE = "\033[4m"


# Non-alphanumeric (str.isalphanum()) characters allowed in a file or folder name
NON_ALPHANUM_FILE_OR_FOLDER_NAME_CHARACTERS = "-_. "
# What to replace bad characters with.
NON_ALPHANUM_FILE_OR_FOLDER_NAME_CHARACTER_REPLACEMENT = "-"


def hash_file(filename: str) -> str:
    """
    Calculate the SHA-256 hash digest of a file.

    Args:
        filename (str): The path to the file.

    Returns:
        str: The hexadecimal representation of the hash digest.

    Raises:
        FileNotFoundError: If the specified file does not exist.
        PermissionError: If the user does not have permission to read the file.
    """
    # make a hash object
    h = hashlib.sha256()

    # open file for reading in binary mode
    with open(filename, "rb") as file:
        # loop till the end of the file
        chunk = 0
        while chunk != b"":
            # read only 1024 bytes at a time
            chunk = file.read(1024)
            h.update(chunk)

    # return the hex representation of digest
    return h.hexdigest()


def normalize_file_or_folder_name(filename: str) -> str:
    """
    Normalize a file or folder name by replacing non-alphanumeric characters.

    Args:
        filename (str): The original file or folder name.

    Returns:
        str: The normalized file or folder name.

    Note:
        If you want to disable normalization, uncomment the return statement at the beginning of the function.

    Example:
        >>> normalize_file_or_folder_name("my_file$%")
        'my_file_'

    """
    return "".join(
        [
            char
            if (char.isalnum() or char in NON_ALPHANUM_FILE_OR_FOLDER_NAME_CHARACTERS)
            else NON_ALPHANUM_FILE_OR_FOLDER_NAME_CHARACTER_REPLACEMENT
            for char in filename
        ]
    )


def print_error(link: str):
    """
    Prints an error message indicating that a file has been deleted or blocked due to being dangerous.

    Parameters:
        link (str): The link to the file or resource that caused the error.

    Returns:
        None

    Example:
        >>> print_error("https://example.com/dangerous_file.txt")
        Deleted file or Dangerous File Blocked
        Take a look if you want to be sure: https://example.com/dangerous_file.txt
    """
    print(
        f"{bcolors.FAIL}Deleted file or Dangerous File Blocked\n"
        f"{bcolors.WARNING}Take a look if you want to be sure: {link}{bcolors.ENDC}"
    )


class PostModel(BaseModel):
    url: str = Field(..., examples=["https://www.mediafire.com/file/example_file.txt"])


tasks = set()

from cycurl.requests import AsyncSession


@app.post("/parse", description="解析下载链接")
async def parse(data: PostModel = Body(..., description="Mediafire链接")):
    """
    Mediafire Bulk Downloader

    Parses command-line arguments to download files or folders from Mediafire.

    Usage:
        python mediafire.py <mediafire_url> [-o <output_path>] [-t <num_threads>]

    Arguments:
        mediafire_url (str): The URL of the file or folder to be downloaded from Mediafire.

    Options:
        -o, --output (str): The path of the desired output folder. Default is the current directory.
        -t, --threads (int): Number of threads to use for downloading. Default is 10.

    Returns:
        None

    Example:
        To download a file:
        $ python mediafire.py https://www.mediafire.com/file/example_file.txt

        To download a folder:
        $ python mediafire.py https://www.mediafire.com/folder/example_folder -o /path/to/output -t 20
    """
    # parser = ArgumentParser(
    #     "mediafire_bulk_downloader", usage="python mediafire.py <mediafire_url>"
    # )
    # parser.add_argument(
    #     "mediafire_url", help="The URL of the file or folder to be downloaded"
    # )
    # parser.add_argument(
    #     "-o",
    #     "--output",
    #     help="The path of the desired output folder",
    #     required=False,
    #     default=".",
    # )
    # parser.add_argument(
    #     "-t",
    #     "--threads",
    #     help="Number of threads to use",
    #     type=int,
    #     default=10,
    #     required=False,
    # )

    # args = parser.parse_args()

    folder_or_file = findall(
        r"mediafire\.com/(folder|file|file_premium)\/([a-zA-Z0-9]+)", data.url
    )

    if not folder_or_file:
        return JSONResponse({"error": f"Invalid link"}, 403)

    t, key = folder_or_file[0]

    if t in {"file", "file_premium"}:
        name, url = await get_file(key)
        return JSONResponse({"name": name, "url": url}, 200)
    elif t == "folder":
        queue = asyncio.Queue()
        task = asyncio.create_task(get_folders(key, "", queue))
        tasks.add(task)
        task.add_done_callback(tasks.discard)

        async def event_resp(q: asyncio.Queue):
            while True:
                folder, filename, url = await q.get()
                yield ServerSentEvent(
                    data=json.dumps(
                        {"folder": folder, "filename": filename, "url": url}
                    )
                )

        return EventSourceResponse(event_resp(queue))
    else:
        return JSONResponse({"error": f"Invalid link"}, 403)


def get_files_or_folders_api_endpoint(
    filefolder: str, folder_key: str, chunk: int = 1, info: bool = False
) -> str:
    """
    Constructs the API endpoint URL for retrieving files or folders information from Mediafire.

    Parameters:
        filefolder (str): Type of content to retrieve. Either 'file' or 'folder'.
        folder_key (str): The unique identifier of the folder.
        chunk (int): The chunk number to retrieve. Default is 1.
        info (bool): If True, gets folder info; otherwise, gets content. Default is False.

    Returns:
        str: The constructed API endpoint URL.

    Example:
        >>> get_files_or_folders_api_endpoint('folder', 'folder_key_123', chunk=2, info=True)
        'https://www.mediafire.com/api/1.4/folder/get_info.php?r=utga&content_type=folder&filter=all&order_by=name&order_direction=asc&chunk=2&version=1.5&folder_key=folder_key_123&response_format=json'
    """
    return (
        f"https://www.mediafire.com/api/1.4/folder"
        f"/{'get_info' if info else 'get_content'}.php?r=utga&content_type={filefolder}"
        f"&filter=all&order_by=name&order_direction=asc&chunk={chunk}"
        f"&version=1.5&folder_key={folder_key}&response_format=json"
    )


def get_info_endpoint(file_key: str) -> str:
    """
    Constructs the API endpoint URL for retrieving information about a specific file from Mediafire.

    Parameters:
        file_key (str): The unique identifier of the file.

    Returns:
        str: The constructed API endpoint URL.

    Example:
        >>> get_info_endpoint('file_key_123')
        'https://www.mediafire.com/api/file/get_info.php?quick_key=file_key_123&response_format=json'
    """
    return f"https://www.mediafire.com/api/file/get_info.php?quick_key={file_key}&response_format=json"


async def get_folders(folder_key: str, folder_name: str, queue: asyncio.Queue) -> None:
    """
    Recursively downloads folders and files from Mediafire.

    Parameters:
        folder_key (str): The unique identifier of the folder.
        folder_name (str): The name of the folder.

    Returns:
        None

    Example:
        >>> get_folders('folder_key_123', '/path/to/download', 5, first=True)
    """
    # if first:
    #     folder_name = path.join(
    #         folder_name,
    #         normalize_file_or_folder_name(
    #             get(
    #                 get_files_or_folders_api_endpoint("folder", folder_key, info=True)
    #             ).json()["response"]["folder_info"]["name"]
    #         ),
    #     )

    # If the folder doesn't exist, create and enter it
    # if not path.exists(folder_name):
    #     makedirs(folder_name)
    # chdir(folder_name)

    # Downloading all the files in the main folder
    await download_folder(folder_key, folder_name, queue)

    # Searching for other folders
    async with httpx.AsyncClient() as client:
        resp = await client.get(
            get_files_or_folders_api_endpoint("folders", folder_key)
        )
        folder_content = resp.json()["response"]["folder_content"]
        # folder_content = get(
        #     get_files_or_folders_api_endpoint("folders", folder_key)
        # ).json()["response"]["folder_content"]

        # Downloading other folders recursively
        if "folders" in folder_content:
            for folder in folder_content["folders"]:
                await get_folders(folder["folderkey"], folder["name"], queue)
                # chdir("..")


async def download_folder(folder_key: str, folder_name, queue: asyncio.Queue) -> None:
    """
    Downloads all files from a Mediafire folder.

    Parameters:
        folder_key (str): The unique identifier of the folder.
        threads_num (int): Number of threads to use for downloading.

    Returns:
        None

    Example:
        >>> download_folder('folder_key_123', 5)
    """
    # Getting all the files
    data = []
    chunk = 1
    more_chunks = True

    async with httpx.AsyncClient() as client:
        try:
            # If there are more than 100 files, make another request
            # and append the result to data
            while more_chunks:
                resp = await client.get(
                    get_files_or_folders_api_endpoint("files", folder_key, chunk=chunk)
                )
                r_json = resp.json()
                # r_json = get(
                #     get_files_or_folders_api_endpoint("files", folder_key, chunk=chunk)
                # ).json()
                more_chunks = (
                    r_json["response"]["folder_content"]["more_chunks"] == "yes"
                )
                data += r_json["response"]["folder_content"]["files"]
                chunk += 1

        except KeyError:
            print("Invalid link")
            return
    for file in data:
        filename, url = await get_file_link(file)
        await queue.put((folder_name, filename, url))
    # event = Event()
    # threadLimiter = BoundedSemaphore(threads_num)
    # total_threads: list[Thread] = []

    # Appending a new thread for downloading every link
    # for file in data:
    #     total_threads.append(
    #         Thread(
    #             target=get_file_link,
    #             args=(
    #                 file,
    #                 event,
    #                 threadLimiter,
    #             ),
    #         )
    #     )

    # Starting all threads
    # for thread in total_threads:
    #     thread.start()

    # Handle being interrupted
    # try:
    #     while True:
    #         if all(not t.is_alive() for t in total_threads):
    #             break
    #         sleep(0.01)
    # except KeyboardInterrupt:
    #     print(f"{bcolors.WARNING}Closing all threads{bcolors.ENDC}")
    #     event.set()
    #     for thread in total_threads:
    #         thread.join()
    #     print(f"{bcolors.WARNING}{bcolors.BOLD}Download interrupted{bcolors.ENDC}")
    #     exit(0)


async def get_file(key: str) -> Tuple[str, str]:
    """
    Downloads a single file from Mediafire using the main thread.

    Parameters:
        key (str): The unique identifier of the file.
        output_path (str): The path where the file will be downloaded. If None, the current directory is used.

    Returns:
        None

    Example:
        >>> get_file('file_key_123', '/path/to/download')
    """
    # Retrieve file information
    async with httpx.AsyncClient() as client:
        resp = await client.get(get_info_endpoint(key))
        file_data = resp.json()["response"]["file_info"]
    # file_data = get(get_info_endpoint(key)).json()["response"]["file_info"]

    # Change directory if output_path is provided
    # if output_path:
    #     chdir(output_path)

    # Download the file
    return await get_file_link(file_data)


async def get_file_link(file: dict) -> Tuple[str, str]:
    """
    Downloads a file from a direct link obtained from Mediafire.

    Parameters:
        file (dict): A dictionary containing file information, including the direct download link.
        event (Event): An optional threading event used for handling interruptions.
        limiter (BoundedSemaphore): An optional semaphore for controlling the number of concurrent downloads.

    Returns:
        None

    Example:
        >>> get_file_link({'filename': 'example_file.txt', 'links': {'normal_download': 'https://www.mediafire.com/download/example_file.txt'}})
    """
    # Acquire semaphore if available
    # if limiter:
    #     limiter.acquire()

    # Extract direct download link from file information
    download_link = file["links"]["normal_download"]

    # Normalize filename
    filename = normalize_file_or_folder_name(file["filename"])

    # Check if file already exists and is not corrupted
    # if path.exists(filename):
    #     if hash_file(filename) == file["hash"]:
    #         print(f"{bcolors.WARNING}{filename}{bcolors.ENDC} already exists, skipping")
    #         if limiter:
    #             limiter.release()
    #         return
    #     else:
    #         print(
    #             f"{bcolors.WARNING}{filename}{bcolors.ENDC} already exists but corrupted, downloading again"
    #         )

    # Start downloading the file
    # print(f"{bcolors.OKBLUE}Downloading {filename}{bcolors.ENDC}")

    # if event:
    #     if event.is_set():
    #         if limiter:
    #             limiter.release()
    #         return

    # Check if the link is not a direct download link and extract the actual download link
    async with httpx.AsyncClient() as client:
        try:
            resp = await client.head(download_link)
            if resp.headers.get("content-encoding") == "gzip":
            # if head(download_link).headers.get("content-encoding") == "gzip":
                # Retrieve the HTML content of the file link
                resp = await client.get(download_link)
                html = resp.text
                # html = get(download_link).text
                # Parse HTML content to extract the actual download link
                soup = Soup(html)
                download_link = (
                    soup.find("div", {"class": "download_link"})
                    .find("a", {"class": "input popsok"})
                    .attrs["href"]
                )
                return filename, download_link
            else:
                raise HTTPException(status_code=404, detail=f"error parse {download_link}, refer {resp.headers.get('location')}")
        except Exception:
            # Handle HTTP errors
            raise HTTPException(status_code=404, detail=f"error parse {download_link}")
            # print_error(download_link)
            # if limiter:
            #     limiter.release()
            # return

    # Download file in chunks
    # with get(download_link, stream=True) as r:
    #     r.raise_for_status()
    #     with open(filename, "wb") as f:
    #         for chunk in r.iter_content(chunk_size=4096):
    #             if event:
    #                 if event.is_set():
    #                     break
    #             if chunk:
    #                 f.write(chunk)

    # Check if download was interrupted
    # if event:
    #     if event.is_set():
    #         remove(filename)
    #         print(
    #             f"{bcolors.WARNING}Partially downloaded {filename} deleted{bcolors.ENDC}"
    #         )
    #         if limiter:
    #             limiter.release()
    #         return

    # Print download success message
    # print(f"{bcolors.OKGREEN}{filename}{bcolors.ENDC} downloaded")

    # Release semaphore if acquired
    # if limiter:
    #     limiter.release()


if __name__ == "__main__":
    from hypercorn.asyncio import serve
    from hypercorn.config import Config

    conf = Config()
    conf.bind = "0.0.0.0:8000"

    asyncio.run(serve(app, conf))