path: root/src/portal/py/modules/pixiv_web.py

import log
import json
import urllib.parse
from typing import Optional, Any
from json.decoder import JSONDecodeError
from base import USER_AGENT, HttpResponse, Search, Module, Method, ParsedJson
from post import Url, PostType, Date, DateType, DateMeta, Post, User, Media, Image, Animation, Tag, TagType
from query_parser import QueryParser
from modules.common import parse_pixiv_date, get_current_utc_time

BASE_URL = 'https://www.pixiv.net/ajax'
LANG = 'en'
VERSION = '5f53980f6b59c9376220b6d86e8feb5ea9e22e10'

class PixivWebBase(Search):
    def __init__(self, userdata: Any) -> None:
        super().__init__()
        self.module: PixivWebModule = userdata

    def check_api_response(self, response: Optional[HttpResponse]) -> Optional[ParsedJson]:
        if not response:
            return None
        try:
            obj = response.json()
        except JSONDecodeError as e:
            log.error(f'JSONDecodeError: {e}.')
            return None
        if obj['error']:
            log.error(obj['message'])
            return None
        return obj['body']

    def parse_pages(self, data: ParsedJson) -> dict[str, Media]:
        media = {}
        for index, m in enumerate(data):
            media[str(index)] = Image(url=Url(m['urls']['original']), thumbnail_url=Url(m['urls']['small']))
        return media

    def parse_ugoira(self, data: ParsedJson, thumbnail_url: str) -> Animation:
        animation = Animation(url=Url(data['originalSrc']), thumbnail_url=Url(thumbnail_url))
        for frame in data['frames']:
            animation.frames.append((frame['file'], frame['delay']))
        return animation

    def search_profile_picture_url(self, user_illusts: ParsedJson) -> Optional[Url]:
        for illust in user_illusts.values():
            if illust is None:
                continue
            if 'profileImageUrl' in illust:
                return Url(illust['profileImageUrl'])
        return None

    def create_tag(self, data: ParsedJson) -> Tag:
        kwargs = {}
        kwargs['type'] = TagType.GENERAL
        kwargs['name'] = data['tag']
        kwargs['alts'] = {}
        if 'romaji' in data:
            kwargs['alts']['romaji'] = data['romaji']
        if 'translation' in data:
            if 'en' in data['translation']:
                kwargs['alts']['en'] = data['translation']['en']
        return Tag(**kwargs)

    def create_post(self, data: ParsedJson, pages: Optional[ParsedJson] = None, ugoira: Optional[ParsedJson] = None, retrieved_date: Optional[float] = None) -> Optional[Post]:
        kwargs = {}
        kwargs['type'] = PostType.POST
        illust_id = data['illustId']
        unique_id = f'pixiv:i:{illust_id}'
        kwargs['unique_id'] = unique_id
        kwargs['raw_responses'] = {}
        kwargs['raw_responses']['api'] = json.dumps(data, separators=(',', ':'), indent=None)
        kwargs['url'] = f'https://www.pixiv.net/{LANG}/artworks/{data['illustId']}'
        create_date = parse_pixiv_date(data['createDate']).timestamp()
        if not retrieved_date:
            retrieved_date = get_current_utc_time().timestamp()
        kwargs['dates'] = [
            Date(DateType.CREATED, (create_date, create_date), DateMeta.NONE),
            Date(DateType.RETRIEVED, (retrieved_date, retrieved_date), DateMeta.NONE)
        ]
        upload_date = parse_pixiv_date(data['uploadDate']).timestamp()
        if upload_date != create_date:
            kwargs['dates'].append(Date(DateType.EDITED, (upload_date, upload_date), DateMeta.NONE))
        user = User(unique_id=f'pixiv:u:{data['userId']}', username=data['userAccount'], display_name=data['userName'])
        profile_picture_url = self.search_profile_picture_url(data['userIllusts'])
        if profile_picture_url:
            user.profile_picture_url = profile_picture_url
        kwargs['author'] = user
        kwargs['title'] = data['illustTitle']
        kwargs['text'] = data['illustComment']
        kwargs['likes'] = data['likeCount']
        kwargs['bookmarks'] = data['bookmarkCount']
        kwargs['comments'] = data['commentCount']
        kwargs['views'] = data['viewCount']
        kwargs['tags'] = [self.create_tag(t) for t in data['tags']['tags']]
        if data['illustType'] == 2: # Ugoira
            if not ugoira:
                url = f'{BASE_URL}/illust/{illust_id}/ugoira_meta?lang={LANG}&version={VERSION}'
                ugoira = self.check_api_response(self.module.do_request(Method.GET, url))
        elif data['pageCount'] > 1:
            if not pages:
                url = f'{BASE_URL}/illust/{illust_id}/pages?lang={LANG}&version={VERSION}'
                pages = self.check_api_response(self.module.do_request(Method.GET, url))
        # Retrieving ugoria or pages could return 404.
        if ugoira:
            kwargs['raw_responses']['ugoira'] = json.dumps(ugoira)
            kwargs['media'] = { '0': self.parse_ugoira(ugoira, data['urls']['original']) }
        elif pages:
            kwargs['raw_responses']['pages'] = json.dumps(pages)
            kwargs['media'] = self.parse_pages(pages)
        else:
            kwargs['media'] = self.parse_pages([{ 'urls': data['urls'] }])
        post = Post(**kwargs)
        self.module.add_to_map(unique_id, post)
        return post

    def create_tombstone(self, data: ParsedJson) -> Post:
        kwargs = {}
        kwargs['type'] = PostType.TOMBSTONE
        illust_id = str(data['id'])
        unique_id = f'pixiv:i:{illust_id}'
        kwargs['unique_id'] = unique_id
        post = Post(**kwargs)
        self.module.add_to_map(unique_id, post)
        return post

    def request_illust(self, illust_id: int) -> Optional[Post]:
        url = f'{BASE_URL}/illust/{illust_id}?lang={LANG}&version={VERSION}'
        obj = self.check_api_response(self.module.do_request(Method.GET, url))
        if not obj:
            return self.create_tombstone({ 'id': illust_id })
        return self.create_post(obj)

    def create_user(self, data: ParsedJson) -> User:
        unique_id=f'pixiv:u:{data['userId']}'
        user = User(unique_id=unique_id, display_name=data['userName'])
        user.profile_picture_url = Url(data['profileImageUrl'])
        self.module.add_to_map(unique_id, user)
        return user

    def remake_post(self, raw_responses: ParsedJson, original_date: float) -> Optional[Post]:
        api = json.loads(raw_responses['api'])
        pages = raw_responses.get('pages', None)
        if pages:
            pages = json.loads(pages)
        ugoira = raw_responses.get('ugoira', None)
        if ugoira:
            ugoira = json.loads(ugoira)
        return self.create_post(api, pages, ugoira, original_date)

