path: root/src/portal/py/tests/archive_query.py

import os
import sys
import signal
import threading
import json
import gzip
from datetime import datetime, timezone
sys.path.append('../')
from base import Method
from post import PostEncoder, PostType, DateType, User
from modules import ALL_MODULES

class Logger():
    def __init__(self, path):
        self.file = open(path, 'a+')
        self.mutex = threading.Lock()

    def write(self, message):
        with self.mutex:
            print(message)
            self.file.write(message + '\n')
            self.file.flush()

    def close(self):
        self.file.close()

#mode = 'pixiv_web'
mode = 'fanbox'
#mode = 'patreon'
#mode = 'instagram'
#mode = 'twitter'
#cmd = 'profile'
cmd = 'user'
#cmd = 'bookmarks'
#cmd = 'search'
module = ALL_MODULES[mode][0]
if not module.init():
    sys.exit(1)

quit_mutex = threading.Lock()
queue_mutex = threading.Lock()
queue_cond = threading.Condition(queue_mutex)
is_running = True

def media_check_disk(obj, post, output_base):
    for key in post.media.keys():
        download_params = module.get_download(post.unique_id, key)
        if not download_params:
            obj.log.write(f'Ignoring media{key} for post {post.unique_id}.')
            continue
        for url in download_params['urls']:
            media_path = f'{output_base}_media{key}.{url.ext}'
            if not os.path.isfile(media_path):
                return False
    return True

def default_media_download(obj, post, output_base):
    for key in post.media.keys():
        download_params = module.get_download(post.unique_id, key)
        if not download_params:
            obj.log.write(f'Skipping media{key} for post {post.unique_id}.')
            continue
        for url in download_params['urls']:
            media_path = f'{output_base}_media{key}.{url.ext}'
            obj.log.write(f'Attempting to download file {media_path}.')
            response = module.do_request(Method.GET, url.url)
            if not response:
                obj.log.write(f'ERROR: Downloading post media failed ({post.unique_id} #{key}).')
                continue
            if not obj.write_to_file(media_path, 'wb+', response.content, False):
                return False

class QueryDownloadThread(threading.Thread):
    def __init__(self, log, complete, query, output_dir, start_page):
        super().__init__()
        self.log = log
        self.complete = complete
        self.query = query
        self.output_dir = output_dir
        self.media_download = default_media_download
        self.start_page = start_page

    def make_path(self, post):
        prefix: str
        if post.type != PostType.TOMBSTONE:
            prefix = post.author.unique_id.replace(':', '_')
        else:
            prefix = 'tombstones'
        path = f'{self.output_dir}/{prefix}'
        if not os.path.isdir(path):
            os.mkdir(path)
        return path

    def write_to_file(self, path, open_method, content, compress):
        try:
            if compress:
                with gzip.open(path, open_method) as f:
                    f.write(content.encode('utf-8'))
            else:
                with open(path, open_method) as f:
                    f.write(content)
        except Exception as e:
            self.log.write(f'ERROR: Writing file failed: {path} ({repr(e)}).')
            return False
        return True

    def write_post_to_disk(self, post):
        if post.post.unique_id:
            original = module.get_item(post.post.unique_id)
            if original:
                self.write_post_to_disk(original)
        if post.quoted.unique_id:
            quoted = module.get_item(post.quoted.unique_id)
            if quoted:
                self.write_post_to_disk(quoted)
        message = f'Downloading post {post.unique_id}'
        if 'detached_api' in post.raw_responses:
            raw_response_path = f'{self.output_dir}/raw_responses'
            if not os.path.isdir(raw_response_path):
                os.mkdir(raw_response_path)
            raw_response_path += f'/{post.raw_responses['hash']}.json.gz'
            if not self.write_to_file(raw_response_path, 'wb+', json.dumps(json.loads(post.raw_responses['detached_api'])), True):
                return False
            del post.raw_responses['detached_api']
        if post.type != PostType.TOMBSTONE:
            message += f':{datetime.fromtimestamp(post.dates[0].range[0], tz=timezone.utc).isoformat()} from {post.author.username}({post.author.display_name}):{post.author.unique_id}.'
        else:
            message += '.'
        output_path = self.make_path(post)
        output_base = f'{output_path}/{post.unique_id.replace(':', '_')}'
        output = output_base + '.json.gz'
        '''
        if not os.path.isfile(output):
            print('done')
            sys.exit(1)
        '''
        '''
        has_mp4 = False
        for m in post.media.values():
            if m.url.ext == 'mp4':
                self.log.write('Rewriting post with mp4.')
                has_mp4 = True
                break
        if post.quoted.unique_id:
            self.log.write('Rewriting post with quote.')
            has_mp4 = True
        '''
        if os.path.isfile(output) and media_check_disk(self, post, output_base):
            skip_message = f'Skipping already downloaded post {post.unique_id}'
            if post.type != PostType.TOMBSTONE:
                skip_message += f' from {post.author.unique_id}.'
            else:
                skip_message += '.'
            self.log.write(skip_message)
            return True
        self.log.write(message)
        self.media_download(self, post, output_base)
        if not self.write_to_file(output, 'wb+', json.dumps(post, cls=PostEncoder), True):
            return False
        return True

    def supply_post(self, post):
        if not self.write_post_to_disk(post):
            return False
        return True

    def run(self):
        global quit_mutex
        global queue_cond
        global is_running
        search = module.search(self.query)
        if not search:
            return
        completed = True
        arg = self.query.split(':')[-1]
        num = self.start_page
        count = 0
        while completed:
            self.log.write(f'Starting page {num} of {arg}.')
            page = search.get_page(num)
            for hash_str, response in module.raw_responses.items():
                raw_response_path = f'{self.output_dir}/raw_responses/{hash_str}.json.gz'
                if not os.path.isfile(raw_response_path):
                    if not self.write_to_file(raw_response_path, 'wb+', response, True):
                        completed = False
                        break
            if not page:
                break
            num += 1
            for unique_id in page:
                #if unique_id == '':
                #    self.log.write('---------Hit target---------')
                post = module.get_item(unique_id)
                if not post:
                    self.log.write(f'Missing post with id {unique_id}.')
                    continue
                if not self.write_post_to_disk(post):
                    continue
                count += 1
            with quit_mutex:
                if not is_running:
                    self.complete.write(f'{mode}:{cmd}:{arg}:i:{str(count)},{str(num)}')
                    completed = False
                    break
        if completed:
            self.complete.write(f'{mode}:{cmd}:{arg}:c:{str(count)},{str(num)}')
        with queue_cond:
            queue_cond.notify()

