gallery-dl/gallery_dl/extractor/shimmie2.py

# -*- coding: utf-8 -*-

# Copyright 2023 Mike Fährmann
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2 as
# published by the Free Software Foundation.

"""Extractors for Shimmie2 instances"""

from .common import BaseExtractor, Message
from .. import text


class Shimmie2Extractor(BaseExtractor):
    """Base class for shimmie2 extractors"""
    basecategory = "shimmie2"
    filename_fmt = "{category}_{id}{md5:?_//}.{extension}"
    archive_fmt = "{id}"

    def _init(self):
        try:
            instance = INSTANCES[self.category]
        except KeyError:
            return

        cookies = instance.get("cookies")
        if cookies:
            domain = self.root.rpartition("/")[2]
            self.cookies_update_dict(cookies, domain=domain)

        file_url = instance.get("file_url")
        if file_url:
            self.file_url_fmt = file_url

        if self.category == "giantessbooru":
            self.posts = self._posts_giantessbooru

    def items(self):
        data = self.metadata()

        for post in self.posts():

            for key in ("id", "width", "height"):
                post[key] = text.parse_int(post[key])
            post["tags"] = text.unquote(post["tags"])
            post.update(data)

            url = post["file_url"]
            if "/index.php?" in url:
                post["filename"], _, post["extension"] = \
                    url.rpartition("/")[2].rpartition(".")
            else:
                text.nameext_from_url(url, post)

            yield Message.Directory, post
            yield Message.Url, url, post

    def metadata(self):
        """Return general metadata"""
        return ()

    def posts(self):
        """Return an iterable containing data of all relevant posts"""
        return ()


INSTANCES = {
    "loudbooru": {
        "root": "https://loudbooru.com",
        "pattern": r"loudbooru\.com",
        "cookies": {"ui-tnc-agreed": "true"},
    },
    "giantessbooru": {
        "root": "https://giantessbooru.com",
        "pattern": r"giantessbooru\.com",
        "cookies": {"agreed": "true"},
    },
    "tentaclerape": {
        "root": "https://tentaclerape.net",
        "pattern": r"tentaclerape\.net",
    },
    "cavemanon": {
        "root": "https://booru.cavemanon.xyz",
        "pattern": r"booru\.cavemanon\.xyz",
        "file_url": "{0}/index.php?q=image/{2}.{4}",
    },
}

BASE_PATTERN = Shimmie2Extractor.update(INSTANCES) + r"/(?:index\.php\?q=/?)?"


class Shimmie2TagExtractor(Shimmie2Extractor):
    """Extractor for shimmie2 posts by tag search"""
    subcategory = "tag"
    directory_fmt = ("{category}", "{search_tags}")
    file_url_fmt = "{}/_images/{}/{}%20-%20{}.{}"
    pattern = BASE_PATTERN + r"post/list/([^/?#]+)(?:/(\d+))?()"
    example = "https://loudbooru.com/post/list/TAG/1"

    def __init__(self, match):
        Shimmie2Extractor.__init__(self, match)
        lastindex = match.lastindex
        self.tags = text.unquote(match.group(lastindex-2))
        self.page = match.group(lastindex-1)

    def metadata(self):
        return {"search_tags": self.tags}

    def posts(self):
        pnum = text.parse_int(self.page, 1)
        file_url_fmt = self.file_url_fmt.format

        init = True
        mime = ""

        while True:
            url = "{}/post/list/{}/{}".format(self.root, self.tags, pnum)
            page = self.request(url).text
            extr = text.extract_from(page)

            if init:
                init = False
                has_mime = ("data-mime='" in page)
                has_pid = ("data-post-id='" in page)

            while True:
                if has_mime:
                    mime = extr("data-mime='", "'")
                if has_pid:
                    pid = extr("data-post-id='", "'")
                else:
                    pid = extr("href='/post/view/", "?")

                if not pid:
                    break

                tags, dimensions, size = extr("title='", "'").split(" // ")
                width, _, height = dimensions.partition("x")
                md5 = extr("/_thumbs/", "/")

                yield {
                    "file_url": file_url_fmt(
                        self.root, md5, pid, text.quote(tags),
                        mime.rpartition("/")[2] if mime else "jpg"),
                    "id": pid,
                    "md5": md5,
                    "tags": tags,
                    "width": width,
                    "height": height,
                    "size": text.parse_bytes(size[:-1]),
                }

            pnum += 1
            if not extr(">Next<", ">"):
                if not extr("/{}'>{}<".format(pnum, pnum), ">"):
                    return

    def _posts_giantessbooru(self):
        pnum = text.parse_int(self.page, 1)
        file_url_fmt = (self.root + "/index.php?q=/image/{}.jpg").format

        while True:
            url = "{}/index.php?q=/post/list/{}/{}".format(
                self.root, self.tags, pnum)
            extr = text.extract_from(self.request(url).text)

            while True:
                pid = extr('href="./index.php?q=/post/view/', '&')
                if not pid:
                    break

                tags, dimensions, size = extr('title="', '"').split(" // ")
                width, _, height = dimensions.partition("x")

                yield {
                    "file_url": file_url_fmt(pid),
                    "id": pid,
                    "md5": "",
                    "tags": tags,
                    "width": width,
                    "height": height,
                    "size": text.parse_bytes(size[:-1]),
                }

            pnum += 1
            if not extr('/{}">{}<'.format(pnum, pnum), ">"):
                return


class Shimmie2PostExtractor(Shimmie2Extractor):
    """Extractor for single shimmie2 posts"""
    subcategory = "post"
    pattern = BASE_PATTERN + r"post/view/(\d+)"
    example = "https://loudbooru.com/post/view/12345"

    def __init__(self, match):
        Shimmie2Extractor.__init__(self, match)
        self.post_id = match.group(match.lastindex)

    def posts(self):
        url = "{}/post/view/{}".format(self.root, self.post_id)
        extr = text.extract_from(self.request(url).text)

        post = {
            "id"      : self.post_id,
            "tags"    : extr(": ", "<").partition(" - ")[0].rstrip(")"),
            "md5"     : extr("/_thumbs/", "/"),
            "file_url": self.root + (
                extr("id='main_image' src='", "'") or
                extr("<source src='", "'")).lstrip("."),
            "width"   : extr("data-width=", " ").strip("\"'"),
            "height"  : extr("data-height=", ">").partition(
                " ")[0].strip("\"'"),
            "size"    : 0,
        }

        if not post["md5"]:
            post["md5"] = text.extr(post["file_url"], "/_images/", "/")

        return (post,)

    def _posts_giantessbooru(self):
        url = "{}/index.php?q=/post/view/{}".format(
            self.root, self.post_id)
        extr = text.extract_from(self.request(url).text)

        return ({
            "id"      : self.post_id,
            "tags"    : extr(": ", "<").partition(" - ")[0].rstrip(")"),
            "md5"     : "",
            "file_url": self.root + extr('id="main_image" src=".', '"'),
            "width"   : extr("orig_width =", ";"),
            "height"  : 0,
            "size"    : 0,
        },)