class PixivWebSearch(PixivWebBase):
    def __init__(self, userdata: Any, arg: str) -> None:
        super().__init__(userdata)
        self.query: str = arg
        self.last_page: Optional[int] = None

    def request_page(self, num: int) -> bool:
        if self.last_page and num >= self.last_page:
            return False
        start = '2019-01-01'
        end = '2019-12-31'
        order = 'popular_d'
        url = f'{BASE_URL}/search/artworks/{urllib.parse.quote(self.query)}?word={self.query}&order={order}&mode=all&scd={start}&ecd={end}&p={num + 1}&csw=0&s_mode=s_tag&type=all&lang={LANG}&version={VERSION}'
        obj = self.check_api_response(self.module.do_request(Method.GET, url))
        if not obj:
            self.errored = True
            return False
        self.last_page = obj['illustManga']['lastPage']
        self.pages[num] = []
        for partial_post in obj['illustManga']['data']:
            post = self.request_illust(int(partial_post['id']))
            if not post:
                self.errored = True
                return False
            self.pages[num].append(post.unique_id)
        return True

class PixivWebUser(PixivWebBase):
    def __init__(self, userdata: Any, arg: str) -> None:
        super().__init__(userdata)
        try:
            self.id: int = int(arg)
        except ValueError:
            self.errored = True

    def request_page(self, num: int) -> bool:
        url = f'{BASE_URL}/user/{self.id}/profile/all?lang={LANG}&version={VERSION}'
        obj = self.check_api_response(self.module.do_request(Method.GET, url))
        if not obj:
            self.errored = True
            return False
        self.pages[num] = []
        keys = []
        if obj['illusts']:
            keys.extend([int(x) for x in obj['illusts'].keys()])
        if obj['manga']:
            keys.extend([int(x) for x in obj['manga'].keys()])
        for illust in keys:
            post = self.request_illust(illust)
            if not post:
                self.errored = True
                return False
            self.pages[num].append(post.unique_id)
        self.completed = True
        return True

