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()