import argparse import importlib import os from flask import Flask, Response from fbin.db import db, User, File parser = argparse.ArgumentParser() parser.add_argument('-s', '--source-config-file', default='fbin/fbin.cfg') parser.add_argument('-t', '--target-config-file', required=True) parser.add_argument('--update-db', action='store_true', help='Update DB as well') parser.add_argument('-n', '--dry-run', action='store_true', help='Do not update anything') args = parser.parse_args() class DummyFile: def __init__(self, stream, filename, content_length): self.stream = stream self.filename = filename self.content_length = content_length def save(self, fp): chunk = self.stream.read(10*1024) while chunk: fp.write(chunk) chunk = self.stream.read(10*1024) class IterStream: def __init__(self, it): self._it = it def read(self, n): try: return next(self._it) except StopIteration: return None def main(): from fbin.file_storage.s3 import Storage as S3Storage from fbin.file_storage.filesystem import Storage as FSStorage source_storage = importlib.import_module(source_app.config.get('STORAGE_MODULE', '.file_storage.filesystem'), package='fbin').Storage(source_app) target_storage = importlib.import_module(target_app.config.get('STORAGE_MODULE', '.file_storage.filesystem'), package='fbin').Storage(target_app) copy_list = [] with source_app.app_context(): db.init_app(source_app) print('Finding existing files') for f in db.session.query(File).all(): if f.user: db.session.refresh(f.user) if source_storage.file_exists(f) and not target_storage.file_exists(f): print('COPY: ', end='') copy_list.append(f) else: print('SKIP: ', end='') print(f.hash, f.filename) if not copy_list: print('No valid files found') return print('Copying {} files'.format(len(copy_list))) for f in copy_list: if args.dry_run: print('Would copy', f.hash, f.filename) continue print('Copying', f.hash, f.filename) with source_app.app_context(): db.init_app(source_app) with source_app.test_request_context(): source = source_storage.get_file(f) if isinstance(source, str) and os.path.exists(source): source = open(source, 'rb') elif isinstance(source, Response): source = IterStream(source.get_app_iter({'REQUEST_METHOD': 'GET'})) df = DummyFile(source, f.filename, f.size) with target_app.app_context(): db.init_app(target_app) if args.update_db: if db.session.query(File).filter(File.hash == f.hash).one(): print(' Cannot copy this file; hash already exists in target DB') else: target_storage.store_file(df, f.hash, f.user, f.ip) else: target_storage.upload_file(df, f.hash, f.user) source_app = Flask('source') target_app = Flask('target') with source_app.app_context(): source_app.config.from_pyfile(args.source_config_file) with target_app.app_context(): target_app.config.from_pyfile(args.target_config_file) main()