class PixivWebBookmarks(PixivWebBase):
    LIMIT = 48

    def __init__(self, userdata: Any, arg: str) -> None:
        super().__init__(userdata)
        try:
            self.id: int = int(arg)
        except ValueError:
            self.errored = True
        self.total: int = 0
        self.count: int = 0

    # rest=show - public
    # rest=hide - private
    def request_page(self, num: int) -> bool:
        url = f'{BASE_URL}/user/{self.id}/illusts/bookmarks?tag=&offset={self.LIMIT * num}&limit={self.LIMIT}&rest=show&lang={LANG}&version={VERSION}'
        obj = self.check_api_response(self.module.do_request(Method.GET, url))
        if not obj:
            self.errored = True
            return False
        self.total = obj['total']
        self.pages[num] = []
        for illust in obj['works']:
            if illust['userId'] == 0:
                # Is it worth assuming this is a tombstone? Should run some tests
                # to see if an illust object is ever returned anyway.
                post = self.create_tombstone(illust)
            else:
                post = self.request_illust(int(illust['id']))
                if not post:
                    self.errored = True
                    return False
            self.pages[num].append(post.unique_id)
            self.count += 1
        if self.count >= self.total:
            self.completed = True
        return True

class PixivWebFollowing(PixivWebBase):
    LIMIT = 24

    def __init__(self, userdata: Any, arg: str) -> None:
        super().__init__(userdata)
        try:
            self.id: int = int(arg)
        except ValueError:
            self.errored = True

    # rest=show - public
    # rest=hide - private
    def request_page(self, num: int) -> bool:
        url = f'{BASE_URL}/user/{self.id}/following?tag=&offset={self.LIMIT * num}&limit={self.LIMIT}&rest=show&acceptingRequests=0&lang={LANG}&version={VERSION}'
        obj = self.check_api_response(self.module.do_request(Method.GET, url))
        if not obj:
            self.errored = True
            return False
        self.pages[num] = []
        for user in obj['users']:
            self.pages[num].append(self.create_user(user).unique_id)
        return True

class PixivWebIllust(PixivWebBase):
    def __init__(self, userdata: Any, arg: str) -> None:
        super().__init__(userdata)
        try:
            self.id: int = int(arg)
        except ValueError:
            self.errored = True

    def request_page(self, num: int) -> bool:
        post = self.request_illust(self.id)
        if not post:
            self.errored = True
            return False
        self.pages[num] = [post.unique_id]
        self.completed = True
        return True

class PixivWebModule(Module):
    def __init__(self, sessid: str, user_id: str) -> None:
        super().__init__()
        self.parser: QueryParser = QueryParser(self, 'search')
        self.parser.add_command('search', PixivWebSearch)
        self.parser.add_command('user', PixivWebUser)
        self.parser.add_command('bookmarks', PixivWebBookmarks)
        self.parser.add_command('following', PixivWebFollowing)
        self.parser.add_command('illust', PixivWebIllust)
        self.headers['User-Agent'] = USER_AGENT
        self.headers['Accept-Encoding'] = 'gzip, deflate, br, zstd'
        self.headers['Accept-Language'] = 'en-US,en;q=0.5'
        self.headers['Referer'] = 'https://www.pixiv.net/'
        self.headers['x-user-id'] = user_id
        self.cookies['PHPSESSID'] = sessid

    def search(self, query: str, *extra_args: Any) -> Optional[Search]:
        return self.parser.parse_query(query)

    def get_download(self, unique_id: str, key: str) -> Optional[dict[str, Any]]:
        if unique_id not in self.unique_id_map:
            return None
        post = self.unique_id_map[unique_id]
        if key not in post.media:
            return None
        return {
            'urls': [post.media[key].url],
            'headers': {
                'User-Agent': USER_AGENT,
                'Referer': 'https://www.pixiv.net/'
            }
        }