#RUNTIME_PATH = './run'
RUNTIME_PATH = '/mnt/store/files/tmp/run'
LOG_FILE = f'{RUNTIME_PATH}/archive5.log'
ARG_FILE = f'{RUNTIME_PATH}/completed_args5.log'

def read_completed_args(path):
    args = []
    if os.path.isfile(path):
        with open(path, 'r') as f:
            for line in f.read().splitlines():
                if line[0] == '-':
                    continue
                s = line.split(':')
                if s[3] == 'c':
                    args.append(f'{s[0]}:{s[1]}:{s[2]}')
    return args

def download_args():
    global quit_mutex
    global queue_cond
    global compat_thread

    completed_args = read_completed_args(ARG_FILE)

    log = Logger(LOG_FILE)
    complete = Logger(ARG_FILE)

    header = f'--------{datetime.now().isoformat()}---------'
    log.write(header)
    complete.write(header)

    output_dir = sys.argv[1]
    if not os.path.isdir(output_dir):
        os.mkdir(output_dir)
    if not os.path.isdir(f'{output_dir}/raw_responses'):
        os.mkdir(f'{output_dir}/raw_responses')

    args = []
    start_page = 0
    threads = []
    start = True

    for arg in args:
        #if mode == 'fanbox':
        #    user = module.get_item(arg)
        #    if not user or not isinstance(user, User):
        #        log.write(f'Failed to retrive information for user {arg}.')
        #        continue
        #    arg = user.username
        with quit_mutex:
            if not is_running:
                break
        if mode == 'pixiv_web' and not start:
            if arg == '':
                start = True
                continue
        if not start:
            continue
        if mode != 'instagram':
            arg = f'{cmd}:{arg}'
        if f'{mode}:{arg}' in completed_args:
            log.write(f'Skipping already downloaded arg {arg}.')
            continue
        if len(threads) >= 1:
            with queue_cond:
                queue_cond.wait()
            threads = [t for t in threads if t.is_alive()]
        with quit_mutex:
            if not is_running:
                break
        log.write(f'***********************\nStarting download of {arg}\n***********************\n')
        arg_thread = QueryDownloadThread(log, complete, arg, output_dir, start_page)
        threads.append(arg_thread)
        arg_thread.start()

    for t in threads:
        t.join()

    log.close()
    complete.close()

def signal_handler(sig, frame):
    global quit_mutex
    global is_running
    print('Attempting to quit...')
    with quit_mutex:
        if not is_running:
            sys.exit(1)
        else:
            is_running = False

signal.signal(signal.SIGINT, signal_handler)
download_args()