get-9anime-episode

#!/usr/bin/env python3

import argparse
import json
import multiprocessing
from dataclasses import dataclass
from typing import Self, Tuple, cast

import cloudscraper
from bs4 import BeautifulSoup, Tag
from bs4.element import ResultSet
from urllib.parse import quote_plus as encode_url_component

from python.helpers.deobfuscator import DefaultPlayerDeobfuscator
from python.helpers.main import (
    DownloadSite,
    DownloadSitesCtx,
    EpisodeNumberNotFoundException,
    downloader_main,
)
from python.log.console import Console

BASE_URL = "https://anix.to"

ALLOWED_SERIES_TYPES = {
    "softsub",
    "sub",
    "dub",
}

UNWANTED_CDN_HOSTNAMES = {
    ".netmagcdn.com",
}


def get_9anime_page_url(anime_name: str) -> str:
    return f"{BASE_URL}/watch/{anime_name}"


def get_anime_id(anime_name: str) -> str | None:
    Console.log_dim("Extracting id from name", return_line=True)
    name_parts = anime_name.split("-")
    anime_id = name_parts[-1]
    if anime_id.isdigit():
        return anime_id
    Console.log_dim("Fetching episode page...", return_line=True)
    response = cloudscraper.create_scraper().get(
        get_9anime_page_url(anime_name),
    )

    if not response.status_code:
        Console.log_dim("Anime not found")
        return None

    page_html = response.text

    Console.log_dim("Parsing HTML...", return_line=True)

    try:
        info = BeautifulSoup(
            page_html,
            "html.parser",
        ).select_one("div[data-id]")

        if not info:
            return None

        return info.attrs["data-id"]
    except Exception:
        return None


@dataclass
class EpisodeInfo:
    number: float
    slug: str
    has_sub: bool
    has_dub: bool
    ids: str


def get_page(url: str):
    return cloudscraper.create_scraper().get(
        url,
        headers={
            "Accept": "*/*",
            "X-Requested-With": "XMLHttpRequest",
        },
    )


def get_site_api_with_vrf(url_base: str, last_segment: str):
    vrf_token = get_encoded_vrf_token(last_segment)
    if not vrf_token:
        return None
    return get_page(
        f"{BASE_URL.rstrip("/")}/{url_base.strip("/")}/{last_segment}?vrf={vrf_token}"
    )


def get_encoded_vrf_token(data: str) -> str | None:
    vrf_resp = DefaultPlayerDeobfuscator.get(
        f"/vrf/9anime?vrf_data={encode_url_component(data)}",
        validate_status=False,
    )

    if vrf_resp is None:
        return None

    if "message" in vrf_resp:
        Console.log_dim(f"Failed to get decoded data: {vrf_resp['message']}")
        return None

    try:
        token = str(vrf_resp["vrf"])
        return encode_url_component(token)
    except Exception:
        Console.log_dim("Failed to get VRF token")
        return None


def get_decoded_vrf_data(data: str) -> str | None:
    vrf_resp = DefaultPlayerDeobfuscator.get(
        f"/vrf/9anime?devrf_data={encode_url_component(data)}",
        validate_status=False,
    )

    if vrf_resp is None:
        return None

    if "message" in vrf_resp:
        Console.log_dim(f"Failed to get decoded data: {vrf_resp['message']}")
        return None

    return str(vrf_resp["data"])


def get_anime_episode_list(anime_id: str) -> dict[float, EpisodeInfo]:
    Console.log_dim("Fetching episode list...", return_line=True)
    response = get_site_api_with_vrf(
        "/ajax/episode/list/",
        anime_id,
    )

    if not response:
        Console.log_dim("No episodes found")
        return {}

    page_json = response.json()

    elements: ResultSet[Tag] = BeautifulSoup(
        page_json["result"],
        "html.parser",
    ).find_all("a")

    if not elements:
        Console.log_dim("No info found")
        return {}

    return {
        float(element.attrs["data-num"]): EpisodeInfo(
            number=float(element.attrs["data-num"]),
            slug=element.attrs["data-slug"],
            has_sub=element.attrs["data-sub"] == "1",
            has_dub=element.attrs["data-dub"] == "1",
            ids=str(element.attrs["data-ids"]),
        )
        for element in elements
        if "data-num" in element.attrs and str(element.attrs["data-num"]).isdigit()
    }


@dataclass
class DownloadServerInfo:
    type: str
    episode_id: str
    link_id: str
    server_id: str
    name: str


@dataclass
class EpisodeSkipData:
    intro: Tuple[int, int] | None
    outro: Tuple[int, int] | None

    @classmethod
    def from_str(cls, skip_data: str | None) -> Self | None:
        if not skip_data:
            return None

        try:
            parsed = json.loads(skip_data)
        except json.decoder.JSONDecodeError:
            return None

        ret = cls(intro=None, outro=None)

        try:
            ret.intro = (parsed["intro"][0], parsed["intro"][1])
        except KeyError:
            pass

        try:
            ret.outro = (parsed["outro"][0], parsed["outro"][1])
        except KeyError:
            pass

        return ret


def get_download_url_from_server_info(
    server: DownloadServerInfo,
) -> DownloadSite | None:
    response = get_site_api_with_vrf("/ajax/server/", server.link_id)

    if not response:
        return None

    page_json = response.json()

    try:
        encoded_url = page_json["result"]["url"]
        encoded_skip_data = (
            page_json["result"]["skip_data"]
            if "skip_data" in page_json["result"]
            else None
        )
    except Exception:
        return None

    decoded_url = get_decoded_vrf_data(encoded_url)
    decoded_skip_data = (
        get_decoded_vrf_data(encoded_skip_data) if encoded_skip_data else None
    )

    if not decoded_url:
        return None

    return DownloadSite(
        name=server.name,
        type=server.type,
        url=decoded_url,
        other={
            "skip_data": EpisodeSkipData.from_str(decoded_skip_data),
        },
    )


def get_anime_episode_download_server_list(
    episode: EpisodeInfo,
    series_types: set[str] | list[str],
) -> list[DownloadSite] | None:
    Console.log_dim("Fetching download servers...", return_line=True)

    response = get_site_api_with_vrf(
        "/ajax/server/list/",
        episode.ids,
    )

    if not response:
        Console.log_dim("No download servers found")
        return None

    page_json = response.json()

    type_groups: ResultSet[Tag] = BeautifulSoup(
        page_json["result"],
        "html.parser",
    ).find_all(
        class_="ani-server-type",
    )

    element_groups = (
        (
            str(e.attrs["data-type"]),
            cast(
                ResultSet[Tag],
                e.find_all(
                    class_="server",
                ),
            ),
        )
        for e in type_groups
    )

    servers_infos = [
        DownloadServerInfo(
            episode_id=element.attrs["data-ep-id"],
            name=next(element.stripped_strings),
            server_id=element.attrs["data-sv-id"],
            link_id=element.attrs["data-link-id"],
            type=ep_type,
        )
        for (ep_type, elements) in element_groups
        for element in elements
        if ep_type in series_types
    ]

    Console.log_dim("Fetching list of source pages...", return_line=True)

    with multiprocessing.Pool() as pool:
        results = [
            item
            for item in pool.imap_unordered(
                get_download_url_from_server_info,
                servers_infos,
            )
            if item and item.type in series_types
        ]

    return results


def main():
    def with_additional_args(parser: argparse.ArgumentParser):
        parser.add_argument(
            "-o, --offset",
            help="Set episode name offset (eg. -o 3 will name the first episode 4). Useful for when a season is split into multiple parts.",
            type=int,
            nargs="?",
            required=False,
            default=0,
            dest="offset",
        )

    def get_download_sites(ctx: DownloadSitesCtx):
        anime_id = get_anime_id(ctx.series_name)
        if anime_id is None:
            Console.log_error(f"Can't fetch page for `{ctx.series_name}'")
            exit(1)

        episodes = get_anime_episode_list(anime_id)
        if ctx.episode_number not in episodes:
            raise EpisodeNumberNotFoundException

        return get_anime_episode_download_server_list(
            episodes[ctx.episode_number],
            ctx.series_types,
        )

    downloader_main(
        site_name="9anime",
        with_additional_args=with_additional_args,
        allowed_series_types=ALLOWED_SERIES_TYPES,
        episode_page_url_fn=get_9anime_page_url,
        download_sites_fn=get_download_sites,
        unwanted_cdn_hostnames=UNWANTED_CDN_HOSTNAMES,
    )


if __name__ == "__main__":
    main()