From 3d37c494aaadea59169cd563011216c444d83569 Mon Sep 17 00:00:00 2001 From: Patrick Robertson Date: Wed, 29 Jan 2025 18:42:12 +0100 Subject: [PATCH 01/62] Tidy ups + unit tests: 1. Allow loading modules from --module_paths=/extra/path/here 2. Improved unit tests for module loading 3. Further small tidy ups/clean ups --- src/auto_archiver/core/config.py | 4 +- src/auto_archiver/core/module.py | 161 ++++++++++++++---- src/auto_archiver/core/orchestrator.py | 29 ++-- .../modules/hash_enricher/hash_enricher.py | 10 -- .../modules/html_formatter/html_formatter.py | 4 +- src/auto_archiver/utils/misc.py | 5 +- tests/conftest.py | 5 +- tests/data/example_module/__init__.py | 1 + tests/data/example_module/__manifest__.py | 10 ++ tests/data/example_module/example_module.py | 4 + tests/enrichers/test_hash_enricher.py | 6 +- tests/extractors/test_extractor_base.py | 3 +- tests/test_modules.py | 55 +++++- 13 files changed, 216 insertions(+), 81 deletions(-) create mode 100644 tests/data/example_module/__init__.py create mode 100644 tests/data/example_module/__manifest__.py create mode 100644 tests/data/example_module/example_module.py diff --git a/src/auto_archiver/core/config.py b/src/auto_archiver/core/config.py index 529e1c2..46dbe28 100644 --- a/src/auto_archiver/core/config.py +++ b/src/auto_archiver/core/config.py @@ -11,7 +11,7 @@ from ruamel.yaml import YAML, CommentedMap, add_representer from loguru import logger from copy import deepcopy -from .module import MODULE_TYPES +from .module import BaseModule from typing import Any, List, Type, Tuple @@ -21,7 +21,7 @@ EMPTY_CONFIG = yaml.load(""" # Auto Archiver Configuration # Steps are the modules that will be run in the order they are defined -steps:""" + "".join([f"\n {module}s: []" for module in MODULE_TYPES]) + \ +steps:""" + "".join([f"\n {module}s: []" for module in BaseModule.MODULE_TYPES]) + \ """ # Global configuration diff --git a/src/auto_archiver/core/module.py b/src/auto_archiver/core/module.py index 18f791b..0888378 100644 --- a/src/auto_archiver/core/module.py +++ b/src/auto_archiver/core/module.py @@ -16,33 +16,53 @@ from importlib.util import find_spec import os from os.path import join, dirname from loguru import logger +import auto_archiver _LAZY_LOADED_MODULES = {} -MODULE_TYPES = [ - 'feeder', - 'extractor', - 'enricher', - 'database', - 'storage', - 'formatter' -] - MANIFEST_FILE = "__manifest__.py" -_DEFAULT_MANIFEST = { - 'name': '', - 'author': 'Bellingcat', - 'type': [], - 'requires_setup': True, - 'description': '', - 'dependencies': {}, - 'entry_point': '', - 'version': '1.0', - 'configs': {} -} class BaseModule(ABC): + """ + Base module class. All modules should inherit from this class. + + The exact methods a class implements will depend on the type of module it is, + however all modules have a .setup(config: dict) method to run any setup code + (e.g. logging in to a site, spinning up a browser etc.) + + See BaseModule.MODULE_TYPES for the types of modules you can create, noting that + a subclass can be of multiple types. For example, a module that extracts data from + a website and stores it in a database would be both an 'extractor' and a 'database' module. + + Each module is a python package, and should have a __manifest__.py file in the + same directory as the module file. The __manifest__.py specifies the module information + like name, author, version, dependencies etc. See BaseModule._DEFAULT_MANIFEST for the + default manifest structure. + + """ + + MODULE_TYPES = [ + 'feeder', + 'extractor', + 'enricher', + 'database', + 'storage', + 'formatter' + ] + + _DEFAULT_MANIFEST = { + 'name': '', # the display name of the module + 'author': 'Bellingcat', # creator of the module, leave this as Bellingcat or set your own name! + 'type': [], # the type of the module, can be one or more of BaseModule.MODULE_TYPES + 'requires_setup': True, # whether or not this module requires additional setup such as setting API Keys or installing additional softare + 'description': '', # a description of the module + 'dependencies': {}, # external dependencies, e.g. python packages or binaries, in dictionary format + 'entry_point': '', # the entry point for the module, in the format 'module_name::ClassName'. This can be left blank to use the default entry point of module_name::ModuleName + 'version': '1.0', # the version of the module + 'configs': {} # any configuration options this module has, these will be exposed to the user in the config file or via the command line +} + config: dict name: str @@ -51,15 +71,51 @@ class BaseModule(ABC): for key, val in config.get(self.name, {}).items(): setattr(self, key, val) -def get_module(module_name: str, additional_paths: List[str] = []) -> LazyBaseModule: + def repr(self): + return f"Module<'{self.display_name}' (config: {self.config[self.name]})>" + + +def setup_paths(paths: list[str]) -> None: + """ + Sets up the paths for the modules to be loaded from + + This is necessary for the modules to be imported correctly + + """ + for path in paths: + # see odoo/module/module.py -> initialize_sys_path + if path not in auto_archiver.modules.__path__: + auto_archiver.modules.__path__.append(path) + + # sort based on the length of the path, so that the longest path is last in the list + auto_archiver.modules.__path__ = sorted(auto_archiver.modules.__path__, key=len, reverse=True) + + +def get_module(module_name: str, config: dict) -> BaseModule: + """ + Gets and sets up a module using the provided config + + This will actually load and instantiate the module, and load all its dependencies (i.e. not lazy) + + """ + return get_module_lazy(module_name).load(config) + +def get_module_lazy(module_name: str, suppress_warnings: bool = False) -> LazyBaseModule: + """ + Lazily loads a module, returning a LazyBaseModule + + This has all the information about the module, but does not load the module itself or its dependencies + + To load an actual module, call .setup() on a laz module + + """ if module_name in _LAZY_LOADED_MODULES: return _LAZY_LOADED_MODULES[module_name] - module = available_modules(additional_paths=additional_paths, limit_to_modules=[module_name])[0] - _LAZY_LOADED_MODULES[module_name] = module + module = available_modules(limit_to_modules=[module_name], suppress_warnings=suppress_warnings)[0] return module -def available_modules(with_manifest: bool=False, limit_to_modules: List[str]= [], additional_paths: List[str] = [], suppress_warnings: bool = False) -> List[LazyBaseModule]: +def available_modules(with_manifest: bool=False, limit_to_modules: List[str]= [], suppress_warnings: bool = False) -> List[LazyBaseModule]: # search through all valid 'modules' paths. Default is 'modules' in the current directory # see odoo/modules/module.py -> get_modules @@ -67,10 +123,9 @@ def available_modules(with_manifest: bool=False, limit_to_modules: List[str]= [] if os.path.isfile(join(module_path, MANIFEST_FILE)): return True - default_path = [join(dirname(dirname((__file__))), "modules")] all_modules = [] - for module_folder in default_path + additional_paths: + for module_folder in auto_archiver.modules.__path__: # walk through each module in module_folder and check if it has a valid manifest try: possible_modules = os.listdir(module_folder) @@ -85,8 +140,12 @@ def available_modules(with_manifest: bool=False, limit_to_modules: List[str]= [] possible_module_path = join(module_folder, possible_module) if not is_really_module(possible_module_path): continue - - all_modules.append(LazyBaseModule(possible_module, possible_module_path)) + if _LAZY_LOADED_MODULES.get(possible_module): + continue + lazy_module = LazyBaseModule(possible_module, possible_module_path) + _LAZY_LOADED_MODULES[possible_module] = lazy_module + + all_modules.append(lazy_module) if not suppress_warnings: for module in limit_to_modules: @@ -97,8 +156,14 @@ def available_modules(with_manifest: bool=False, limit_to_modules: List[str]= [] @dataclass class LazyBaseModule: + + """ + A lazy module class, which only loads the manifest and does not load the module itself. + + This is useful for getting information about a module without actually loading it. + + """ name: str - display_name: str type: list description: str path: str @@ -129,6 +194,10 @@ class LazyBaseModule: @property def requires_setup(self) -> bool: return self.manifest['requires_setup'] + + @property + def display_name(self) -> str: + return self.manifest['name'] @property def manifest(self) -> dict: @@ -136,7 +205,7 @@ class LazyBaseModule: return self._manifest # print(f"Loading manifest for module {module_path}") # load the manifest file - manifest = copy.deepcopy(_DEFAULT_MANIFEST) + manifest = copy.deepcopy(BaseModule._DEFAULT_MANIFEST) with open(join(self.path, MANIFEST_FILE)) as f: try: @@ -145,7 +214,6 @@ class LazyBaseModule: logger.error(f"Error loading manifest from file {self.path}/{MANIFEST_FILE}: {e}") self._manifest = manifest - self.display_name = manifest['name'] self.type = manifest['type'] self._entry_point = manifest['entry_point'] self.description = manifest['description'] @@ -153,7 +221,7 @@ class LazyBaseModule: return manifest - def load(self) -> BaseModule: + def load(self, config) -> BaseModule: if self._instance: return self._instance @@ -162,10 +230,27 @@ class LazyBaseModule: def check_deps(deps, check): for dep in deps: if not check(dep): - logger.error(f"Module '{self.name}' requires external dependency '{dep}' which is not available. Have you installed the required dependencies for the '{self.name}' module? See the README for more information.") + logger.error(f"Module '{self.name}' requires external dependency '{dep}' which is not available/setup. Have you installed the required dependencies for the '{self.name}' module? See the README for more information.") exit(1) - check_deps(self.dependencies.get('python', []), lambda dep: find_spec(dep)) + def check_python_dep(dep): + # first check if it's a module: + try: + m = get_module_lazy(dep, suppress_warnings=True) + try: + # we must now load this module and set it up with the config + m.load(config) + return True + except: + logger.error(f"Unable to setup module '{dep}' for use in module '{self.name}'") + return False + except IndexError: + # not a module, continue + pass + + return find_spec(dep) + + check_deps(self.dependencies.get('python', []), check_python_dep) check_deps(self.dependencies.get('bin', []), lambda dep: shutil.which(dep)) @@ -184,9 +269,8 @@ class LazyBaseModule: sub_qualname = f'{qualname}.{file_name}' __import__(f'{qualname}.{file_name}', fromlist=[self.entry_point]) - # finally, get the class instance - instance = getattr(sys.modules[sub_qualname], class_name)() + instance: BaseModule = getattr(sys.modules[sub_qualname], class_name)() if not getattr(instance, 'name', None): instance.name = self.name @@ -194,6 +278,11 @@ class LazyBaseModule: instance.display_name = self.display_name self._instance = instance + + # merge the default config with the user config + default_config = dict((k, v['default']) for k, v in self.configs.items() if v.get('default')) + config[self.name] = default_config | config.get(self.name, {}) + instance.setup(config) return instance def __repr__(self): diff --git a/src/auto_archiver/core/orchestrator.py b/src/auto_archiver/core/orchestrator.py index 16cf9c4..dc15809 100644 --- a/src/auto_archiver/core/orchestrator.py +++ b/src/auto_archiver/core/orchestrator.py @@ -19,7 +19,7 @@ from .context import ArchivingContext from .metadata import Metadata from ..version import __version__ from .config import read_yaml, store_yaml, to_dot_notation, merge_dicts, EMPTY_CONFIG, DefaultValidatingParser -from .module import available_modules, LazyBaseModule, MODULE_TYPES, get_module +from .module import available_modules, LazyBaseModule, get_module, setup_paths from . import validators from .module import BaseModule @@ -57,6 +57,7 @@ class ArchivingOrchestrator: # override the default 'help' so we can inject all the configs and show those parser.add_argument('-h', '--help', action='store_true', dest='help', help='show this help message and exit') parser.add_argument('-s', '--store', dest='store', default=False, help='Store the created config in the config file', action=argparse.BooleanOptionalAction) + parser.add_argument('--module_paths', dest='module_paths', nargs='+', default=[], help='additional paths to search for modules', action=UniqueAppendAction) self.basic_parser = parser @@ -72,19 +73,21 @@ class ArchivingOrchestrator: # if full, we'll load all modules # TODO: BUG** - basic_config won't have steps in it, since these args aren't added to 'basic_parser' # but should we add them? Or should we just add them to the 'complete' parser? + if yaml_config != EMPTY_CONFIG: # only load the modules enabled in config # TODO: if some steps are empty (e.g. 'feeders' is empty), should we default to the 'simple' ones? Or only if they are ALL empty? enabled_modules = [] - for module_type in MODULE_TYPES: + for module_type in BaseModule.MODULE_TYPES: enabled_modules.extend(yaml_config['steps'].get(f"{module_type}s", [])) # add in any extra modules that have been passed on the command line for 'feeders', 'enrichers', 'archivers', 'databases', 'storages', 'formatter' - for module_type in MODULE_TYPES: + for module_type in BaseModule.MODULE_TYPES: if modules := getattr(basic_config, f"{module_type}s", []): enabled_modules.extend(modules) - self.add_module_args(available_modules(with_manifest=True, limit_to_modules=set(enabled_modules), suppress_warnings=True), parser) + avail_modules = available_modules(with_manifest=True, limit_to_modules=list(dict.fromkeys(enabled_modules)), suppress_warnings=True) + self.add_module_args(avail_modules, parser) elif basic_config.mode == 'simple': simple_modules = [module for module in available_modules(with_manifest=True) if not module.requires_setup] self.add_module_args(simple_modules, parser) @@ -135,10 +138,7 @@ class ArchivingOrchestrator: parser.add_argument('--logging.file', action='store', dest='logging.file', help='the logging file to write to', default=None) parser.add_argument('--logging.rotation', action='store', dest='logging.rotation', help='the logging rotation to use', default=None) - # additional modules - parser.add_argument('--additional-modules', dest='additional_modules', nargs='+', help='additional paths to search for modules', action=UniqueAppendAction) - - def add_module_args(self, modules: list[LazyBaseModule] = None, parser: argparse.ArgumentParser = None): + def add_module_args(self, modules: list[LazyBaseModule] = None, parser: argparse.ArgumentParser = None) -> None: if not modules: modules = available_modules(with_manifest=True) @@ -173,7 +173,7 @@ class ArchivingOrchestrator: arg = group.add_argument(f"--{module.name}.{name}", **kwargs) arg.should_store = should_store - def show_help(self): + def show_help(self, basic_config: dict): # for the help message, we want to load *all* possible modules and show the help # add configs as arg parser arguments @@ -198,7 +198,7 @@ class ArchivingOrchestrator: """ invalid_modules = [] - for module_type in MODULE_TYPES: + for module_type in BaseModule.MODULE_TYPES: step_items = [] modules_to_load = self.config['steps'][f"{module_type}s"] @@ -216,9 +216,8 @@ class ArchivingOrchestrator: for module in modules_to_load: if module in invalid_modules: continue - loaded_module: BaseModule = get_module(module).load() try: - loaded_module.setup(self.config) + loaded_module: BaseModule = get_module(module, self.config) except (KeyboardInterrupt, Exception) as e: logger.error(f"Error during setup of archivers: {e}\n{traceback.format_exc()}") if module_type == 'extractor': @@ -249,9 +248,11 @@ class ArchivingOrchestrator: # load the config file to get the list of enabled items basic_config, unused_args = self.basic_parser.parse_known_args() + setup_paths(basic_config.module_paths) + # if help flag was called, then show the help if basic_config.help: - self.show_help() + self.show_help(basic_config) # load the config file yaml_config = {} @@ -268,7 +269,7 @@ class ArchivingOrchestrator: self.install_modules() # log out the modules that were loaded - for module_type in MODULE_TYPES: + for module_type in BaseModule.MODULE_TYPES: logger.info(f"{module_type.upper()}S: " + ", ".join(m.display_name for m in self.config['steps'][f"{module_type}s"])) for item in self.feed(): diff --git a/src/auto_archiver/modules/hash_enricher/hash_enricher.py b/src/auto_archiver/modules/hash_enricher/hash_enricher.py index 827b65f..94b5dce 100644 --- a/src/auto_archiver/modules/hash_enricher/hash_enricher.py +++ b/src/auto_archiver/modules/hash_enricher/hash_enricher.py @@ -19,16 +19,6 @@ class HashEnricher(Enricher): Calculates hashes for Media instances """ - def __init__(self, config: dict = None): - """ - Initialize the HashEnricher with a configuration dictionary. - """ - super().__init__() - # TODO set these from the manifest? - # Set default values - self.algorithm = config.get("algorithm", "SHA-256") if config else "SHA-256" - self.chunksize = config.get("chunksize", int(1.6e7)) if config else int(1.6e7) - def enrich(self, to_enrich: Metadata) -> None: url = to_enrich.get_url() diff --git a/src/auto_archiver/modules/html_formatter/html_formatter.py b/src/auto_archiver/modules/html_formatter/html_formatter.py index e6e5e58..570fc6f 100644 --- a/src/auto_archiver/modules/html_formatter/html_formatter.py +++ b/src/auto_archiver/modules/html_formatter/html_formatter.py @@ -12,7 +12,7 @@ from auto_archiver.core import Metadata, Media, ArchivingContext from auto_archiver.core import Formatter from auto_archiver.modules.hash_enricher import HashEnricher from auto_archiver.utils.misc import random_str - +from auto_archiver.core.module import get_module @dataclass class HtmlFormatter(Formatter): @@ -53,7 +53,7 @@ class HtmlFormatter(Formatter): outf.write(content) final_media = Media(filename=html_path, _mimetype="text/html") - he = HashEnricher({"hash_enricher": {"algorithm": ArchivingContext.get("hash_enricher.algorithm"), "chunksize": 1.6e7}}) + he = get_module('hash_enricher', self.config) if len(hd := he.calculate_hash(final_media.filename)): final_media.set("hash", f"{he.algorithm}:{hd}") diff --git a/src/auto_archiver/utils/misc.py b/src/auto_archiver/utils/misc.py index e985e3e..300a710 100644 --- a/src/auto_archiver/utils/misc.py +++ b/src/auto_archiver/utils/misc.py @@ -1,7 +1,10 @@ -import os, json, requests + +import os +import json import uuid from datetime import datetime +import requests from loguru import logger diff --git a/tests/conftest.py b/tests/conftest.py index c2c74f2..af0fd6d 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -23,10 +23,7 @@ def setup_module(request): # if the class does not have a .name, use the name of the parent folder module_name = module_name.__module__.rsplit(".",2)[-2] - m = get_module(module_name).load() - m.name = module_name - m.setup({module_name : config}) - + m = get_module(module_name, {module_name: config}) def cleanup(): _LAZY_LOADED_MODULES.pop(module_name) diff --git a/tests/data/example_module/__init__.py b/tests/data/example_module/__init__.py new file mode 100644 index 0000000..560a9b9 --- /dev/null +++ b/tests/data/example_module/__init__.py @@ -0,0 +1 @@ +from .example_module import ExampleModule \ No newline at end of file diff --git a/tests/data/example_module/__manifest__.py b/tests/data/example_module/__manifest__.py new file mode 100644 index 0000000..ca3a678 --- /dev/null +++ b/tests/data/example_module/__manifest__.py @@ -0,0 +1,10 @@ +{ + "name": "Example Module", + "type": ["extractor"], + "requires_setup": False, + "external_dependencies": {"python": ["loguru"] + }, + "configs": { + "csv_file": {"default": "db.csv", "help": "CSV file name"} + }, +} \ No newline at end of file diff --git a/tests/data/example_module/example_module.py b/tests/data/example_module/example_module.py new file mode 100644 index 0000000..b752743 --- /dev/null +++ b/tests/data/example_module/example_module.py @@ -0,0 +1,4 @@ +from auto_archiver.core.extractor import Extractor + +class ExampleModule(Extractor): + pass \ No newline at end of file diff --git a/tests/enrichers/test_hash_enricher.py b/tests/enrichers/test_hash_enricher.py index 63e4824..4b61fc2 100644 --- a/tests/enrichers/test_hash_enricher.py +++ b/tests/enrichers/test_hash_enricher.py @@ -2,7 +2,7 @@ import pytest from auto_archiver.modules.hash_enricher import HashEnricher from auto_archiver.core import Metadata, Media -from auto_archiver.core.module import get_module +from auto_archiver.core.module import get_module_lazy @pytest.mark.parametrize("algorithm, filename, expected_hash", [ ("SHA-256", "tests/data/testfile_1.txt", "1b4f0e9851971998e732078544c96b36c3d01cedf7caa332359d6f1d83567014"), @@ -12,7 +12,7 @@ from auto_archiver.core.module import get_module ]) def test_calculate_hash(algorithm, filename, expected_hash, setup_module): # test SHA-256 - he = setup_module(HashEnricher, {"algorithm": algorithm, "chunksize": 1}) + he = setup_module(HashEnricher, {"algorithm": algorithm, "chunksize": 100}) assert he.calculate_hash(filename) == expected_hash def test_default_config_values(setup_module): @@ -22,7 +22,7 @@ def test_default_config_values(setup_module): def test_config(): # test default config - c = get_module('hash_enricher').configs + c = get_module_lazy('hash_enricher').configs assert c["algorithm"]["default"] == "SHA-256" assert c["chunksize"]["default"] == 16000000 assert c["algorithm"]["choices"] == ["SHA-256", "SHA3-512"] diff --git a/tests/extractors/test_extractor_base.py b/tests/extractors/test_extractor_base.py index bb78794..f6be70b 100644 --- a/tests/extractors/test_extractor_base.py +++ b/tests/extractors/test_extractor_base.py @@ -2,7 +2,7 @@ import pytest from auto_archiver.core.metadata import Metadata from auto_archiver.core.extractor import Extractor -from auto_archiver.core.module import get_module + class TestExtractorBase(object): extractor_module: str = None @@ -12,6 +12,7 @@ class TestExtractorBase(object): def setup_archiver(self, setup_module): assert self.extractor_module is not None, "self.extractor_module must be set on the subclass" assert self.config is not None, "self.config must be a dict set on the subclass" + self.extractor: Extractor = setup_module(self.extractor_module, self.config) def assertValidResponseMetadata(self, test_response: Metadata, title: str, timestamp: str, status: str = ""): diff --git a/tests/test_modules.py b/tests/test_modules.py index 619906b..109bc52 100644 --- a/tests/test_modules.py +++ b/tests/test_modules.py @@ -1,26 +1,65 @@ +import sys import pytest -from auto_archiver.core.module import get_module, BaseModule, LazyBaseModule +from auto_archiver.core.module import get_module_lazy, BaseModule, LazyBaseModule, _LAZY_LOADED_MODULES +from auto_archiver.core.extractor import Extractor + +@pytest.fixture +def example_module(): + yield get_module_lazy("example_module", ["tests/data/"]) + # cleanup + _LAZY_LOADED_MODULES.pop("example_module") + +def test_get_module_lazy(example_module): + assert example_module.name == "example_module" + assert example_module.display_name == "Example Module" + + assert example_module.manifest is not None + + +def test_load_module_abc_check(example_module): + + # example_module is an extractor but doesn't have the 'download' method, should raise an ABC error + with pytest.raises(TypeError) as load_error: + example_module.load({}) + assert "Can't instantiate abstract class ExampleModule with abstract method download" in str(load_error.value) + + +def test_load_module(example_module, monkeypatch): + # hack - remove the 'download' method from the required methods of Extractor + monkeypatch.setattr(Extractor, "__abstractmethods__", set()) + + # setup the module, and check that config is set to the default values + loaded_module = example_module.load({}) + assert loaded_module is not None + assert isinstance(loaded_module, BaseModule) + assert loaded_module.name == "example_module" + assert loaded_module.display_name == "Example Module" + assert loaded_module.config["example_module"] == {"csv_file" : "db.csv"} + + # check that the vlaue is set on the module itself + assert loaded_module.csv_file == "db.csv" @pytest.mark.parametrize("module_name", ["cli_feeder", "local_storage", "generic_extractor", "html_formatter", "csv_db"]) def test_load_modules(module_name): # test that specific modules can be loaded - module = get_module(module_name) + module = get_module_lazy(module_name) assert module is not None assert isinstance(module, LazyBaseModule) assert module.name == module_name - loaded_module = module.load() + loaded_module = module.load({}) assert isinstance(loaded_module, BaseModule) + assert loaded_module.name == module_name + assert loaded_module.display_name == module.display_name - # test module setup - loaded_module.setup(config={}) - - assert loaded_module.config == {} + # check that default settings are applied + default_config = module.configs + assert loaded_module.name in loaded_module.config.keys() @pytest.mark.parametrize("module_name", ["cli_feeder", "local_storage", "generic_extractor", "html_formatter", "csv_db"]) def test_lazy_base_module(module_name): - lazy_module = get_module(module_name) + lazy_module = get_module_lazy(module_name) assert lazy_module is not None assert isinstance(lazy_module, LazyBaseModule) From 00a7018f365b651b9b36da93319b94c4c71a375e Mon Sep 17 00:00:00 2001 From: Patrick Robertson Date: Wed, 29 Jan 2025 19:25:22 +0100 Subject: [PATCH 02/62] Fix up dependency checking (use 'dependencies' instead of 'external_dependencies' -> simpler/easier to remember --- src/auto_archiver/core/module.py | 4 ++ .../modules/api_db/__manifest__.py | 2 +- .../modules/atlos/__manifest__.py | 2 +- .../modules/atlos_db/__manifest__.py | 2 +- .../modules/atlos_feeder/__manifest__.py | 2 +- .../modules/cli_feeder/__manifest__.py | 2 +- .../modules/console_db/__manifest__.py | 2 +- .../modules/csv_db/__manifest__.py | 2 +- .../modules/csv_feeder/__manifest__.py | 2 +- .../modules/gdrive_storage/__manifest__.py | 2 +- .../modules/gsheet_db/__manifest__.py | 2 +- .../modules/gsheet_feeder/__manifest__.py | 2 +- .../modules/hash_enricher/__manifest__.py | 2 +- .../modules/html_formatter/__manifest__.py | 4 +- .../modules/html_formatter/html_formatter.py | 1 + .../instagram_api_extractor/__manifest__.py | 2 +- .../instagram_extractor/__manifest__.py | 2 +- .../instagram_tbot_extractor/__manifest__.py | 2 +- .../modules/local_storage/__manifest__.py | 2 +- .../modules/meta_enricher/__manifest__.py | 2 +- .../modules/metadata_enricher/__manifest__.py | 2 +- .../modules/mute_formatter/__manifest__.py | 2 +- .../modules/pdq_hash_enricher/__manifest__.py | 2 +- .../modules/s3_storage/__manifest__.py | 2 +- .../screenshot_enricher/__manifest__.py | 2 +- .../modules/ssl_enricher/__manifest__.py | 2 +- .../telegram_extractor/__manifest__.py | 2 +- .../telethon_extractor/__manifest__.py | 2 +- .../thumbnail_enricher/__manifest__.py | 2 +- .../timestamping_enricher/__manifest__.py | 2 +- .../twitter_api_extractor/__manifest__.py | 2 +- .../modules/vk_extractor/__manifest__.py | 2 +- .../modules/wacz_enricher/__manifest__.py | 2 +- .../modules/wayback_enricher/__manifest__.py | 2 +- .../modules/whisper_enricher/__manifest__.py | 2 +- tests/data/example_module/__manifest__.py | 2 +- tests/data/example_module/example_module.py | 4 +- tests/test_modules.py | 51 ++++++++++++++----- 38 files changed, 81 insertions(+), 49 deletions(-) diff --git a/src/auto_archiver/core/module.py b/src/auto_archiver/core/module.py index 0888378..cb380cf 100644 --- a/src/auto_archiver/core/module.py +++ b/src/auto_archiver/core/module.py @@ -143,6 +143,7 @@ def available_modules(with_manifest: bool=False, limit_to_modules: List[str]= [] if _LAZY_LOADED_MODULES.get(possible_module): continue lazy_module = LazyBaseModule(possible_module, possible_module_path) + _LAZY_LOADED_MODULES[possible_module] = lazy_module all_modules.append(lazy_module) @@ -229,6 +230,9 @@ class LazyBaseModule: # check external dependencies are installed def check_deps(deps, check): for dep in deps: + if not len(dep): + # clear out any empty strings that a user may have erroneously added + continue if not check(dep): logger.error(f"Module '{self.name}' requires external dependency '{dep}' which is not available/setup. Have you installed the required dependencies for the '{self.name}' module? See the README for more information.") exit(1) diff --git a/src/auto_archiver/modules/api_db/__manifest__.py b/src/auto_archiver/modules/api_db/__manifest__.py index c89165f..d22fa59 100644 --- a/src/auto_archiver/modules/api_db/__manifest__.py +++ b/src/auto_archiver/modules/api_db/__manifest__.py @@ -3,7 +3,7 @@ "type": ["database"], "entry_point": "api_db:AAApiDb", "requires_setup": True, - "external_dependencies": { + "dependencies": { "python": ["requests", "loguru"], }, diff --git a/src/auto_archiver/modules/atlos/__manifest__.py b/src/auto_archiver/modules/atlos/__manifest__.py index 459fefe..7ba2f72 100644 --- a/src/auto_archiver/modules/atlos/__manifest__.py +++ b/src/auto_archiver/modules/atlos/__manifest__.py @@ -2,7 +2,7 @@ "name": "atlos_storage", "type": ["storage"], "requires_setup": True, - "external_dependencies": {"python": ["loguru", "requests"], "bin": [""]}, + "dependencies": {"python": ["loguru", "requests"], "bin": [""]}, "configs": { "path_generator": { "default": "url", diff --git a/src/auto_archiver/modules/atlos_db/__manifest__.py b/src/auto_archiver/modules/atlos_db/__manifest__.py index 42ce560..8f9473f 100644 --- a/src/auto_archiver/modules/atlos_db/__manifest__.py +++ b/src/auto_archiver/modules/atlos_db/__manifest__.py @@ -3,7 +3,7 @@ "type": ["database"], "entry_point": "atlos_db:AtlosDb", "requires_setup": True, - "external_dependencies": + "dependencies": {"python": ["loguru", ""], "bin": [""]}, diff --git a/src/auto_archiver/modules/atlos_feeder/__manifest__.py b/src/auto_archiver/modules/atlos_feeder/__manifest__.py index 0d90c8b..f2772f2 100644 --- a/src/auto_archiver/modules/atlos_feeder/__manifest__.py +++ b/src/auto_archiver/modules/atlos_feeder/__manifest__.py @@ -2,7 +2,7 @@ "name": "Atlos Feeder", "type": ["feeder"], "requires_setup": True, - "external_dependencies": { + "dependencies": { "python": ["loguru", "requests"], }, "configs": { diff --git a/src/auto_archiver/modules/cli_feeder/__manifest__.py b/src/auto_archiver/modules/cli_feeder/__manifest__.py index fe784c3..cf5c1b7 100644 --- a/src/auto_archiver/modules/cli_feeder/__manifest__.py +++ b/src/auto_archiver/modules/cli_feeder/__manifest__.py @@ -2,7 +2,7 @@ "name": "CLI Feeder", "type": ["feeder"], "requires_setup": False, - "external_dependencies": { + "dependencies": { "python": ["loguru"], }, 'entry_point': 'cli_feeder::CLIFeeder', diff --git a/src/auto_archiver/modules/console_db/__manifest__.py b/src/auto_archiver/modules/console_db/__manifest__.py index cd40496..a1d0d48 100644 --- a/src/auto_archiver/modules/console_db/__manifest__.py +++ b/src/auto_archiver/modules/console_db/__manifest__.py @@ -2,7 +2,7 @@ "name": "Console Database", "type": ["database"], "requires_setup": False, - "external_dependencies": { + "dependencies": { "python": ["loguru"], }, "description": """ diff --git a/src/auto_archiver/modules/csv_db/__manifest__.py b/src/auto_archiver/modules/csv_db/__manifest__.py index 3131188..507ce14 100644 --- a/src/auto_archiver/modules/csv_db/__manifest__.py +++ b/src/auto_archiver/modules/csv_db/__manifest__.py @@ -2,7 +2,7 @@ "name": "CSV Database", "type": ["database"], "requires_setup": False, - "external_dependencies": {"python": ["loguru"] + "dependencies": {"python": ["loguru"] }, 'entry_point': 'csv_db::CSVDb', "configs": { diff --git a/src/auto_archiver/modules/csv_feeder/__manifest__.py b/src/auto_archiver/modules/csv_feeder/__manifest__.py index 81c4dcd..b062ee6 100644 --- a/src/auto_archiver/modules/csv_feeder/__manifest__.py +++ b/src/auto_archiver/modules/csv_feeder/__manifest__.py @@ -2,7 +2,7 @@ "name": "CSV Feeder", "type": ["feeder"], "requires_setup": False, - "external_dependencies": { + "dependencies": { "python": ["loguru"], "bin": [""] }, diff --git a/src/auto_archiver/modules/gdrive_storage/__manifest__.py b/src/auto_archiver/modules/gdrive_storage/__manifest__.py index b81b717..e24f21b 100644 --- a/src/auto_archiver/modules/gdrive_storage/__manifest__.py +++ b/src/auto_archiver/modules/gdrive_storage/__manifest__.py @@ -2,7 +2,7 @@ "name": "Google Drive Storage", "type": ["storage"], "requires_setup": True, - "external_dependencies": { + "dependencies": { "python": [ "loguru", "google-api-python-client", diff --git a/src/auto_archiver/modules/gsheet_db/__manifest__.py b/src/auto_archiver/modules/gsheet_db/__manifest__.py index f2f1c35..f926adc 100644 --- a/src/auto_archiver/modules/gsheet_db/__manifest__.py +++ b/src/auto_archiver/modules/gsheet_db/__manifest__.py @@ -3,7 +3,7 @@ "type": ["database"], "entry_point": "gsheet_db::GsheetsDb", "requires_setup": True, - "external_dependencies": { + "dependencies": { "python": ["loguru", "gspread", "python-slugify"], }, "configs": { diff --git a/src/auto_archiver/modules/gsheet_feeder/__manifest__.py b/src/auto_archiver/modules/gsheet_feeder/__manifest__.py index 3d9cb08..1c9acab 100644 --- a/src/auto_archiver/modules/gsheet_feeder/__manifest__.py +++ b/src/auto_archiver/modules/gsheet_feeder/__manifest__.py @@ -3,7 +3,7 @@ "type": ["feeder"], "entry_point": "gsheet_feeder::GsheetsFeeder", "requires_setup": True, - "external_dependencies": { + "dependencies": { "python": ["loguru", "gspread", "python-slugify"], }, "configs": { diff --git a/src/auto_archiver/modules/hash_enricher/__manifest__.py b/src/auto_archiver/modules/hash_enricher/__manifest__.py index f306808..c7a023e 100644 --- a/src/auto_archiver/modules/hash_enricher/__manifest__.py +++ b/src/auto_archiver/modules/hash_enricher/__manifest__.py @@ -2,7 +2,7 @@ "name": "Hash Enricher", "type": ["enricher"], "requires_setup": False, - "external_dependencies": { + "dependencies": { "python": ["loguru"], }, "configs": { diff --git a/src/auto_archiver/modules/html_formatter/__manifest__.py b/src/auto_archiver/modules/html_formatter/__manifest__.py index 259a3d1..ec19cf8 100644 --- a/src/auto_archiver/modules/html_formatter/__manifest__.py +++ b/src/auto_archiver/modules/html_formatter/__manifest__.py @@ -2,8 +2,8 @@ "name": "HTML Formatter", "type": ["formatter"], "requires_setup": False, - "external_dependencies": { - "python": ["loguru", "jinja2"], + "dependencies": { + "python": ["hash_enricher", "loguru", "jinja2"], "bin": [""] }, "configs": { diff --git a/src/auto_archiver/modules/html_formatter/html_formatter.py b/src/auto_archiver/modules/html_formatter/html_formatter.py index 570fc6f..8f006e0 100644 --- a/src/auto_archiver/modules/html_formatter/html_formatter.py +++ b/src/auto_archiver/modules/html_formatter/html_formatter.py @@ -53,6 +53,7 @@ class HtmlFormatter(Formatter): outf.write(content) final_media = Media(filename=html_path, _mimetype="text/html") + # get the already instantiated hash_enricher module he = get_module('hash_enricher', self.config) if len(hd := he.calculate_hash(final_media.filename)): final_media.set("hash", f"{he.algorithm}:{hd}") diff --git a/src/auto_archiver/modules/instagram_api_extractor/__manifest__.py b/src/auto_archiver/modules/instagram_api_extractor/__manifest__.py index cdaf635..57f378e 100644 --- a/src/auto_archiver/modules/instagram_api_extractor/__manifest__.py +++ b/src/auto_archiver/modules/instagram_api_extractor/__manifest__.py @@ -1,7 +1,7 @@ { "name": "Instagram API Extractor", "type": ["extractor"], - "external_dependencies": + "dependencies": {"python": ["requests", "loguru", "retrying", diff --git a/src/auto_archiver/modules/instagram_extractor/__manifest__.py b/src/auto_archiver/modules/instagram_extractor/__manifest__.py index f1857c2..6e7518e 100644 --- a/src/auto_archiver/modules/instagram_extractor/__manifest__.py +++ b/src/auto_archiver/modules/instagram_extractor/__manifest__.py @@ -1,7 +1,7 @@ { "name": "Instagram Extractor", "type": ["extractor"], - "external_dependencies": { + "dependencies": { "python": [ "instaloader", "loguru", diff --git a/src/auto_archiver/modules/instagram_tbot_extractor/__manifest__.py b/src/auto_archiver/modules/instagram_tbot_extractor/__manifest__.py index 95d6808..8a1f74f 100644 --- a/src/auto_archiver/modules/instagram_tbot_extractor/__manifest__.py +++ b/src/auto_archiver/modules/instagram_tbot_extractor/__manifest__.py @@ -1,7 +1,7 @@ { "name": "Instagram Telegram Bot Extractor", "type": ["extractor"], - "external_dependencies": {"python": ["loguru", + "dependencies": {"python": ["loguru", "telethon",], }, "requires_setup": True, diff --git a/src/auto_archiver/modules/local_storage/__manifest__.py b/src/auto_archiver/modules/local_storage/__manifest__.py index ce00953..6d9cf53 100644 --- a/src/auto_archiver/modules/local_storage/__manifest__.py +++ b/src/auto_archiver/modules/local_storage/__manifest__.py @@ -2,7 +2,7 @@ "name": "Local Storage", "type": ["storage"], "requires_setup": False, - "external_dependencies": { + "dependencies": { "python": ["loguru"], }, "configs": { diff --git a/src/auto_archiver/modules/meta_enricher/__manifest__.py b/src/auto_archiver/modules/meta_enricher/__manifest__.py index 10acf71..37c9201 100644 --- a/src/auto_archiver/modules/meta_enricher/__manifest__.py +++ b/src/auto_archiver/modules/meta_enricher/__manifest__.py @@ -2,7 +2,7 @@ "name": "Archive Metadata Enricher", "type": ["enricher"], "requires_setup": False, - "external_dependencies": { + "dependencies": { "python": ["loguru"], }, "description": """ diff --git a/src/auto_archiver/modules/metadata_enricher/__manifest__.py b/src/auto_archiver/modules/metadata_enricher/__manifest__.py index 50064e9..f8ccdc6 100644 --- a/src/auto_archiver/modules/metadata_enricher/__manifest__.py +++ b/src/auto_archiver/modules/metadata_enricher/__manifest__.py @@ -2,7 +2,7 @@ "name": "Media Metadata Enricher", "type": ["enricher"], "requires_setup": True, - "external_dependencies": { + "dependencies": { "python": ["loguru"], "bin": ["exiftool"] }, diff --git a/src/auto_archiver/modules/mute_formatter/__manifest__.py b/src/auto_archiver/modules/mute_formatter/__manifest__.py index 77f2784..e81dc4c 100644 --- a/src/auto_archiver/modules/mute_formatter/__manifest__.py +++ b/src/auto_archiver/modules/mute_formatter/__manifest__.py @@ -2,7 +2,7 @@ "name": "Mute Formatter", "type": ["formatter"], "requires_setup": True, - "external_dependencies": { + "dependencies": { }, "description": """ Default formatter. """, diff --git a/src/auto_archiver/modules/pdq_hash_enricher/__manifest__.py b/src/auto_archiver/modules/pdq_hash_enricher/__manifest__.py index 7b418b1..6353d12 100644 --- a/src/auto_archiver/modules/pdq_hash_enricher/__manifest__.py +++ b/src/auto_archiver/modules/pdq_hash_enricher/__manifest__.py @@ -2,7 +2,7 @@ "name": "PDQ Hash Enricher", "type": ["enricher"], "requires_setup": False, - "external_dependencies": { + "dependencies": { "python": ["loguru", "pdqhash", "numpy", "Pillow"], }, "description": """ diff --git a/src/auto_archiver/modules/s3_storage/__manifest__.py b/src/auto_archiver/modules/s3_storage/__manifest__.py index 811c703..16ac7bd 100644 --- a/src/auto_archiver/modules/s3_storage/__manifest__.py +++ b/src/auto_archiver/modules/s3_storage/__manifest__.py @@ -2,7 +2,7 @@ "name": "S3 Storage", "type": ["storage"], "requires_setup": True, - "external_dependencies": { + "dependencies": { "python": ["boto3", "loguru"], }, "configs": { diff --git a/src/auto_archiver/modules/screenshot_enricher/__manifest__.py b/src/auto_archiver/modules/screenshot_enricher/__manifest__.py index c1a30e7..52842c9 100644 --- a/src/auto_archiver/modules/screenshot_enricher/__manifest__.py +++ b/src/auto_archiver/modules/screenshot_enricher/__manifest__.py @@ -2,7 +2,7 @@ "name": "Screenshot Enricher", "type": ["enricher"], "requires_setup": True, - "external_dependencies": { + "dependencies": { "python": ["loguru", "selenium"], "bin": ["chromedriver"] }, diff --git a/src/auto_archiver/modules/ssl_enricher/__manifest__.py b/src/auto_archiver/modules/ssl_enricher/__manifest__.py index ccde957..0fb7cd9 100644 --- a/src/auto_archiver/modules/ssl_enricher/__manifest__.py +++ b/src/auto_archiver/modules/ssl_enricher/__manifest__.py @@ -2,7 +2,7 @@ "name": "SSL Certificate Enricher", "type": ["enricher"], "requires_setup": False, - "external_dependencies": { + "dependencies": { "python": ["loguru", "python-slugify"], }, 'entry_point': 'ssl_enricher::SSLEnricher', diff --git a/src/auto_archiver/modules/telegram_extractor/__manifest__.py b/src/auto_archiver/modules/telegram_extractor/__manifest__.py index 86b5e0f..e1c49c2 100644 --- a/src/auto_archiver/modules/telegram_extractor/__manifest__.py +++ b/src/auto_archiver/modules/telegram_extractor/__manifest__.py @@ -2,7 +2,7 @@ "name": "Telegram Extractor", "type": ["extractor"], "requires_setup": False, - "external_dependencies": { + "dependencies": { "python": [ "requests", "bs4", diff --git a/src/auto_archiver/modules/telethon_extractor/__manifest__.py b/src/auto_archiver/modules/telethon_extractor/__manifest__.py index 2cf1e42..6b37654 100644 --- a/src/auto_archiver/modules/telethon_extractor/__manifest__.py +++ b/src/auto_archiver/modules/telethon_extractor/__manifest__.py @@ -2,7 +2,7 @@ "name": "telethon_extractor", "type": ["extractor"], "requires_setup": True, - "external_dependencies": { + "dependencies": { "python": ["telethon", "loguru", "tqdm", diff --git a/src/auto_archiver/modules/thumbnail_enricher/__manifest__.py b/src/auto_archiver/modules/thumbnail_enricher/__manifest__.py index 2b0f167..bd7836d 100644 --- a/src/auto_archiver/modules/thumbnail_enricher/__manifest__.py +++ b/src/auto_archiver/modules/thumbnail_enricher/__manifest__.py @@ -2,7 +2,7 @@ "name": "Thumbnail Enricher", "type": ["enricher"], "requires_setup": False, - "external_dependencies": { + "dependencies": { "python": ["loguru", "ffmpeg-python"], "bin": ["ffmpeg"] }, diff --git a/src/auto_archiver/modules/timestamping_enricher/__manifest__.py b/src/auto_archiver/modules/timestamping_enricher/__manifest__.py index 496d211..6ad9c57 100644 --- a/src/auto_archiver/modules/timestamping_enricher/__manifest__.py +++ b/src/auto_archiver/modules/timestamping_enricher/__manifest__.py @@ -2,7 +2,7 @@ "name": "Timestamping Enricher", "type": ["enricher"], "requires_setup": True, - "external_dependencies": { + "dependencies": { "python": [ "loguru", "slugify", diff --git a/src/auto_archiver/modules/twitter_api_extractor/__manifest__.py b/src/auto_archiver/modules/twitter_api_extractor/__manifest__.py index 02d0d6c..05d1ac0 100644 --- a/src/auto_archiver/modules/twitter_api_extractor/__manifest__.py +++ b/src/auto_archiver/modules/twitter_api_extractor/__manifest__.py @@ -2,7 +2,7 @@ "name": "Twitter API Extractor", "type": ["extractor"], "requires_setup": True, - "external_dependencies": { + "dependencies": { "python": ["requests", "loguru", "pytwitter", diff --git a/src/auto_archiver/modules/vk_extractor/__manifest__.py b/src/auto_archiver/modules/vk_extractor/__manifest__.py index bdcaf99..116b430 100644 --- a/src/auto_archiver/modules/vk_extractor/__manifest__.py +++ b/src/auto_archiver/modules/vk_extractor/__manifest__.py @@ -3,7 +3,7 @@ "type": ["extractor"], "requires_setup": True, "depends": ["core", "utils"], - "external_dependencies": { + "dependencies": { "python": ["loguru", "vk_url_scraper"], }, diff --git a/src/auto_archiver/modules/wacz_enricher/__manifest__.py b/src/auto_archiver/modules/wacz_enricher/__manifest__.py index 07983d9..bb9d290 100644 --- a/src/auto_archiver/modules/wacz_enricher/__manifest__.py +++ b/src/auto_archiver/modules/wacz_enricher/__manifest__.py @@ -2,7 +2,7 @@ "name": "WACZ Enricher", "type": ["enricher", "archiver"], "requires_setup": True, - "external_dependencies": { + "dependencies": { "python": [ "loguru", "jsonlines", diff --git a/src/auto_archiver/modules/wayback_enricher/__manifest__.py b/src/auto_archiver/modules/wayback_enricher/__manifest__.py index bff10af..5d1fe25 100644 --- a/src/auto_archiver/modules/wayback_enricher/__manifest__.py +++ b/src/auto_archiver/modules/wayback_enricher/__manifest__.py @@ -2,7 +2,7 @@ "name": "Wayback Machine Enricher", "type": ["enricher", "archiver"], "requires_setup": True, - "external_dependencies": { + "dependencies": { "python": ["loguru", "requests"], }, "entry_point": "wayback_enricher::WaybackExtractorEnricher", diff --git a/src/auto_archiver/modules/whisper_enricher/__manifest__.py b/src/auto_archiver/modules/whisper_enricher/__manifest__.py index 25eae25..0adf9ff 100644 --- a/src/auto_archiver/modules/whisper_enricher/__manifest__.py +++ b/src/auto_archiver/modules/whisper_enricher/__manifest__.py @@ -2,7 +2,7 @@ "name": "Whisper Enricher", "type": ["enricher"], "requires_setup": True, - "external_dependencies": { + "dependencies": { "python": ["loguru", "requests"], }, "configs": { diff --git a/tests/data/example_module/__manifest__.py b/tests/data/example_module/__manifest__.py index ca3a678..19a85f9 100644 --- a/tests/data/example_module/__manifest__.py +++ b/tests/data/example_module/__manifest__.py @@ -2,7 +2,7 @@ "name": "Example Module", "type": ["extractor"], "requires_setup": False, - "external_dependencies": {"python": ["loguru"] + "dependencies": {"python": ["loguru"] }, "configs": { "csv_file": {"default": "db.csv", "help": "CSV file name"} diff --git a/tests/data/example_module/example_module.py b/tests/data/example_module/example_module.py index b752743..bce8ba4 100644 --- a/tests/data/example_module/example_module.py +++ b/tests/data/example_module/example_module.py @@ -1,4 +1,4 @@ from auto_archiver.core.extractor import Extractor - class ExampleModule(Extractor): - pass \ No newline at end of file + def download(self, item): + print("do something") \ No newline at end of file diff --git a/tests/test_modules.py b/tests/test_modules.py index 109bc52..decc616 100644 --- a/tests/test_modules.py +++ b/tests/test_modules.py @@ -1,13 +1,24 @@ import sys import pytest from auto_archiver.core.module import get_module_lazy, BaseModule, LazyBaseModule, _LAZY_LOADED_MODULES -from auto_archiver.core.extractor import Extractor @pytest.fixture def example_module(): - yield get_module_lazy("example_module", ["tests/data/"]) + import auto_archiver + + previous_path = auto_archiver.modules.__path__ + auto_archiver.modules.__path__.append("tests/data/") + + module = get_module_lazy("example_module") + yield module # cleanup - _LAZY_LOADED_MODULES.pop("example_module") + try: + del module._manifest + except AttributeError: + pass + del _LAZY_LOADED_MODULES["example_module"] + sys.modules.pop("auto_archiver.modules.example_module.example_module", None) + auto_archiver.modules.__path__ = previous_path def test_get_module_lazy(example_module): assert example_module.name == "example_module" @@ -15,18 +26,34 @@ def test_get_module_lazy(example_module): assert example_module.manifest is not None +def test_python_dependency_check(example_module): + # example_module requires loguru, which is not installed + # monkey patch the manifest to include a nonexistnet dependency + example_module.manifest["dependencies"]["python"] = ["does_not_exist"] -def test_load_module_abc_check(example_module): - - # example_module is an extractor but doesn't have the 'download' method, should raise an ABC error - with pytest.raises(TypeError) as load_error: + with pytest.raises(SystemExit) as load_error: example_module.load({}) - assert "Can't instantiate abstract class ExampleModule with abstract method download" in str(load_error.value) - -def test_load_module(example_module, monkeypatch): - # hack - remove the 'download' method from the required methods of Extractor - monkeypatch.setattr(Extractor, "__abstractmethods__", set()) + assert load_error.value.code == 1 + +def test_binary_dependency_check(example_module): + # example_module requires ffmpeg, which is not installed + # monkey patch the manifest to include a nonexistnet dependency + example_module.manifest["dependencies"]["binary"] = ["does_not_exist"] + +def test_module_dependency_check_loads_module(example_module): + # example_module requires cli_feeder, which is not installed + # monkey patch the manifest to include a nonexistnet dependency + example_module.manifest["dependencies"]["python"] = ["hash_enricher"] + + loaded_module = example_module.load({}) + assert loaded_module is not None + + # check the dependency is loaded + assert _LAZY_LOADED_MODULES["hash_enricher"] is not None + assert _LAZY_LOADED_MODULES["hash_enricher"]._instance is not None + +def test_load_module(example_module): # setup the module, and check that config is set to the default values loaded_module = example_module.load({}) From 18ff36ce154dad2f083c9f4488b0f8027bcbe278 Mon Sep 17 00:00:00 2001 From: Patrick Robertson Date: Wed, 29 Jan 2025 19:37:41 +0100 Subject: [PATCH 03/62] Add ruamel to dependencies (replaces pyyaml) --- poetry.lock | 80 ++++++++++++++++++++++++++++++++++++++++++++++++-- pyproject.toml | 2 +- 2 files changed, 79 insertions(+), 3 deletions(-) diff --git a/poetry.lock b/poetry.lock index 6d6ad8c..e8a899a 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1818,7 +1818,7 @@ version = "6.0.2" description = "YAML parser and emitter for Python" optional = false python-versions = ">=3.8" -groups = ["main", "docs"] +groups = ["docs"] files = [ {file = "PyYAML-6.0.2-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:0a9a2848a5b7feac301353437eb7d5957887edbf81d56e903999a75a3d743086"}, {file = "PyYAML-6.0.2-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:29717114e51c84ddfba879543fb232a6ed60086602313ca38cce623c1d62cfbf"}, @@ -2086,6 +2086,82 @@ files = [ [package.dependencies] pyasn1 = ">=0.1.3" +[[package]] +name = "ruamel-yaml" +version = "0.18.10" +description = "ruamel.yaml is a YAML parser/emitter that supports roundtrip preservation of comments, seq/map flow style, and map key order" +optional = false +python-versions = ">=3.7" +groups = ["main"] +files = [ + {file = "ruamel.yaml-0.18.10-py3-none-any.whl", hash = "sha256:30f22513ab2301b3d2b577adc121c6471f28734d3d9728581245f1e76468b4f1"}, + {file = "ruamel.yaml-0.18.10.tar.gz", hash = "sha256:20c86ab29ac2153f80a428e1254a8adf686d3383df04490514ca3b79a362db58"}, +] + +[package.dependencies] +"ruamel.yaml.clib" = {version = ">=0.2.7", markers = "platform_python_implementation == \"CPython\" and python_version < \"3.13\""} + +[package.extras] +docs = ["mercurial (>5.7)", "ryd"] +jinja2 = ["ruamel.yaml.jinja2 (>=0.2)"] + +[[package]] +name = "ruamel-yaml-clib" +version = "0.2.12" +description = "C version of reader, parser and emitter for ruamel.yaml derived from libyaml" +optional = false +python-versions = ">=3.9" +groups = ["main"] +markers = "platform_python_implementation == \"CPython\"" +files = [ + {file = "ruamel.yaml.clib-0.2.12-cp310-cp310-macosx_13_0_arm64.whl", hash = "sha256:11f891336688faf5156a36293a9c362bdc7c88f03a8a027c2c1d8e0bcde998e5"}, + {file = "ruamel.yaml.clib-0.2.12-cp310-cp310-manylinux2014_aarch64.whl", hash = "sha256:a606ef75a60ecf3d924613892cc603b154178ee25abb3055db5062da811fd969"}, + {file = "ruamel.yaml.clib-0.2.12-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fd5415dded15c3822597455bc02bcd66e81ef8b7a48cb71a33628fc9fdde39df"}, + {file = "ruamel.yaml.clib-0.2.12-cp310-cp310-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:f66efbc1caa63c088dead1c4170d148eabc9b80d95fb75b6c92ac0aad2437d76"}, + {file = "ruamel.yaml.clib-0.2.12-cp310-cp310-musllinux_1_1_i686.whl", hash = "sha256:22353049ba4181685023b25b5b51a574bce33e7f51c759371a7422dcae5402a6"}, + {file = "ruamel.yaml.clib-0.2.12-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:932205970b9f9991b34f55136be327501903f7c66830e9760a8ffb15b07f05cd"}, + {file = "ruamel.yaml.clib-0.2.12-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:a52d48f4e7bf9005e8f0a89209bf9a73f7190ddf0489eee5eb51377385f59f2a"}, + {file = "ruamel.yaml.clib-0.2.12-cp310-cp310-win32.whl", hash = "sha256:3eac5a91891ceb88138c113f9db04f3cebdae277f5d44eaa3651a4f573e6a5da"}, + {file = "ruamel.yaml.clib-0.2.12-cp310-cp310-win_amd64.whl", hash = "sha256:ab007f2f5a87bd08ab1499bdf96f3d5c6ad4dcfa364884cb4549aa0154b13a28"}, + {file = "ruamel.yaml.clib-0.2.12-cp311-cp311-macosx_13_0_arm64.whl", hash = "sha256:4a6679521a58256a90b0d89e03992c15144c5f3858f40d7c18886023d7943db6"}, + {file = "ruamel.yaml.clib-0.2.12-cp311-cp311-manylinux2014_aarch64.whl", hash = "sha256:d84318609196d6bd6da0edfa25cedfbabd8dbde5140a0a23af29ad4b8f91fb1e"}, + {file = "ruamel.yaml.clib-0.2.12-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:bb43a269eb827806502c7c8efb7ae7e9e9d0573257a46e8e952f4d4caba4f31e"}, + {file = "ruamel.yaml.clib-0.2.12-cp311-cp311-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:811ea1594b8a0fb466172c384267a4e5e367298af6b228931f273b111f17ef52"}, + {file = "ruamel.yaml.clib-0.2.12-cp311-cp311-musllinux_1_1_i686.whl", hash = "sha256:cf12567a7b565cbf65d438dec6cfbe2917d3c1bdddfce84a9930b7d35ea59642"}, + {file = "ruamel.yaml.clib-0.2.12-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:7dd5adc8b930b12c8fc5b99e2d535a09889941aa0d0bd06f4749e9a9397c71d2"}, + {file = "ruamel.yaml.clib-0.2.12-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:1492a6051dab8d912fc2adeef0e8c72216b24d57bd896ea607cb90bb0c4981d3"}, + {file = "ruamel.yaml.clib-0.2.12-cp311-cp311-win32.whl", hash = "sha256:bd0a08f0bab19093c54e18a14a10b4322e1eacc5217056f3c063bd2f59853ce4"}, + {file = "ruamel.yaml.clib-0.2.12-cp311-cp311-win_amd64.whl", hash = "sha256:a274fb2cb086c7a3dea4322ec27f4cb5cc4b6298adb583ab0e211a4682f241eb"}, + {file = "ruamel.yaml.clib-0.2.12-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:20b0f8dc160ba83b6dcc0e256846e1a02d044e13f7ea74a3d1d56ede4e48c632"}, + {file = "ruamel.yaml.clib-0.2.12-cp312-cp312-manylinux2014_aarch64.whl", hash = "sha256:943f32bc9dedb3abff9879edc134901df92cfce2c3d5c9348f172f62eb2d771d"}, + {file = "ruamel.yaml.clib-0.2.12-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:95c3829bb364fdb8e0332c9931ecf57d9be3519241323c5274bd82f709cebc0c"}, + {file = "ruamel.yaml.clib-0.2.12-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:749c16fcc4a2b09f28843cda5a193e0283e47454b63ec4b81eaa2242f50e4ccd"}, + {file = "ruamel.yaml.clib-0.2.12-cp312-cp312-musllinux_1_1_i686.whl", hash = "sha256:bf165fef1f223beae7333275156ab2022cffe255dcc51c27f066b4370da81e31"}, + {file = "ruamel.yaml.clib-0.2.12-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:32621c177bbf782ca5a18ba4d7af0f1082a3f6e517ac2a18b3974d4edf349680"}, + {file = "ruamel.yaml.clib-0.2.12-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:b82a7c94a498853aa0b272fd5bc67f29008da798d4f93a2f9f289feb8426a58d"}, + {file = "ruamel.yaml.clib-0.2.12-cp312-cp312-win32.whl", hash = "sha256:e8c4ebfcfd57177b572e2040777b8abc537cdef58a2120e830124946aa9b42c5"}, + {file = "ruamel.yaml.clib-0.2.12-cp312-cp312-win_amd64.whl", hash = "sha256:0467c5965282c62203273b838ae77c0d29d7638c8a4e3a1c8bdd3602c10904e4"}, + {file = "ruamel.yaml.clib-0.2.12-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:4c8c5d82f50bb53986a5e02d1b3092b03622c02c2eb78e29bec33fd9593bae1a"}, + {file = "ruamel.yaml.clib-0.2.12-cp313-cp313-manylinux2014_aarch64.whl", hash = "sha256:e7e3736715fbf53e9be2a79eb4db68e4ed857017344d697e8b9749444ae57475"}, + {file = "ruamel.yaml.clib-0.2.12-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0b7e75b4965e1d4690e93021adfcecccbca7d61c7bddd8e22406ef2ff20d74ef"}, + {file = "ruamel.yaml.clib-0.2.12-cp313-cp313-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:96777d473c05ee3e5e3c3e999f5d23c6f4ec5b0c38c098b3a5229085f74236c6"}, + {file = "ruamel.yaml.clib-0.2.12-cp313-cp313-musllinux_1_1_i686.whl", hash = "sha256:3bc2a80e6420ca8b7d3590791e2dfc709c88ab9152c00eeb511c9875ce5778bf"}, + {file = "ruamel.yaml.clib-0.2.12-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:e188d2699864c11c36cdfdada94d781fd5d6b0071cd9c427bceb08ad3d7c70e1"}, + {file = "ruamel.yaml.clib-0.2.12-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:4f6f3eac23941b32afccc23081e1f50612bdbe4e982012ef4f5797986828cd01"}, + {file = "ruamel.yaml.clib-0.2.12-cp313-cp313-win32.whl", hash = "sha256:6442cb36270b3afb1b4951f060eccca1ce49f3d087ca1ca4563a6eb479cb3de6"}, + {file = "ruamel.yaml.clib-0.2.12-cp313-cp313-win_amd64.whl", hash = "sha256:e5b8daf27af0b90da7bb903a876477a9e6d7270be6146906b276605997c7e9a3"}, + {file = "ruamel.yaml.clib-0.2.12-cp39-cp39-macosx_12_0_arm64.whl", hash = "sha256:fc4b630cd3fa2cf7fce38afa91d7cfe844a9f75d7f0f36393fa98815e911d987"}, + {file = "ruamel.yaml.clib-0.2.12-cp39-cp39-manylinux2014_aarch64.whl", hash = "sha256:bc5f1e1c28e966d61d2519f2a3d451ba989f9ea0f2307de7bc45baa526de9e45"}, + {file = "ruamel.yaml.clib-0.2.12-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:5a0e060aace4c24dcaf71023bbd7d42674e3b230f7e7b97317baf1e953e5b519"}, + {file = "ruamel.yaml.clib-0.2.12-cp39-cp39-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:e2f1c3765db32be59d18ab3953f43ab62a761327aafc1594a2a1fbe038b8b8a7"}, + {file = "ruamel.yaml.clib-0.2.12-cp39-cp39-musllinux_1_1_i686.whl", hash = "sha256:d85252669dc32f98ebcd5d36768f5d4faeaeaa2d655ac0473be490ecdae3c285"}, + {file = "ruamel.yaml.clib-0.2.12-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:e143ada795c341b56de9418c58d028989093ee611aa27ffb9b7f609c00d813ed"}, + {file = "ruamel.yaml.clib-0.2.12-cp39-cp39-musllinux_1_2_aarch64.whl", hash = "sha256:2c59aa6170b990d8d2719323e628aaf36f3bfbc1c26279c0eeeb24d05d2d11c7"}, + {file = "ruamel.yaml.clib-0.2.12-cp39-cp39-win32.whl", hash = "sha256:beffaed67936fbbeffd10966a4eb53c402fafd3d6833770516bf7314bc6ffa12"}, + {file = "ruamel.yaml.clib-0.2.12-cp39-cp39-win_amd64.whl", hash = "sha256:040ae85536960525ea62868b642bdb0c2cc6021c9f9d507810c0c604e66f5a7b"}, + {file = "ruamel.yaml.clib-0.2.12.tar.gz", hash = "sha256:6c8fbb13ec503f99a91901ab46e0b07ae7941cd527393187039aec586fdfd36f"}, +] + [[package]] name = "s3transfer" version = "0.11.2" @@ -3006,4 +3082,4 @@ test = ["pytest (>=8.1,<9.0)", "pytest-rerunfailures (>=14.0,<15.0)"] [metadata] lock-version = "2.1" python-versions = ">=3.10,<3.13" -content-hash = "d1af74e7fc7c919eda55dd383208edab906508353b4a9eff8e979967484823f8" +content-hash = "1556d53c5a94392c120ebaafc495d3b322daf64dac4a19f9726588c7f3d84bca" diff --git a/pyproject.toml b/pyproject.toml index ec78212..b3a2456 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,7 +37,6 @@ dependencies = [ "pdqhash (>=0.0.0)", "pillow (>=0.0.0)", "python-slugify (>=0.0.0)", - "pyyaml (>=0.0.0)", "dateparser (>=0.0.0)", "python-twitter-v2 (>=0.0.0)", "instaloader (>=0.0.0)", @@ -58,6 +57,7 @@ dependencies = [ "tsp-client (>=0.0.0)", "certvalidator (>=0.0.0)", "rich-argparse (>=1.6.0,<2.0.0)", + "ruamel-yaml (>=0.18.10,<0.19.0)", ] [tool.poetry.group.dev.dependencies] From cddae65a90a0dc225f9fcac26cdb5fce21448ccc Mon Sep 17 00:00:00 2001 From: erinhmclark Date: Thu, 30 Jan 2025 08:42:23 +0000 Subject: [PATCH 04/62] Update modules for new core structure. --- scripts/create_update_gdrive_oauth_token.py | 43 +++++---- scripts/telegram_setup.py | 29 ++++++ src/auto_archiver/core/orchestrator.py | 2 +- src/auto_archiver/core/storage.py | 2 +- src/auto_archiver/modules/api_db/__init__.py | 2 +- .../modules/api_db/__manifest__.py | 42 ++++++-- src/auto_archiver/modules/api_db/api_db.py | 16 +--- src/auto_archiver/modules/atlos/__init__.py | 1 - .../modules/atlos/__manifest__.py | 40 -------- .../modules/atlos_db/atlos_db.py | 8 +- .../modules/atlos_feeder/__manifest__.py | 3 +- .../modules/atlos_feeder/atlos_feeder.py | 11 +-- .../atlos_storage.py} | 10 +- .../modules/gdrive_storage/__manifest__.py | 73 ++++++++++++-- .../modules/gdrive_storage/gdrive_storage.py | 96 +++++++++---------- .../modules/gsheet_db/__manifest__.py | 3 +- .../modules/gsheet_feeder/__manifest__.py | 2 +- .../instagram_api_extractor/__manifest__.py | 13 ++- .../instagram_api_extractor.py | 9 +- .../instagram_extractor/__manifest__.py | 13 ++- .../instagram_extractor.py | 12 +-- .../instagram_tbot_extractor/__manifest__.py | 15 ++- .../instagram_tbot_extractor.py | 7 +- .../modules/pdq_hash_enricher/__manifest__.py | 2 +- .../modules/s3_storage/__init__.py | 2 +- .../modules/s3_storage/__manifest__.py | 12 ++- .../s3_storage/{s3.py => s3_storage.py} | 29 +++--- .../modules/ssl_enricher/__manifest__.py | 2 +- .../thumbnail_enricher/__manifest__.py | 2 +- .../modules/vk_extractor/__manifest__.py | 19 ++-- .../modules/vk_extractor/vk_extractor.py | 6 +- .../modules/wacz_enricher/__manifest__.py | 4 +- .../modules/wacz_enricher/wacz_enricher.py | 4 +- .../modules/whisper_enricher/__manifest__.py | 2 +- .../whisper_enricher/whisper_enricher.py | 4 +- 35 files changed, 307 insertions(+), 233 deletions(-) create mode 100644 scripts/telegram_setup.py delete mode 100644 src/auto_archiver/modules/atlos/__init__.py delete mode 100644 src/auto_archiver/modules/atlos/__manifest__.py rename src/auto_archiver/modules/{atlos/atlos.py => atlos_storage/atlos_storage.py} (96%) rename src/auto_archiver/modules/s3_storage/{s3.py => s3_storage.py} (88%) diff --git a/scripts/create_update_gdrive_oauth_token.py b/scripts/create_update_gdrive_oauth_token.py index ec8a120..eb6fdbe 100644 --- a/scripts/create_update_gdrive_oauth_token.py +++ b/scripts/create_update_gdrive_oauth_token.py @@ -12,7 +12,7 @@ from googleapiclient.errors import HttpError # Code below from https://developers.google.com/drive/api/quickstart/python # Example invocation: py scripts/create_update_gdrive_oauth_token.py -c secrets/credentials.json -t secrets/gd-token.json -SCOPES = ['https://www.googleapis.com/auth/drive'] +SCOPES = ["https://www.googleapis.com/auth/drive.file"] @click.command( @@ -23,7 +23,7 @@ SCOPES = ['https://www.googleapis.com/auth/drive'] "-c", type=click.Path(exists=True), help="path to the credentials.json file downloaded from https://console.cloud.google.com/apis/credentials", - required=True + required=True, ) @click.option( "--token", @@ -31,59 +31,62 @@ SCOPES = ['https://www.googleapis.com/auth/drive'] type=click.Path(exists=False), default="gd-token.json", help="file where to place the OAuth token, defaults to gd-token.json which you must then move to where your orchestration file points to, defaults to gd-token.json", - required=True + required=True, ) def main(credentials, token): # The file token.json stores the user's access and refresh tokens, and is # created automatically when the authorization flow completes for the first time. creds = None if os.path.exists(token): - with open(token, 'r') as stream: + with open(token, "r") as stream: creds_json = json.load(stream) # creds = Credentials.from_authorized_user_file(creds_json, SCOPES) - creds_json['refresh_token'] = creds_json.get("refresh_token", "") + creds_json["refresh_token"] = creds_json.get("refresh_token", "") creds = Credentials.from_authorized_user_info(creds_json, SCOPES) # If there are no (valid) credentials available, let the user log in. if not creds or not creds.valid: if creds and creds.expired and creds.refresh_token: - print('Requesting new token') + print("Requesting new token") creds.refresh(Request()) else: - print('First run through so putting up login dialog') + print("First run through so putting up login dialog") # credentials.json downloaded from https://console.cloud.google.com/apis/credentials flow = InstalledAppFlow.from_client_secrets_file(credentials, SCOPES) creds = flow.run_local_server(port=55192) # Save the credentials for the next run - with open(token, 'w') as token: - print('Saving new token') + with open(token, "w") as token: + print("Saving new token") token.write(creds.to_json()) else: - print('Token valid') + print("Token valid") try: - service = build('drive', 'v3', credentials=creds) + service = build("drive", "v3", credentials=creds) # About the user results = service.about().get(fields="*").execute() - emailAddress = results['user']['emailAddress'] + emailAddress = results["user"]["emailAddress"] print(emailAddress) # Call the Drive v3 API and return some files - results = service.files().list( - pageSize=10, fields="nextPageToken, files(id, name)").execute() - items = results.get('files', []) + results = ( + service.files() + .list(pageSize=10, fields="nextPageToken, files(id, name)") + .execute() + ) + items = results.get("files", []) if not items: - print('No files found.') + print("No files found.") return - print('Files:') + print("Files:") for item in items: - print(u'{0} ({1})'.format(item['name'], item['id'])) + print("{0} ({1})".format(item["name"], item["id"])) except HttpError as error: - print(f'An error occurred: {error}') + print(f"An error occurred: {error}") -if __name__ == '__main__': +if __name__ == "__main__": main() diff --git a/scripts/telegram_setup.py b/scripts/telegram_setup.py new file mode 100644 index 0000000..e6fa43c --- /dev/null +++ b/scripts/telegram_setup.py @@ -0,0 +1,29 @@ +""" +This script is used to create a new session file for the Telegram client. +To do this you must first create a Telegram application at https://my.telegram.org/apps +And store your id and hash in the environment variables TELEGRAM_API_ID and TELEGRAM_API_HASH. +Create a .env file, or add the following to your environment : +``` +export TELEGRAM_API_ID=[YOUR_ID_HERE] +export TELEGRAM_API_HASH=[YOUR_HASH_HERE] +``` +Then run this script to create a new session file. + +You will need to provide your phone number and a 2FA code the first time you run this script. +""" + + +import os +from telethon.sync import TelegramClient +from loguru import logger + + +# Create a +API_ID = os.getenv("TELEGRAM_API_ID") +API_HASH = os.getenv("TELEGRAM_API_HASH") +SESSION_FILE = "secrets/anon-insta" + +os.makedirs("secrets", exist_ok=True) +with TelegramClient(SESSION_FILE, API_ID, API_HASH) as client: + logger.success(f"New session file created: {SESSION_FILE}.session") + diff --git a/src/auto_archiver/core/orchestrator.py b/src/auto_archiver/core/orchestrator.py index dc15809..b305963 100644 --- a/src/auto_archiver/core/orchestrator.py +++ b/src/auto_archiver/core/orchestrator.py @@ -220,7 +220,7 @@ class ArchivingOrchestrator: loaded_module: BaseModule = get_module(module, self.config) except (KeyboardInterrupt, Exception) as e: logger.error(f"Error during setup of archivers: {e}\n{traceback.format_exc()}") - if module_type == 'extractor': + if module_type == 'extractor' and loaded_module.name == module: loaded_module.cleanup() exit() diff --git a/src/auto_archiver/core/storage.py b/src/auto_archiver/core/storage.py index e167024..5274204 100644 --- a/src/auto_archiver/core/storage.py +++ b/src/auto_archiver/core/storage.py @@ -30,7 +30,7 @@ class Storage(BaseModule): def uploadf(self, file: IO[bytes], key: str, **kwargs: dict) -> bool: pass def upload(self, media: Media, **kwargs) -> bool: - logger.debug(f'[{self.__class__.name}] storing file {media.filename} with key {media.key}') + logger.debug(f'[{self.__class__.__name__}] storing file {media.filename} with key {media.key}') with open(media.filename, 'rb') as f: return self.uploadf(f, media, **kwargs) diff --git a/src/auto_archiver/modules/api_db/__init__.py b/src/auto_archiver/modules/api_db/__init__.py index 2070b06..a4f39a1 100644 --- a/src/auto_archiver/modules/api_db/__init__.py +++ b/src/auto_archiver/modules/api_db/__init__.py @@ -1 +1 @@ -from api_db import AAApiDb \ No newline at end of file +from .api_db import AAApiDb \ No newline at end of file diff --git a/src/auto_archiver/modules/api_db/__manifest__.py b/src/auto_archiver/modules/api_db/__manifest__.py index d22fa59..3874496 100644 --- a/src/auto_archiver/modules/api_db/__manifest__.py +++ b/src/auto_archiver/modules/api_db/__manifest__.py @@ -4,19 +4,41 @@ "entry_point": "api_db:AAApiDb", "requires_setup": True, "dependencies": { - "python": ["requests", - "loguru"], + "python": ["requests", "loguru"], }, "configs": { - "api_endpoint": {"default": None, "help": "API endpoint where calls are made to"}, - "api_token": {"default": None, "help": "API Bearer token."}, - "public": {"default": False, "help": "whether the URL should be publicly available via the API"}, - "author_id": {"default": None, "help": "which email to assign as author"}, - "group_id": {"default": None, "help": "which group of users have access to the archive in case public=false as author"}, - "allow_rearchive": {"default": True, "help": "if False then the API database will be queried prior to any archiving operations and stop if the link has already been archived", "type": "bool",}, - "store_results": {"default": True, "help": "when set, will send the results to the API database.", "type": "bool",}, - "tags": {"default": [], "help": "what tags to add to the archived URL",} + "api_endpoint": { + "default": None, + "required": True, + "help": "API endpoint where calls are made to", }, + "api_token": {"default": None, + "help": "API Bearer token."}, + "public": { + "default": False, + "type": "bool", + "help": "whether the URL should be publicly available via the API", + }, + "author_id": {"default": None, "help": "which email to assign as author"}, + "group_id": { + "default": None, + "help": "which group of users have access to the archive in case public=false as author", + }, + "allow_rearchive": { + "default": True, + "type": "bool", + "help": "if False then the API database will be queried prior to any archiving operations and stop if the link has already been archived", + }, + "store_results": { + "default": True, + "type": "bool", + "help": "when set, will send the results to the API database.", + }, + "tags": { + "default": [], + "help": "what tags to add to the archived URL", + }, + }, "description": """ Provides integration with the Auto-Archiver API for querying and storing archival data. diff --git a/src/auto_archiver/modules/api_db/api_db.py b/src/auto_archiver/modules/api_db/api_db.py index a893aee..e1f67ce 100644 --- a/src/auto_archiver/modules/api_db/api_db.py +++ b/src/auto_archiver/modules/api_db/api_db.py @@ -1,5 +1,7 @@ from typing import Union -import requests, os + +import os +import requests from loguru import logger from auto_archiver.core import Database @@ -7,17 +9,7 @@ from auto_archiver.core import Metadata class AAApiDb(Database): - """ - Connects to auto-archiver-api instance - """ - - def __init__(self, config: dict) -> None: - # without this STEP.__init__ is not called - super().__init__(config) - self.allow_rearchive = bool(self.allow_rearchive) - self.store_results = bool(self.store_results) - self.assert_valid_string("api_endpoint") - + """Connects to auto-archiver-api instance""" def fetch(self, item: Metadata) -> Union[Metadata, bool]: """ query the database for the existence of this item. diff --git a/src/auto_archiver/modules/atlos/__init__.py b/src/auto_archiver/modules/atlos/__init__.py deleted file mode 100644 index de7fead..0000000 --- a/src/auto_archiver/modules/atlos/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from .atlos import AtlosStorage \ No newline at end of file diff --git a/src/auto_archiver/modules/atlos/__manifest__.py b/src/auto_archiver/modules/atlos/__manifest__.py deleted file mode 100644 index 7ba2f72..0000000 --- a/src/auto_archiver/modules/atlos/__manifest__.py +++ /dev/null @@ -1,40 +0,0 @@ -{ - "name": "atlos_storage", - "type": ["storage"], - "requires_setup": True, - "dependencies": {"python": ["loguru", "requests"], "bin": [""]}, - "configs": { - "path_generator": { - "default": "url", - "help": "how to store the file in terms of directory structure: 'flat' sets to root; 'url' creates a directory based on the provided URL; 'random' creates a random directory.", - }, - "filename_generator": { - "default": "random", - "help": "how to name stored files: 'random' creates a random string; 'static' uses a replicable strategy such as a hash.", - }, - "api_token": { - "default": None, - "help": "An Atlos API token. For more information, see https://docs.atlos.org/technical/api/", - "type": "str", - }, - "atlos_url": { - "default": "https://platform.atlos.org", - "help": "The URL of your Atlos instance (e.g., https://platform.atlos.org), without a trailing slash.", - "type": "str", - }, - }, - "description": """ - AtlosStorage: A storage module for saving media files to the Atlos platform. - - ### Features - - Uploads media files to Atlos using Atlos-specific APIs. - - Automatically calculates SHA-256 hashes of media files for integrity verification. - - Skips uploads for files that already exist on Atlos with the same hash. - - Supports attaching metadata, such as `atlos_id`, to the uploaded files. - - Provides CDN-like URLs for accessing uploaded media. - - ### Notes - - Requires Atlos API configuration, including `atlos_url` and `api_token`. - - Files are linked to an `atlos_id` in the metadata, ensuring proper association with Atlos source materials. - """, -} diff --git a/src/auto_archiver/modules/atlos_db/atlos_db.py b/src/auto_archiver/modules/atlos_db/atlos_db.py index c45e215..baa9fef 100644 --- a/src/auto_archiver/modules/atlos_db/atlos_db.py +++ b/src/auto_archiver/modules/atlos_db/atlos_db.py @@ -1,14 +1,10 @@ -import os - from typing import Union -from loguru import logger -from csv import DictWriter -from dataclasses import asdict + import requests +from loguru import logger from auto_archiver.core import Database from auto_archiver.core import Metadata -from auto_archiver.utils import get_atlos_config_options class AtlosDb(Database): diff --git a/src/auto_archiver/modules/atlos_feeder/__manifest__.py b/src/auto_archiver/modules/atlos_feeder/__manifest__.py index f2772f2..5ae3540 100644 --- a/src/auto_archiver/modules/atlos_feeder/__manifest__.py +++ b/src/auto_archiver/modules/atlos_feeder/__manifest__.py @@ -8,8 +8,9 @@ "configs": { "api_token": { "default": None, + "type": "str", + "required": True, "help": "An Atlos API token. For more information, see https://docs.atlos.org/technical/api/", - "type": "str" }, "atlos_url": { "default": "https://platform.atlos.org", diff --git a/src/auto_archiver/modules/atlos_feeder/atlos_feeder.py b/src/auto_archiver/modules/atlos_feeder/atlos_feeder.py index 9811a82..bbf06f6 100644 --- a/src/auto_archiver/modules/atlos_feeder/atlos_feeder.py +++ b/src/auto_archiver/modules/atlos_feeder/atlos_feeder.py @@ -1,19 +1,12 @@ -from loguru import logger import requests +from loguru import logger from auto_archiver.core import Feeder -from auto_archiver.core import Metadata, ArchivingContext -from auto_archiver.utils import get_atlos_config_options +from auto_archiver.core import Metadata class AtlosFeeder(Feeder): - def __init__(self, config: dict) -> None: - # without this STEP.__init__ is not called - super().__init__(config) - if type(self.api_token) != str: - raise Exception("Atlos Feeder did not receive an Atlos API token") - def __iter__(self) -> Metadata: # Get all the urls from the Atlos API count = 0 diff --git a/src/auto_archiver/modules/atlos/atlos.py b/src/auto_archiver/modules/atlos_storage/atlos_storage.py similarity index 96% rename from src/auto_archiver/modules/atlos/atlos.py rename to src/auto_archiver/modules/atlos_storage/atlos_storage.py index abc8a1a..f8eef68 100644 --- a/src/auto_archiver/modules/atlos/atlos.py +++ b/src/auto_archiver/modules/atlos_storage/atlos_storage.py @@ -1,12 +1,12 @@ -import os -from typing import IO, List, Optional -from loguru import logger -import requests import hashlib +import os +from typing import IO, Optional + +import requests +from loguru import logger from auto_archiver.core import Media, Metadata from auto_archiver.core import Storage -from auto_archiver.utils import get_atlos_config_options class AtlosStorage(Storage): diff --git a/src/auto_archiver/modules/gdrive_storage/__manifest__.py b/src/auto_archiver/modules/gdrive_storage/__manifest__.py index e24f21b..2ca7e27 100644 --- a/src/auto_archiver/modules/gdrive_storage/__manifest__.py +++ b/src/auto_archiver/modules/gdrive_storage/__manifest__.py @@ -1,14 +1,14 @@ { "name": "Google Drive Storage", "type": ["storage"], + "author": "Dave Mateer", + "entry_point": "gdrive_storage::GDriveStorage", "requires_setup": True, "dependencies": { "python": [ "loguru", - "google-api-python-client", - "google-auth", - "google-auth-oauthlib", - "google-auth-httplib2" + "googleapiclient", + "google", ], }, "configs": { @@ -18,17 +18,24 @@ "choices": ["flat", "url", "random"], }, "filename_generator": { - "default": "random", + "default": "static", "help": "how to name stored files: 'random' creates a random string; 'static' uses a replicable strategy such as a hash.", "choices": ["random", "static"], }, - "root_folder_id": {"default": None, "help": "root google drive folder ID to use as storage, found in URL: 'https://drive.google.com/drive/folders/FOLDER_ID'"}, - "oauth_token": {"default": None, "help": "JSON filename with Google Drive OAuth token: check auto-archiver repository scripts folder for create_update_gdrive_oauth_token.py. NOTE: storage used will count towards owner of GDrive folder, therefore it is best to use oauth_token_filename over service_account."}, + "root_folder_id": {"default": None, + # "required": True, + "help": "root google drive folder ID to use as storage, found in URL: 'https://drive.google.com/drive/folders/FOLDER_ID'"}, + "oauth_token": {"default": None, + "help": "JSON filename with Google Drive OAuth token: check auto-archiver repository scripts folder for create_update_gdrive_oauth_token.py. NOTE: storage used will count towards owner of GDrive folder, therefore it is best to use oauth_token_filename over service_account."}, "service_account": {"default": "secrets/service_account.json", "help": "service account JSON file path, same as used for Google Sheets. NOTE: storage used will count towards the developer account."}, }, "description": """ + GDriveStorage: A storage module for saving archived content to Google Drive. + Author: Dave Mateer, (And maintained by: ) + Source Documentation: https://davemateer.com/2022/04/28/google-drive-with-python + ### Features - Saves media files to Google Drive, organizing them into folders based on the provided path structure. - Supports OAuth token-based authentication or service account credentials for API access. @@ -39,5 +46,55 @@ - Requires setup with either a Google OAuth token or a service account JSON file. - Files are uploaded to the specified `root_folder_id` and organized by the `media.key` structure. - Automatically handles Google Drive API token refreshes for long-running jobs. - """ + + ## Overview +This module integrates Google Drive as a storage backend, enabling automatic folder creation and file uploads. It supports authentication via **service accounts** (recommended for automation) or **OAuth tokens** (for user-based authentication). + +## Features +- Saves files to Google Drive, organizing them into structured folders. +- Supports both **service account** and **OAuth token** authentication. +- Automatically creates folders if they don't exist. +- Generates public URLs for easy file sharing. + +## Setup Guide +1. **Enable Google Drive API** + - Create a Google Cloud project at [Google Cloud Console](https://console.cloud.google.com/) + - Enable the **Google Drive API**. + +2. **Set Up a Google Drive Folder** + - Create a folder in **Google Drive** and copy its **folder ID** from the URL. + - Add the **folder ID** to your configuration (`orchestration.yaml`): + ```yaml + root_folder_id: "FOLDER_ID" + ``` + +3. **Authentication Options** + - **Option 1: Service Account (Recommended)** + - Create a **service account** in Google Cloud IAM. + - Download the JSON key file and save it as: + ``` + secrets/service_account.json + ``` + - **Share your Drive folder** with the service account’s `client_email` (found in the JSON file). + + - **Option 2: OAuth Token (User Authentication)** + - Create OAuth **Desktop App credentials** in Google Cloud. + - Save the credentials as: + ``` + secrets/oauth_credentials.json + ``` + - Generate an OAuth token by running: + ```sh + python scripts/create_update_gdrive_oauth_token.py -c secrets/oauth_credentials.json + ``` + + + Notes on the OAuth token: + Tokens are refreshed after 1 hour however keep working for 7 days (tbc) + so as long as the job doesn't last for 7 days then this method of refreshing only once per run will work + see this link for details on the token: + https://davemateer.com/2022/04/28/google-drive-with-python#tokens + + +""" } diff --git a/src/auto_archiver/modules/gdrive_storage/gdrive_storage.py b/src/auto_archiver/modules/gdrive_storage/gdrive_storage.py index c2d326d..b764f1d 100644 --- a/src/auto_archiver/modules/gdrive_storage/gdrive_storage.py +++ b/src/auto_archiver/modules/gdrive_storage/gdrive_storage.py @@ -1,68 +1,69 @@ -import shutil, os, time, json +import json +import os +import time from typing import IO -from loguru import logger -from googleapiclient.discovery import build -from googleapiclient.http import MediaFileUpload +from google.auth.transport.requests import Request from google.oauth2 import service_account from google.oauth2.credentials import Credentials -from google.auth.transport.requests import Request +from googleapiclient.discovery import build +from googleapiclient.http import MediaFileUpload +from loguru import logger from auto_archiver.core import Media from auto_archiver.core import Storage + + class GDriveStorage(Storage): - def __init__(self, config: dict) -> None: - super().__init__(config) + def setup(self, config: dict) -> None: + # Step 1: Call the BaseModule setup to dynamically assign configs + super().setup(config) + self.scopes = ['https://www.googleapis.com/auth/drive'] + # Initialize Google Drive service + self._setup_google_drive_service() - SCOPES = ['https://www.googleapis.com/auth/drive'] - - if self.oauth_token is not None: - """ - Tokens are refreshed after 1 hour - however keep working for 7 days (tbc) - so as long as the job doesn't last for 7 days - then this method of refreshing only once per run will work - see this link for details on the token - https://davemateer.com/2022/04/28/google-drive-with-python#tokens - """ - logger.debug(f'Using GD OAuth token {self.oauth_token}') - # workaround for missing 'refresh_token' in from_authorized_user_file - with open(self.oauth_token, 'r') as stream: - creds_json = json.load(stream) - creds_json['refresh_token'] = creds_json.get("refresh_token", "") - creds = Credentials.from_authorized_user_info(creds_json, SCOPES) - # creds = Credentials.from_authorized_user_file(self.oauth_token, SCOPES) - - if not creds or not creds.valid: - if creds and creds.expired and creds.refresh_token: - logger.debug('Requesting new GD OAuth token') - creds.refresh(Request()) - else: - raise Exception("Problem with creds - create the token again") - - # Save the credentials for the next run - with open(self.oauth_token, 'w') as token: - logger.debug('Saving new GD OAuth token') - token.write(creds.to_json()) - else: - logger.debug('GD OAuth Token valid') + def _setup_google_drive_service(self): + """Initialize Google Drive service based on provided credentials.""" + if self.oauth_token: + logger.debug(f"Using Google Drive OAuth token: {self.oauth_token}") + self.service = self._initialize_with_oauth_token() + elif self.service_account: + logger.debug(f"Using Google Drive service account: {self.service_account}") + self.service = self._initialize_with_service_account() else: - gd_service_account = self.service_account - logger.debug(f'Using GD Service Account {gd_service_account}') - creds = service_account.Credentials.from_service_account_file(gd_service_account, scopes=SCOPES) + raise ValueError("Missing credentials: either `oauth_token` or `service_account` must be provided.") - self.service = build('drive', 'v3', credentials=creds) + def _initialize_with_oauth_token(self): + """Initialize Google Drive service with OAuth token.""" + with open(self.oauth_token, 'r') as stream: + creds_json = json.load(stream) + creds_json['refresh_token'] = creds_json.get("refresh_token", "") + + creds = Credentials.from_authorized_user_info(creds_json, self.scopes) + if not creds.valid and creds.expired and creds.refresh_token: + creds.refresh(Request()) + with open(self.oauth_token, 'w') as token_file: + logger.debug("Saving refreshed OAuth token.") + token_file.write(creds.to_json()) + elif not creds.valid: + raise ValueError("Invalid OAuth token. Please regenerate the token.") + + return build('drive', 'v3', credentials=creds) + + def _initialize_with_service_account(self): + """Initialize Google Drive service with service account.""" + creds = service_account.Credentials.from_service_account_file(self.service_account, scopes=self.scopes) + return build('drive', 'v3', credentials=creds) def get_cdn_url(self, media: Media) -> str: """ only support files saved in a folder for GD S3 supports folder and all stored in the root """ - # full_name = os.path.join(self.folder, media.key) parent_id, folder_id = self.root_folder_id, None path_parts = media.key.split(os.path.sep) @@ -77,7 +78,7 @@ class GDriveStorage(Storage): return f"https://drive.google.com/file/d/{file_id}/view?usp=sharing" def upload(self, media: Media, **kwargs) -> bool: - logger.debug(f'[{self.__class__.name}] storing file {media.filename} with key {media.key}') + logger.debug(f'[{self.__class__.__name__}] storing file {media.filename} with key {media.key}') """ 1. for each sub-folder in the path check if exists or create 2. upload file to root_id/other_paths.../filename @@ -168,8 +169,3 @@ class GDriveStorage(Storage): gd_folder = self.service.files().create(supportsAllDrives=True, body=file_metadata, fields='id').execute() return gd_folder.get('id') - # def exists(self, key): - # try: - # self.get_cdn_url(key) - # return True - # except: return False diff --git a/src/auto_archiver/modules/gsheet_db/__manifest__.py b/src/auto_archiver/modules/gsheet_db/__manifest__.py index f926adc..cf95245 100644 --- a/src/auto_archiver/modules/gsheet_db/__manifest__.py +++ b/src/auto_archiver/modules/gsheet_db/__manifest__.py @@ -4,7 +4,7 @@ "entry_point": "gsheet_db::GsheetsDb", "requires_setup": True, "dependencies": { - "python": ["loguru", "gspread", "python-slugify"], + "python": ["loguru", "gspread", "slugify"], }, "configs": { "allow_worksheets": { @@ -17,6 +17,7 @@ }, "use_sheet_names_in_stored_paths": { "default": True, + "type": "bool", "help": "if True the stored files path will include 'workbook_name/worksheet_name/...'", } }, diff --git a/src/auto_archiver/modules/gsheet_feeder/__manifest__.py b/src/auto_archiver/modules/gsheet_feeder/__manifest__.py index 1c9acab..7b74072 100644 --- a/src/auto_archiver/modules/gsheet_feeder/__manifest__.py +++ b/src/auto_archiver/modules/gsheet_feeder/__manifest__.py @@ -4,7 +4,7 @@ "entry_point": "gsheet_feeder::GsheetsFeeder", "requires_setup": True, "dependencies": { - "python": ["loguru", "gspread", "python-slugify"], + "python": ["loguru", "gspread", "slugify"], }, "configs": { "sheet": {"default": None, "help": "name of the sheet to archive"}, diff --git a/src/auto_archiver/modules/instagram_api_extractor/__manifest__.py b/src/auto_archiver/modules/instagram_api_extractor/__manifest__.py index 57f378e..a958a99 100644 --- a/src/auto_archiver/modules/instagram_api_extractor/__manifest__.py +++ b/src/auto_archiver/modules/instagram_api_extractor/__manifest__.py @@ -1,6 +1,7 @@ { "name": "Instagram API Extractor", "type": ["extractor"], + "entry_point": "instagram_api_extractor::InstagramAPIExtractor", "dependencies": {"python": ["requests", "loguru", @@ -9,24 +10,32 @@ }, "requires_setup": True, "configs": { - "access_token": {"default": None, "help": "a valid instagrapi-api token"}, - "api_endpoint": {"default": None, "help": "API endpoint to use"}, + "access_token": {"default": None, + "help": "a valid instagrapi-api token"}, + "api_endpoint": {"default": None, + # "required": True, + "help": "API endpoint to use"}, "full_profile": { "default": False, + "type": "bool", "help": "if true, will download all posts, tagged posts, stories, and highlights for a profile, if false, will only download the profile pic and information.", }, "full_profile_max_posts": { "default": 0, + "type": "int", "help": "Use to limit the number of posts to download when full_profile is true. 0 means no limit. limit is applied softly since posts are fetched in batch, once to: posts, tagged posts, and highlights", }, "minimize_json_output": { "default": True, + "type": "bool", "help": "if true, will remove empty values from the json output", }, }, "description": """ Archives various types of Instagram content using the Instagrapi API. +Requires setting up an Instagrapi API deployment and providing an access token and API endpoint. + ### Features - Connects to an Instagrapi API deployment to fetch Instagram profiles, posts, stories, highlights, reels, and tagged content. - Supports advanced configuration options, including: diff --git a/src/auto_archiver/modules/instagram_api_extractor/instagram_api_extractor.py b/src/auto_archiver/modules/instagram_api_extractor/instagram_api_extractor.py index 3d7f9e5..4a18228 100644 --- a/src/auto_archiver/modules/instagram_api_extractor/instagram_api_extractor.py +++ b/src/auto_archiver/modules/instagram_api_extractor/instagram_api_extractor.py @@ -32,16 +32,11 @@ class InstagramAPIExtractor(Extractor): r"(?:(?:http|https):\/\/)?(?:www.)?(?:instagram.com)\/(stories(?:\/highlights)?|p|reel)?\/?([^\/\?]*)\/?(\d+)?" ) - def __init__(self, config: dict) -> None: - super().__init__(config) - self.assert_valid_string("access_token") - self.assert_valid_string("api_endpoint") - self.full_profile_max_posts = int(self.full_profile_max_posts) + def setup(self, config: dict) -> None: + super().setup(config) if self.api_endpoint[-1] == "/": self.api_endpoint = self.api_endpoint[:-1] - self.full_profile = bool(self.full_profile) - self.minimize_json_output = bool(self.minimize_json_output) def download(self, item: Metadata) -> Metadata: url = item.get_url() diff --git a/src/auto_archiver/modules/instagram_extractor/__manifest__.py b/src/auto_archiver/modules/instagram_extractor/__manifest__.py index 6e7518e..d8e4a9b 100644 --- a/src/auto_archiver/modules/instagram_extractor/__manifest__.py +++ b/src/auto_archiver/modules/instagram_extractor/__manifest__.py @@ -9,9 +9,12 @@ }, "requires_setup": True, "configs": { - "username": {"default": None, "help": "a valid Instagram username"}, + "username": {"default": None, + "required": True, + "help": "a valid Instagram username"}, "password": { "default": None, + "required": True, "help": "the corresponding Instagram account password", }, "download_folder": { @@ -25,9 +28,11 @@ # TODO: fine-grain # "download_stories": {"default": True, "help": "if the link is to a user profile: whether to get stories information"}, }, - "description": """Uses the Instaloader library to download content from Instagram. This class handles both individual posts - and user profiles, downloading as much information as possible, including images, videos, text, stories, - highlights, and tagged posts. Authentication is required via username/password or a session file. + "description": """ + Uses the [Instaloader library](https://instaloader.github.io/as-module.html) to download content from Instagram. This class handles both individual posts + and user profiles, downloading as much information as possible, including images, videos, text, stories, + highlights, and tagged posts. + Authentication is required via username/password or a session file. """, } diff --git a/src/auto_archiver/modules/instagram_extractor/instagram_extractor.py b/src/auto_archiver/modules/instagram_extractor/instagram_extractor.py index 1a246fb..1cdb0b1 100644 --- a/src/auto_archiver/modules/instagram_extractor/instagram_extractor.py +++ b/src/auto_archiver/modules/instagram_extractor/instagram_extractor.py @@ -4,7 +4,7 @@ """ import re, os, shutil, traceback -import instaloader # https://instaloader.github.io/as-module.html +import instaloader from loguru import logger from auto_archiver.core import Extractor @@ -22,13 +22,9 @@ class InstagramExtractor(Extractor): profile_pattern = re.compile(r"(?:(?:http|https):\/\/)?(?:www.)?(?:instagram.com|instagr.am|instagr.com)\/(\w+)") # TODO: links to stories - def __init__(self, config: dict) -> None: - super().__init__(config) - # TODO: refactor how configuration validation is done - self.assert_valid_string("username") - self.assert_valid_string("password") - self.assert_valid_string("download_folder") - self.assert_valid_string("session_file") + def setup(self, config: dict) -> None: + super().setup(config) + self.insta = instaloader.Instaloader( download_geotags=True, download_comments=True, compress_json=False, dirname_pattern=self.download_folder, filename_pattern="{date_utc}_UTC_{target}__{typename}" ) diff --git a/src/auto_archiver/modules/instagram_tbot_extractor/__manifest__.py b/src/auto_archiver/modules/instagram_tbot_extractor/__manifest__.py index 8a1f74f..a24a864 100644 --- a/src/auto_archiver/modules/instagram_tbot_extractor/__manifest__.py +++ b/src/auto_archiver/modules/instagram_tbot_extractor/__manifest__.py @@ -1,15 +1,16 @@ { "name": "Instagram Telegram Bot Extractor", "type": ["extractor"], - "dependencies": {"python": ["loguru", - "telethon",], + "dependencies": {"python": ["loguru", "telethon",], }, "requires_setup": True, "configs": { "api_id": {"default": None, "help": "telegram API_ID value, go to https://my.telegram.org/apps"}, "api_hash": {"default": None, "help": "telegram API_HASH value, go to https://my.telegram.org/apps"}, "session_file": {"default": "secrets/anon-insta", "help": "optional, records the telegram login session for future usage, '.session' will be appended to the provided value."}, - "timeout": {"default": 45, "help": "timeout to fetch the instagram content in seconds."}, + "timeout": {"default": 45, + "type": "int", + "help": "timeout to fetch the instagram content in seconds."}, }, "description": """ The `InstagramTbotExtractor` module uses a Telegram bot (`instagram_load_bot`) to fetch and archive Instagram content, @@ -28,6 +29,12 @@ returned as part of a `Metadata` object. To use the `InstagramTbotExtractor`, you need to provide the following configuration settings: - **API ID and Hash**: Telegram API credentials obtained from [my.telegram.org/apps](https://my.telegram.org/apps). - **Session File**: Optional path to store the Telegram session file for future use. - +- The session file is created automatically and should be unique for each instance. +- You may need to enter your Telegram credentials (phone) and use the a 2FA code sent to you the first time you run the extractor.: +```2025-01-30 00:43:49.348 | INFO | auto_archiver.modules.instagram_tbot_extractor.instagram_tbot_extractor:setup:36 - SETUP instagram_tbot_extractor checking login... +Please enter your phone (or bot token): +447123456789 +Please enter the code you received: 00000 +Signed in successfully as E C; remember to not break the ToS or you will risk an account ban! +``` """, } diff --git a/src/auto_archiver/modules/instagram_tbot_extractor/instagram_tbot_extractor.py b/src/auto_archiver/modules/instagram_tbot_extractor/instagram_tbot_extractor.py index 60fa397..791b9c0 100644 --- a/src/auto_archiver/modules/instagram_tbot_extractor/instagram_tbot_extractor.py +++ b/src/auto_archiver/modules/instagram_tbot_extractor/instagram_tbot_extractor.py @@ -27,15 +27,19 @@ class InstagramTbotExtractor(Extractor): https://t.me/instagram_load_bot """ - def setup(self) -> None: + def setup(self, configs) -> None: """ 1. makes a copy of session_file that is removed in cleanup 2. checks if the session file is valid """ + super().setup(configs) logger.info(f"SETUP {self.name} checking login...") # make a copy of the session that is used exclusively with this archiver instance new_session_file = os.path.join("secrets/", f"instabot-{time.strftime('%Y-%m-%d')}{random_str(8)}.session") + if not os.path.exists(f"{self.session_file}.session"): + raise FileNotFoundError(f"session file {self.session_file}.session not found, " + f"to set this up run the setup script in scripts/telegram_setup.py") shutil.copy(self.session_file + ".session", new_session_file) self.session_file = new_session_file.replace(".session", "") @@ -43,7 +47,6 @@ class InstagramTbotExtractor(Extractor): self.client = TelegramClient(self.session_file, self.api_id, self.api_hash) except OperationalError as e: logger.error(f"Unable to access the {self.session_file} session, please make sure you don't use the same session file here and in telethon_extractor. if you do then disable at least one of the archivers for the 1st time you setup telethon session: {e}") - with self.client.start(): logger.success(f"SETUP {self.name} login works.") diff --git a/src/auto_archiver/modules/pdq_hash_enricher/__manifest__.py b/src/auto_archiver/modules/pdq_hash_enricher/__manifest__.py index 6353d12..133fef7 100644 --- a/src/auto_archiver/modules/pdq_hash_enricher/__manifest__.py +++ b/src/auto_archiver/modules/pdq_hash_enricher/__manifest__.py @@ -3,7 +3,7 @@ "type": ["enricher"], "requires_setup": False, "dependencies": { - "python": ["loguru", "pdqhash", "numpy", "Pillow"], + "python": ["loguru", "pdqhash", "numpy", "PIL"], }, "description": """ PDQ Hash Enricher for generating perceptual hashes of media files. diff --git a/src/auto_archiver/modules/s3_storage/__init__.py b/src/auto_archiver/modules/s3_storage/__init__.py index 1c826fd..cbf3237 100644 --- a/src/auto_archiver/modules/s3_storage/__init__.py +++ b/src/auto_archiver/modules/s3_storage/__init__.py @@ -1 +1 @@ -from .s3 import S3Storage \ No newline at end of file +from .s3_storage import S3Storage \ No newline at end of file diff --git a/src/auto_archiver/modules/s3_storage/__manifest__.py b/src/auto_archiver/modules/s3_storage/__manifest__.py index 16ac7bd..df05055 100644 --- a/src/auto_archiver/modules/s3_storage/__manifest__.py +++ b/src/auto_archiver/modules/s3_storage/__manifest__.py @@ -7,12 +7,12 @@ }, "configs": { "path_generator": { - "default": "url", + "default": "flat", "help": "how to store the file in terms of directory structure: 'flat' sets to root; 'url' creates a directory based on the provided URL; 'random' creates a random directory.", "choices": ["flat", "url", "random"], }, "filename_generator": { - "default": "random", + "default": "static", "help": "how to name stored files: 'random' creates a random string; 'static' uses a replicable strategy such as a hash.", "choices": ["random", "static"], }, @@ -20,7 +20,9 @@ "region": {"default": None, "help": "S3 region name"}, "key": {"default": None, "help": "S3 API key"}, "secret": {"default": None, "help": "S3 API secret"}, - "random_no_duplicate": {"default": False, "help": "if set, it will override `path_generator`, `filename_generator` and `folder`. It will check if the file already exists and if so it will not upload it again. Creates a new root folder path `no-dups/`"}, + "random_no_duplicate": {"default": False, + "type": "bool", + "help": "if set, it will override `path_generator`, `filename_generator` and `folder`. It will check if the file already exists and if so it will not upload it again. Creates a new root folder path `no-dups/`"}, "endpoint_url": { "default": 'https://{region}.digitaloceanspaces.com', "help": "S3 bucket endpoint, {region} are inserted at runtime" @@ -29,7 +31,9 @@ "default": 'https://{bucket}.{region}.cdn.digitaloceanspaces.com/{key}', "help": "S3 CDN url, {bucket}, {region} and {key} are inserted at runtime" }, - "private": {"default": False, "help": "if true S3 files will not be readable online"}, + "private": {"default": False, + "type": "bool", + "help": "if true S3 files will not be readable online"}, }, "description": """ S3Storage: A storage module for saving media files to an S3-compatible object storage. diff --git a/src/auto_archiver/modules/s3_storage/s3.py b/src/auto_archiver/modules/s3_storage/s3_storage.py similarity index 88% rename from src/auto_archiver/modules/s3_storage/s3.py rename to src/auto_archiver/modules/s3_storage/s3_storage.py index 10d5f61..f324d5c 100644 --- a/src/auto_archiver/modules/s3_storage/s3.py +++ b/src/auto_archiver/modules/s3_storage/s3_storage.py @@ -1,19 +1,21 @@ from typing import IO -import boto3, os -from auto_archiver.utils.misc import random_str -from auto_archiver.core import Media -from auto_archiver.core import Storage - -from auto_archiver.modules.hash_enricher import HashEnricher +import boto3 +import os from loguru import logger -NO_DUPLICATES_FOLDER = "no-dups/" -class S3Storage(Storage): +from auto_archiver.core import Media +from auto_archiver.core import Storage +from auto_archiver.modules.hash_enricher import HashEnricher +from auto_archiver.utils.misc import random_str - def __init__(self, config: dict) -> None: - super().__init__(config) +NO_DUPLICATES_FOLDER = "no-dups/" + +class S3Storage(Storage, HashEnricher): + + def setup(self, config: dict) -> None: + super().setup(config) self.s3 = boto3.client( 's3', region_name=self.region, @@ -21,7 +23,6 @@ class S3Storage(Storage): aws_access_key_id=self.key, aws_secret_access_key=self.secret ) - self.random_no_duplicate = bool(self.random_no_duplicate) if self.random_no_duplicate: logger.warning("random_no_duplicate is set to True, this will override `path_generator`, `filename_generator` and `folder`.") @@ -48,8 +49,7 @@ class S3Storage(Storage): def is_upload_needed(self, media: Media) -> bool: if self.random_no_duplicate: # checks if a folder with the hash already exists, if so it skips the upload - he = HashEnricher({"hash_enricher": {"algorithm": "SHA-256", "chunksize": 1.6e7}}) - hd = he.calculate_hash(media.filename) + hd = self.calculate_hash(media.filename) path = os.path.join(NO_DUPLICATES_FOLDER, hd[:24]) if existing_key:=self.file_in_folder(path): @@ -61,8 +61,7 @@ class S3Storage(Storage): _, ext = os.path.splitext(media.key) media.key = os.path.join(path, f"{random_str(24)}{ext}") return True - - + def file_in_folder(self, path:str) -> str: # checks if path exists and is not an empty folder if not path.endswith('/'): diff --git a/src/auto_archiver/modules/ssl_enricher/__manifest__.py b/src/auto_archiver/modules/ssl_enricher/__manifest__.py index 0fb7cd9..9028f14 100644 --- a/src/auto_archiver/modules/ssl_enricher/__manifest__.py +++ b/src/auto_archiver/modules/ssl_enricher/__manifest__.py @@ -3,7 +3,7 @@ "type": ["enricher"], "requires_setup": False, "dependencies": { - "python": ["loguru", "python-slugify"], + "python": ["loguru", "slugify"], }, 'entry_point': 'ssl_enricher::SSLEnricher', "configs": { diff --git a/src/auto_archiver/modules/thumbnail_enricher/__manifest__.py b/src/auto_archiver/modules/thumbnail_enricher/__manifest__.py index bd7836d..e47397f 100644 --- a/src/auto_archiver/modules/thumbnail_enricher/__manifest__.py +++ b/src/auto_archiver/modules/thumbnail_enricher/__manifest__.py @@ -3,7 +3,7 @@ "type": ["enricher"], "requires_setup": False, "dependencies": { - "python": ["loguru", "ffmpeg-python"], + "python": ["loguru", "ffmpeg"], "bin": ["ffmpeg"] }, "configs": { diff --git a/src/auto_archiver/modules/vk_extractor/__manifest__.py b/src/auto_archiver/modules/vk_extractor/__manifest__.py index 116b430..033fe50 100644 --- a/src/auto_archiver/modules/vk_extractor/__manifest__.py +++ b/src/auto_archiver/modules/vk_extractor/__manifest__.py @@ -4,14 +4,20 @@ "requires_setup": True, "depends": ["core", "utils"], "dependencies": { - "python": ["loguru", - "vk_url_scraper"], + "python": ["loguru", "vk_url_scraper"], }, "configs": { - "username": {"default": None, "help": "valid VKontakte username"}, - "password": {"default": None, "help": "valid VKontakte password"}, - "session_file": {"default": "secrets/vk_config.v2.json", "help": "valid VKontakte password"}, + "username": {"default": None, + "required": True, + "help": "valid VKontakte username"}, + "password": {"default": None, + "required": True, + "help": "valid VKontakte password"}, + "session_file": { + "default": "secrets/vk_config.v2.json", + "help": "valid VKontakte password", }, + }, "description": """ The `VkExtractor` fetches posts, text, and images from VK (VKontakte) social media pages. This archiver is specialized for `/wall` posts and uses the `VkScraper` library to extract @@ -31,6 +37,5 @@ To use the `VkArchiver`, you must provide valid VKontakte login credentials and Credentials can be set in the configuration file or directly via environment variables. Ensure you have access to the VKontakte API by creating an account at [VKontakte](https://vk.com/). -""" -, +""", } diff --git a/src/auto_archiver/modules/vk_extractor/vk_extractor.py b/src/auto_archiver/modules/vk_extractor/vk_extractor.py index 1bce167..301fa89 100644 --- a/src/auto_archiver/modules/vk_extractor/vk_extractor.py +++ b/src/auto_archiver/modules/vk_extractor/vk_extractor.py @@ -12,10 +12,8 @@ class VkExtractor(Extractor): Currently only works for /wall posts """ - def __init__(self, config: dict) -> None: - super().__init__(config) - self.assert_valid_string("username") - self.assert_valid_string("password") + def setup(self, config: dict) -> None: + super().setup(config) self.vks = VkScraper(self.username, self.password, session_file=self.session_file) def download(self, item: Metadata) -> Metadata: diff --git a/src/auto_archiver/modules/wacz_enricher/__manifest__.py b/src/auto_archiver/modules/wacz_enricher/__manifest__.py index bb9d290..46ce05e 100644 --- a/src/auto_archiver/modules/wacz_enricher/__manifest__.py +++ b/src/auto_archiver/modules/wacz_enricher/__manifest__.py @@ -1,6 +1,7 @@ { "name": "WACZ Enricher", "type": ["enricher", "archiver"], + "entry_point": "wacz_enricher::WaczExtractorEnricher", "requires_setup": True, "dependencies": { "python": [ @@ -25,6 +26,7 @@ }, "description": """ Creates .WACZ archives of web pages using the `browsertrix-crawler` tool, with options for media extraction and screenshot saving. + [Browsertrix-crawler](https://crawler.docs.browsertrix.com/user-guide/) is a headless browser-based crawler that archives web pages in WACZ format. ### Features - Archives web pages into .WACZ format using Docker or direct invocation of `browsertrix-crawler`. @@ -33,7 +35,7 @@ - Generates metadata from the archived page's content and structure (e.g., titles, text). ### Notes - - Requires Docker for running `browsertrix-crawler` unless explicitly disabled. + - Requires Docker for running `browsertrix-crawler` . - Configurable via parameters for timeout, media extraction, screenshots, and proxy settings. """ } diff --git a/src/auto_archiver/modules/wacz_enricher/wacz_enricher.py b/src/auto_archiver/modules/wacz_enricher/wacz_enricher.py index 1eb7398..8810b84 100644 --- a/src/auto_archiver/modules/wacz_enricher/wacz_enricher.py +++ b/src/auto_archiver/modules/wacz_enricher/wacz_enricher.py @@ -18,7 +18,9 @@ class WaczExtractorEnricher(Enricher, Extractor): When used as an archiver it will extract the media from the .WACZ archive so it can be enriched. """ - def setup(self) -> None: + def setup(self, configs) -> None: + super().setup(configs) + self.use_docker = os.environ.get('WACZ_ENABLE_DOCKER') or not os.environ.get('RUNNING_IN_DOCKER') self.docker_in_docker = os.environ.get('WACZ_ENABLE_DOCKER') and os.environ.get('RUNNING_IN_DOCKER') diff --git a/src/auto_archiver/modules/whisper_enricher/__manifest__.py b/src/auto_archiver/modules/whisper_enricher/__manifest__.py index 0adf9ff..f7ad1b3 100644 --- a/src/auto_archiver/modules/whisper_enricher/__manifest__.py +++ b/src/auto_archiver/modules/whisper_enricher/__manifest__.py @@ -3,7 +3,7 @@ "type": ["enricher"], "requires_setup": True, "dependencies": { - "python": ["loguru", "requests"], + "python": ["s3_storage", "loguru", "requests"], }, "configs": { "api_endpoint": {"default": None, "help": "WhisperApi api endpoint, eg: https://whisperbox-api.com/api/v1, a deployment of https://github.com/bellingcat/whisperbox-transcribe."}, diff --git a/src/auto_archiver/modules/whisper_enricher/whisper_enricher.py b/src/auto_archiver/modules/whisper_enricher/whisper_enricher.py index 09eb3db..b8fe634 100644 --- a/src/auto_archiver/modules/whisper_enricher/whisper_enricher.py +++ b/src/auto_archiver/modules/whisper_enricher/whisper_enricher.py @@ -5,7 +5,7 @@ from loguru import logger from auto_archiver.core import Enricher from auto_archiver.core import Metadata, Media, ArchivingContext from auto_archiver.modules.s3_storage import S3Storage - +from auto_archiver.core.module import get_module class WhisperEnricher(Enricher): """ @@ -53,7 +53,7 @@ class WhisperEnricher(Enricher): to_enrich.set_content(f"\n[automatic video transcript]: {v}") def submit_job(self, media: Media): - s3 = self._get_s3_storage() + s3 = get_module("s3_storage", self.config) s3_url = s3.get_cdn_url(media) assert s3_url in media.urls, f"Could not find S3 url ({s3_url}) in list of stored media urls " payload = { From b7d9145f6c1b9c50af572156f76df764f4373182 Mon Sep 17 00:00:00 2001 From: Patrick Robertson Date: Thu, 30 Jan 2025 13:21:10 +0100 Subject: [PATCH 05/62] Further tidyups + refactoring for new structure * Add implementation tests for orchestrator + logging tests * Standardise method/class vars for extractors to see if they are suitable * Fix bugs with removing default loguru logger (allows further customisation) * Fix bug loading required fields from file * --- poetry.lock | 24 +++- pyproject.toml | 1 + src/auto_archiver/__main__.py | 3 +- src/auto_archiver/core/authentication.py | 0 src/auto_archiver/core/config.py | 4 + src/auto_archiver/core/extractor.py | 14 +- src/auto_archiver/core/module.py | 5 + src/auto_archiver/core/orchestrator.py | 60 +++++---- .../instagram_api_extractor.py | 4 +- .../instagram_extractor.py | 7 +- .../telethon_extractor/telethon_extractor.py | 4 +- .../twitter_api_extractor.py | 4 +- src/auto_archiver/utils/url.py | 12 +- tests/data/example_module/example_module.py | 4 - .../example_module/__init__.py | 0 .../example_module/__manifest__.py | 5 +- .../example_module/example_module.py | 28 ++++ tests/data/test_orchestration.yaml | 16 +++ tests/extractors/test_extractor_base.py | 2 +- tests/extractors/test_instagram_extractor.py | 21 +++ tests/test_modules.py | 2 +- tests/test_orchestrator.py | 123 ++++++++++++++++++ 22 files changed, 292 insertions(+), 51 deletions(-) create mode 100644 src/auto_archiver/core/authentication.py delete mode 100644 tests/data/example_module/example_module.py rename tests/data/{ => test_modules}/example_module/__init__.py (100%) rename tests/data/{ => test_modules}/example_module/__manifest__.py (55%) create mode 100644 tests/data/test_modules/example_module/example_module.py create mode 100644 tests/data/test_orchestration.yaml create mode 100644 tests/extractors/test_instagram_extractor.py create mode 100644 tests/test_orchestrator.py diff --git a/poetry.lock b/poetry.lock index e8a899a..088fc70 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1025,7 +1025,7 @@ version = "0.7.3" description = "Python logging made (stupidly) simple" optional = false python-versions = "<4.0,>=3.5" -groups = ["main"] +groups = ["main", "dev"] files = [ {file = "loguru-0.7.3-py3-none-any.whl", hash = "sha256:31a33c10c8e1e10422bfd431aeb5d351c7cf7fa671e3c4df004162264b28220c"}, {file = "loguru-0.7.3.tar.gz", hash = "sha256:19480589e77d47b8d85b2c827ad95d49bf31b0dcde16593892eb51dd18706eb6"}, @@ -1750,6 +1750,24 @@ tomli = {version = ">=1", markers = "python_version < \"3.11\""} [package.extras] dev = ["argcomplete", "attrs (>=19.2)", "hypothesis (>=3.56)", "mock", "pygments (>=2.7.2)", "requests", "setuptools", "xmlschema"] +[[package]] +name = "pytest-loguru" +version = "0.4.0" +description = "Pytest Loguru" +optional = false +python-versions = ">=3.8" +groups = ["dev"] +files = [ + {file = "pytest_loguru-0.4.0-py3-none-any.whl", hash = "sha256:3cc7b9c6b22cb158209ccbabf0d678dacd3f3c7497d6f46f1c338c13bee1ac77"}, + {file = "pytest_loguru-0.4.0.tar.gz", hash = "sha256:0d9e4e72ae9bfd92f774c666e7353766af11b0b78edd59c290e89be116050f03"}, +] + +[package.dependencies] +loguru = "*" + +[package.extras] +test = ["pytest", "pytest-cov"] + [[package]] name = "python-dateutil" version = "2.9.0.post0" @@ -3032,7 +3050,7 @@ version = "1.2.0" description = "A small Python utility to set file creation time on Windows" optional = false python-versions = ">=3.5" -groups = ["main"] +groups = ["main", "dev"] markers = "sys_platform == \"win32\"" files = [ {file = "win32_setctime-1.2.0-py3-none-any.whl", hash = "sha256:95d644c4e708aba81dc3704a116d8cbc974d70b3bdb8be1d150e36be6e9d1390"}, @@ -3082,4 +3100,4 @@ test = ["pytest (>=8.1,<9.0)", "pytest-rerunfailures (>=14.0,<15.0)"] [metadata] lock-version = "2.1" python-versions = ">=3.10,<3.13" -content-hash = "1556d53c5a94392c120ebaafc495d3b322daf64dac4a19f9726588c7f3d84bca" +content-hash = "5a54c84ba388db7b77d1c28973b710fc99aa3822a2860b30acaf5b02ba1927bd" diff --git a/pyproject.toml b/pyproject.toml index b3a2456..3cd47e7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -63,6 +63,7 @@ dependencies = [ [tool.poetry.group.dev.dependencies] pytest = "^8.3.4" autopep8 = "^2.3.1" +pytest-loguru = "^0.4.0" [tool.poetry.group.docs.dependencies] sphinx = "^8.1.3" diff --git a/src/auto_archiver/__main__.py b/src/auto_archiver/__main__.py index d31ec5c..0e2f54f 100644 --- a/src/auto_archiver/__main__.py +++ b/src/auto_archiver/__main__.py @@ -1,8 +1,9 @@ """ Entry point for the auto_archiver package. """ from auto_archiver.core.orchestrator import ArchivingOrchestrator +import sys def main(): - ArchivingOrchestrator().run() + ArchivingOrchestrator().run(sys.argv) if __name__ == "__main__": main() diff --git a/src/auto_archiver/core/authentication.py b/src/auto_archiver/core/authentication.py new file mode 100644 index 0000000..e69de29 diff --git a/src/auto_archiver/core/config.py b/src/auto_archiver/core/config.py index 46dbe28..ca8ed25 100644 --- a/src/auto_archiver/core/config.py +++ b/src/auto_archiver/core/config.py @@ -48,6 +48,10 @@ class DefaultValidatingParser(argparse.ArgumentParser): """ for action in self._actions: if not namespace or action.dest not in namespace: + # for actions that are required and already have a default value, remove the 'required' check + if action.required and action.default is not None: + action.required = False + if action.default is not None: try: self._check_value(action, action.default) diff --git a/src/auto_archiver/core/extractor.py b/src/auto_archiver/core/extractor.py index 8d509ec..51d784f 100644 --- a/src/auto_archiver/core/extractor.py +++ b/src/auto_archiver/core/extractor.py @@ -11,9 +11,12 @@ from abc import abstractmethod from dataclasses import dataclass import mimetypes import os -import mimetypes, requests +import mimetypes + +import requests from loguru import logger from retrying import retry +import re from ..core import Metadata, ArchivingContext, BaseModule @@ -25,6 +28,8 @@ class Extractor(BaseModule): Subclasses must implement the `download` method to define platform-specific behavior. """ + valid_url: re.Pattern = None + def cleanup(self) -> None: # called when extractors are done, or upon errors, cleanup any resources pass @@ -32,13 +37,20 @@ class Extractor(BaseModule): def sanitize_url(self, url: str) -> str: # used to clean unnecessary URL parameters OR unfurl redirect links return url + + def match_link(self, url: str) -> re.Match: + return self.valid_url.match(url) def suitable(self, url: str) -> bool: """ Returns True if this extractor can handle the given URL Should be overridden by subclasses + """ + if self.valid_url: + return self.match_link(url) is not None + return True def _guess_file_type(self, path: str) -> str: diff --git a/src/auto_archiver/core/module.py b/src/auto_archiver/core/module.py index cb380cf..4542b88 100644 --- a/src/auto_archiver/core/module.py +++ b/src/auto_archiver/core/module.py @@ -83,6 +83,11 @@ def setup_paths(paths: list[str]) -> None: """ for path in paths: + # check path exists, if it doesn't, log a warning + if not os.path.exists(path): + logger.warning(f"Path '{path}' does not exist. Skipping...") + continue + # see odoo/module/module.py -> initialize_sys_path if path not in auto_archiver.modules.__path__: auto_archiver.modules.__path__.append(path) diff --git a/src/auto_archiver/core/orchestrator.py b/src/auto_archiver/core/orchestrator.py index b305963..ba46492 100644 --- a/src/auto_archiver/core/orchestrator.py +++ b/src/auto_archiver/core/orchestrator.py @@ -43,6 +43,7 @@ class ArchivingOrchestrator: def setup_basic_parser(self): parser = argparse.ArgumentParser( + prog="auto-archiver", add_help=False, description=""" Auto Archiver is a CLI tool to archive media/metadata from online URLs; @@ -51,15 +52,16 @@ class ArchivingOrchestrator: epilog="Check the code at https://github.com/bellingcat/auto-archiver", formatter_class=RichHelpFormatter, ) - parser.add_argument('--config', action='store', dest="config_file", help='the filename of the YAML configuration file (defaults to \'config.yaml\')', default=DEFAULT_CONFIG_FILE) + parser.add_argument('--help', '-h', action='store_true', dest='help', help='show this help message and exit') parser.add_argument('--version', action='version', version=__version__) + parser.add_argument('--config', action='store', dest="config_file", help='the filename of the YAML configuration file (defaults to \'config.yaml\')', default=DEFAULT_CONFIG_FILE) parser.add_argument('--mode', action='store', dest='mode', type=str, choices=['simple', 'full'], help='the mode to run the archiver in', default='simple') # override the default 'help' so we can inject all the configs and show those - parser.add_argument('-h', '--help', action='store_true', dest='help', help='show this help message and exit') parser.add_argument('-s', '--store', dest='store', default=False, help='Store the created config in the config file', action=argparse.BooleanOptionalAction) parser.add_argument('--module_paths', dest='module_paths', nargs='+', default=[], help='additional paths to search for modules', action=UniqueAppendAction) self.basic_parser = parser + return parser def setup_complete_parser(self, basic_config: dict, yaml_config: dict, unused_args: list[str]) -> None: parser = DefaultValidatingParser( @@ -78,15 +80,15 @@ class ArchivingOrchestrator: # only load the modules enabled in config # TODO: if some steps are empty (e.g. 'feeders' is empty), should we default to the 'simple' ones? Or only if they are ALL empty? enabled_modules = [] - for module_type in BaseModule.MODULE_TYPES: - enabled_modules.extend(yaml_config['steps'].get(f"{module_type}s", [])) - # add in any extra modules that have been passed on the command line for 'feeders', 'enrichers', 'archivers', 'databases', 'storages', 'formatter' - for module_type in BaseModule.MODULE_TYPES: - if modules := getattr(basic_config, f"{module_type}s", []): - enabled_modules.extend(modules) + # first loads the modules from the config file, then from the command line + for config in [yaml_config['steps'], basic_config.__dict__]: + for module_type in BaseModule.MODULE_TYPES: + enabled_modules.extend(config.get(f"{module_type}s", [])) - avail_modules = available_modules(with_manifest=True, limit_to_modules=list(dict.fromkeys(enabled_modules)), suppress_warnings=True) + # clear out duplicates, but keep the order + enabled_modules = list(dict.fromkeys(enabled_modules)) + avail_modules = available_modules(with_manifest=True, limit_to_modules=enabled_modules, suppress_warnings=True) self.add_module_args(avail_modules, parser) elif basic_config.mode == 'simple': simple_modules = [module for module in available_modules(with_manifest=True) if not module.requires_setup] @@ -163,6 +165,10 @@ class ArchivingOrchestrator: # make a nicer metavar, metavar is what's used in the help, e.g. --cli_feeder.urls [METAVAR] kwargs['metavar'] = name.upper() + if kwargs.get('required', False): + # required args shouldn't have a 'default' value, remove it + kwargs.pop('default', None) + kwargs.pop('cli_set', None) should_store = kwargs.pop('should_store', False) kwargs['dest'] = f"{module.name}.{kwargs.pop('dest', name)}" @@ -179,13 +185,12 @@ class ArchivingOrchestrator: self.add_additional_args(self.basic_parser) self.add_module_args(parser=self.basic_parser) - self.basic_parser.print_help() - exit() + self.basic_parser.exit() def setup_logging(self): # setup loguru logging - logger.remove() # remove the default logger + logger.remove(0) # remove the default logger logging_config = self.config['logging'] logger.add(sys.stderr, level=logging_config['level']) if log_file := logging_config['file']: @@ -194,14 +199,18 @@ class ArchivingOrchestrator: def install_modules(self): """ - Swaps out the previous 'strings' in the config with the actual modules + Swaps out the previous 'strings' in the config with the actual modules and loads them """ invalid_modules = [] for module_type in BaseModule.MODULE_TYPES: + step_items = [] modules_to_load = self.config['steps'][f"{module_type}s"] + assert modules_to_load, f"No {module_type}s were configured. Make sure to set at least one {module_type} \ + in your configuration file or on the command line (using --{module_type}s)" + def check_steps_ok(): if not len(step_items): logger.error(f"NO {module_type.upper()}S LOADED. Please check your configuration and try again.") @@ -239,30 +248,29 @@ class ArchivingOrchestrator: assert len(step_items) > 0, f"No {module_type}s were loaded. Please check your configuration file and try again." self.config['steps'][f"{module_type}s"] = step_items + + def load_config(self, config_file: str) -> dict: + if not os.path.exists(config_file) and config_file != DEFAULT_CONFIG_FILE: + logger.error(f"The configuration file {config_file} was not found. Make sure the file exists and try again, or run without the --config file to use the default settings.") + exit() - def run(self) -> None: + return read_yaml(config_file) + + def run(self, args: list) -> None: + self.setup_basic_parser() # parse the known arguments for now (basically, we want the config file) + basic_config, unused_args = self.basic_parser.parse_known_args(args) - # load the config file to get the list of enabled items - basic_config, unused_args = self.basic_parser.parse_known_args() - + # setup any custom module paths, so they'll show in the help and for arg parsing setup_paths(basic_config.module_paths) # if help flag was called, then show the help if basic_config.help: self.show_help(basic_config) - # load the config file - yaml_config = {} - - if not os.path.exists(basic_config.config_file) and basic_config.config_file != DEFAULT_CONFIG_FILE: - logger.error(f"The configuration file {basic_config.config_file} was not found. Make sure the file exists and try again, or run without the --config file to use the default settings.") - exit() - - - yaml_config = read_yaml(basic_config.config_file) + yaml_config = self.load_config(basic_config.config_file) self.setup_complete_parser(basic_config, yaml_config, unused_args) logger.info(f"======== Welcome to the AUTO ARCHIVER ({__version__}) ==========") diff --git a/src/auto_archiver/modules/instagram_api_extractor/instagram_api_extractor.py b/src/auto_archiver/modules/instagram_api_extractor/instagram_api_extractor.py index 4a18228..5dad0ba 100644 --- a/src/auto_archiver/modules/instagram_api_extractor/instagram_api_extractor.py +++ b/src/auto_archiver/modules/instagram_api_extractor/instagram_api_extractor.py @@ -28,7 +28,7 @@ class InstagramAPIExtractor(Extractor): # TODO: improvement collect aggregates of locations[0].location and mentions for all posts """ - global_pattern = re.compile( + valid_url = re.compile( r"(?:(?:http|https):\/\/)?(?:www.)?(?:instagram.com)\/(stories(?:\/highlights)?|p|reel)?\/?([^\/\?]*)\/?(\d+)?" ) @@ -44,7 +44,7 @@ class InstagramAPIExtractor(Extractor): url.replace("instagr.com", "instagram.com").replace( "instagr.am", "instagram.com" ) - insta_matches = self.global_pattern.findall(url) + insta_matches = self.valid_url.findall(url) logger.info(f"{insta_matches=}") if not len(insta_matches) or len(insta_matches[0]) != 3: return diff --git a/src/auto_archiver/modules/instagram_extractor/instagram_extractor.py b/src/auto_archiver/modules/instagram_extractor/instagram_extractor.py index 1cdb0b1..3cf0362 100644 --- a/src/auto_archiver/modules/instagram_extractor/instagram_extractor.py +++ b/src/auto_archiver/modules/instagram_extractor/instagram_extractor.py @@ -16,10 +16,13 @@ class InstagramExtractor(Extractor): Uses Instaloader to download either a post (inc images, videos, text) or as much as possible from a profile (posts, stories, highlights, ...) """ # NB: post regex should be tested before profile + + valid_url = re.compile(r"(?:(?:http|https):\/\/)?(?:www.)?(?:instagram.com|instagr.am|instagr.com)\/") + # https://regex101.com/r/MGPquX/1 - post_pattern = re.compile(r"(?:(?:http|https):\/\/)?(?:www.)?(?:instagram.com|instagr.am|instagr.com)\/(?:p|reel)\/(\w+)") + post_pattern = re.compile(r"{valid_url}(?:p|reel)\/(\w+)".format(valid_url=valid_url)) # https://regex101.com/r/6Wbsxa/1 - profile_pattern = re.compile(r"(?:(?:http|https):\/\/)?(?:www.)?(?:instagram.com|instagr.am|instagr.com)\/(\w+)") + profile_pattern = re.compile(r"{valid_url}(\w+)".format(valid_url=valid_url)) # TODO: links to stories def setup(self, config: dict) -> None: diff --git a/src/auto_archiver/modules/telethon_extractor/telethon_extractor.py b/src/auto_archiver/modules/telethon_extractor/telethon_extractor.py index f378e7e..8a08954 100644 --- a/src/auto_archiver/modules/telethon_extractor/telethon_extractor.py +++ b/src/auto_archiver/modules/telethon_extractor/telethon_extractor.py @@ -14,7 +14,7 @@ from auto_archiver.utils import random_str class TelethonArchiver(Extractor): - link_pattern = re.compile(r"https:\/\/t\.me(\/c){0,1}\/(.+)\/(\d+)") + valid_url = re.compile(r"https:\/\/t\.me(\/c){0,1}\/(.+)\/(\d+)") invite_pattern = re.compile(r"t.me(\/joinchat){0,1}\/\+?(.+)") @@ -92,7 +92,7 @@ class TelethonArchiver(Extractor): """ url = item.get_url() # detect URLs that we definitely cannot handle - match = self.link_pattern.search(url) + match = self.valid_url.search(url) logger.debug(f"TELETHON: {match=}") if not match: return False diff --git a/src/auto_archiver/modules/twitter_api_extractor/twitter_api_extractor.py b/src/auto_archiver/modules/twitter_api_extractor/twitter_api_extractor.py index ede0239..0434190 100644 --- a/src/auto_archiver/modules/twitter_api_extractor/twitter_api_extractor.py +++ b/src/auto_archiver/modules/twitter_api_extractor/twitter_api_extractor.py @@ -12,7 +12,7 @@ from auto_archiver.core import Extractor from auto_archiver.core import Metadata,Media class TwitterApiExtractor(Extractor): - link_pattern = re.compile(r"(?:twitter|x).com\/(?:\#!\/)?(\w+)\/status(?:es)?\/(\d+)") + valid_url = re.compile(r"(?:twitter|x).com\/(?:\#!\/)?(\w+)\/status(?:es)?\/(\d+)") def setup(self, config: dict) -> None: super().setup(config) @@ -54,7 +54,7 @@ class TwitterApiExtractor(Extractor): def get_username_tweet_id(self, url): # detect URLs that we definitely cannot handle - matches = self.link_pattern.findall(url) + matches = self.valid_url.findall(url) if not len(matches): return False, False username, tweet_id = matches[0] # only one URL supported diff --git a/src/auto_archiver/utils/url.py b/src/auto_archiver/utils/url.py index 7586cca..3b67514 100644 --- a/src/auto_archiver/utils/url.py +++ b/src/auto_archiver/utils/url.py @@ -2,8 +2,11 @@ import re from urllib.parse import urlparse, urlunparse class UrlUtil: - telegram_private = re.compile(r"https:\/\/t\.me(\/c)\/(.+)\/(\d+)") - is_istagram = re.compile(r"https:\/\/www\.instagram\.com") + + AUTHWALL_URLS = [ + re.compile(r"https:\/\/t\.me(\/c)\/(.+)\/(\d+)"), # telegram private channels + re.compile(r"https:\/\/www\.instagram\.com"), # instagram + ] @staticmethod def clean(url: str) -> str: return url @@ -13,8 +16,9 @@ class UrlUtil: """ checks if URL is behind an authentication wall meaning steps like wayback, wacz, ... may not work """ - if UrlUtil.telegram_private.match(url): return True - if UrlUtil.is_istagram.match(url): return True + for regex in UrlUtil.AUTHWALL_URLS: + if regex.match(url): + return True return False diff --git a/tests/data/example_module/example_module.py b/tests/data/example_module/example_module.py deleted file mode 100644 index bce8ba4..0000000 --- a/tests/data/example_module/example_module.py +++ /dev/null @@ -1,4 +0,0 @@ -from auto_archiver.core.extractor import Extractor -class ExampleModule(Extractor): - def download(self, item): - print("do something") \ No newline at end of file diff --git a/tests/data/example_module/__init__.py b/tests/data/test_modules/example_module/__init__.py similarity index 100% rename from tests/data/example_module/__init__.py rename to tests/data/test_modules/example_module/__init__.py diff --git a/tests/data/example_module/__manifest__.py b/tests/data/test_modules/example_module/__manifest__.py similarity index 55% rename from tests/data/example_module/__manifest__.py rename to tests/data/test_modules/example_module/__manifest__.py index 19a85f9..f2ebdbf 100644 --- a/tests/data/example_module/__manifest__.py +++ b/tests/data/test_modules/example_module/__manifest__.py @@ -1,10 +1,11 @@ { "name": "Example Module", - "type": ["extractor"], + "type": ["extractor", "feeder", "formatter", "storage", "enricher", "database"], "requires_setup": False, "dependencies": {"python": ["loguru"] }, "configs": { - "csv_file": {"default": "db.csv", "help": "CSV file name"} + "csv_file": {"default": "db.csv", "help": "CSV file name"}, + "required_field": {"required": True, "help": "required field in the CSV file"}, }, } \ No newline at end of file diff --git a/tests/data/test_modules/example_module/example_module.py b/tests/data/test_modules/example_module/example_module.py new file mode 100644 index 0000000..7def054 --- /dev/null +++ b/tests/data/test_modules/example_module/example_module.py @@ -0,0 +1,28 @@ +from auto_archiver.core import Extractor, Enricher, Feeder, Database, Storage, Formatter, Metadata + +class ExampleModule(Extractor, Enricher, Feeder, Database, Storage, Formatter): + def download(self, item): + print("download") + + def __iter__(self): + yield Metadata().set_url("https://example.com") + + + def done(self, result): + print("done") + + def enrich(self, to_enrich): + print("enrich") + + def get_cdn_url(self, media): + return "nice_url" + + def save(self, item): + print("save") + + def uploadf(self, file, key, **kwargs): + print("uploadf") + + + def format(self, item): + print("format") diff --git a/tests/data/test_orchestration.yaml b/tests/data/test_orchestration.yaml new file mode 100644 index 0000000..ec6af35 --- /dev/null +++ b/tests/data/test_orchestration.yaml @@ -0,0 +1,16 @@ +steps: + feeders: + - example_module + extractors: + - example_module + formatters: + - example_module + storages: + - example_module + databases: + - example_module + enrichers: + - example_module + + +# Global configuration \ No newline at end of file diff --git a/tests/extractors/test_extractor_base.py b/tests/extractors/test_extractor_base.py index f6be70b..24689b4 100644 --- a/tests/extractors/test_extractor_base.py +++ b/tests/extractors/test_extractor_base.py @@ -9,7 +9,7 @@ class TestExtractorBase(object): config: dict = None @pytest.fixture(autouse=True) - def setup_archiver(self, setup_module): + def setup_extractor(self, setup_module): assert self.extractor_module is not None, "self.extractor_module must be set on the subclass" assert self.config is not None, "self.config must be a dict set on the subclass" diff --git a/tests/extractors/test_instagram_extractor.py b/tests/extractors/test_instagram_extractor.py new file mode 100644 index 0000000..7efe1b1 --- /dev/null +++ b/tests/extractors/test_instagram_extractor.py @@ -0,0 +1,21 @@ +import pytest + +from auto_archiver.modules.instagram_extractor import InstagramExtractor +from .test_extractor_base import TestExtractorBase + +class TestInstagramExtractor(TestExtractorBase): + + extractor_module: str = 'instagram_extractor' + config: dict = {} + + @pytest.mark.parametrize("url", [ + "https://www.instagram.com/p/", + "https://www.instagram.com/p/1234567890/", + "https://www.instagram.com/reel/1234567890/", + "https://www.instagram.com/username/", + "https://www.instagram.com/username/stories/", + "https://www.instagram.com/username/highlights/", + ]) + def test_regex_matches(self, url): + # post + assert InstagramExtractor.valid_url.match(url) diff --git a/tests/test_modules.py b/tests/test_modules.py index decc616..a4c0ec8 100644 --- a/tests/test_modules.py +++ b/tests/test_modules.py @@ -7,7 +7,7 @@ def example_module(): import auto_archiver previous_path = auto_archiver.modules.__path__ - auto_archiver.modules.__path__.append("tests/data/") + auto_archiver.modules.__path__.append("tests/data/test_modules/") module = get_module_lazy("example_module") yield module diff --git a/tests/test_orchestrator.py b/tests/test_orchestrator.py new file mode 100644 index 0000000..9e81df7 --- /dev/null +++ b/tests/test_orchestrator.py @@ -0,0 +1,123 @@ +import pytest +import sys +from argparse import ArgumentParser +from auto_archiver.core.orchestrator import ArchivingOrchestrator +from auto_archiver.version import __version__ +from auto_archiver.core.config import read_yaml, store_yaml + +TEST_ORCHESTRATION = "tests/data/test_orchestration.yaml" +TEST_MODULES = "tests/data/test_modules/" + +@pytest.fixture +def test_args(): + return ["--config", TEST_ORCHESTRATION, + "--module_paths", TEST_MODULES, + "--example_module.required_field", "some_value"] # just set this for normal testing, we will remove it later + +@pytest.fixture +def orchestrator(): + yield ArchivingOrchestrator() + # hack - the loguru logger starts with one logger, but if orchestrator has run before + # it'll remove the default logger, add it back in: + + from loguru import logger + + if not logger._core.handlers.get(0): + logger._core.handlers_count = 0 + logger.add(sys.stderr) + # and remove the custom logger + if logger._core.handlers.get(1): + logger.remove(1) + +@pytest.fixture +def basic_parser(orchestrator) -> ArgumentParser: + return orchestrator.setup_basic_parser() + +def test_setup_orchestrator(orchestrator): + assert orchestrator is not None + +def test_parse_config(): + pass + +def test_parse_basic(basic_parser): + args = basic_parser.parse_args(["--config", TEST_ORCHESTRATION]) + assert args.config_file == TEST_ORCHESTRATION + +@pytest.mark.parametrize("mode", ["simple", "full"]) +def test_mode(basic_parser, mode): + args = basic_parser.parse_args(["--mode", mode]) + assert args.mode == mode + +def test_mode_invalid(basic_parser, capsys): + with pytest.raises(SystemExit) as exit_error: + basic_parser.parse_args(["--mode", "invalid"]) + assert exit_error.value.code == 2 + assert "invalid choice" in capsys.readouterr().err + +def test_version(basic_parser, capsys): + with pytest.raises(SystemExit) as exit_error: + basic_parser.parse_args(["--version"]) + assert exit_error.value.code == 0 + assert capsys.readouterr().out == f"{__version__}\n" + +def test_help(orchestrator, basic_parser, capsys): + + args = basic_parser.parse_args(["--help"]) + assert args.help == True + + # test the show_help() on orchestrator + with pytest.raises(SystemExit) as exit_error: + orchestrator.show_help(args) + + assert exit_error.value.code == 0 + assert "Usage: auto-archiver [--help] [--version] [--config CONFIG_FILE]" in capsys.readouterr().out + + +def test_add_custom_modules_path(orchestrator, test_args): + orchestrator.run(test_args) + + import auto_archiver + assert "tests/data/test_modules/" in auto_archiver.modules.__path__ + +def test_add_custom_modules_path_invalid(orchestrator, caplog, test_args): + + orchestrator.run(test_args + # we still need to load the real path to get the example_module + ["--module_paths", "tests/data/invalid_test_modules/"]) + + # assert False + assert caplog.records[0].message == "Path 'tests/data/invalid_test_modules/' does not exist. Skipping..." + + +def test_check_required_values(orchestrator, caplog, test_args): + # drop the example_module.required_field from the test_args + test_args = test_args[:-2] + + with pytest.raises(SystemExit) as exit_error: + orchestrator.run(test_args) + + assert caplog.records[1].message == "the following arguments are required: --example_module.required_field" + +def test_get_required_values_from_config(orchestrator, test_args, tmp_path): + + # load the default example yaml, add a required field, then run the orchestrator + test_yaml = read_yaml(TEST_ORCHESTRATION) + test_yaml['example_module'] = {'required_field': 'some_value'} + # write it to a temp file + tmp_file = (tmp_path / "temp_config.yaml").as_posix() + store_yaml(test_yaml, tmp_file) + + # run the orchestrator + orchestrator.run(["--config", tmp_file, "--module_paths", TEST_MODULES]) + + # should run OK, since there are no missing required fields + + # basic_args = basic_parser.parse_known_args(test_args) + # test_yaml = read_yaml(TEST_ORCHESTRATION) + # test_yaml['example_module'] = {'required_field': 'some_value'} + + # # monkey patch the example_module to have a 'configs' setting of 'my_var' with required=True + # # load the module first + # m = get_module_lazy("example_module") + + # orchestrator.setup_complete_parser(basic_args, test_yaml, unused_args=[]) + # assert orchestrator.config is not None \ No newline at end of file From fade68c6f48bcc6cc69c6dcf05e4b398e5439dd0 Mon Sep 17 00:00:00 2001 From: Patrick Robertson Date: Thu, 30 Jan 2025 13:45:24 +0100 Subject: [PATCH 06/62] Fix up unit tests - dataclass + subclasses not having @dataclass was breaking it --- src/auto_archiver/core/extractor.py | 1 - .../modules/twitter_api_extractor/twitter_api_extractor.py | 5 +++-- tests/extractors/test_extractor_base.py | 5 ++++- tests/test_orchestrator.py | 5 +++++ 4 files changed, 12 insertions(+), 4 deletions(-) diff --git a/src/auto_archiver/core/extractor.py b/src/auto_archiver/core/extractor.py index 51d784f..ed261eb 100644 --- a/src/auto_archiver/core/extractor.py +++ b/src/auto_archiver/core/extractor.py @@ -21,7 +21,6 @@ import re from ..core import Metadata, ArchivingContext, BaseModule -@dataclass class Extractor(BaseModule): """ Base class for implementing extractors in the media archiving framework. diff --git a/src/auto_archiver/modules/twitter_api_extractor/twitter_api_extractor.py b/src/auto_archiver/modules/twitter_api_extractor/twitter_api_extractor.py index 0434190..6573475 100644 --- a/src/auto_archiver/modules/twitter_api_extractor/twitter_api_extractor.py +++ b/src/auto_archiver/modules/twitter_api_extractor/twitter_api_extractor.py @@ -9,10 +9,11 @@ from pytwitter import Api from slugify import slugify from auto_archiver.core import Extractor -from auto_archiver.core import Metadata,Media +from auto_archiver.core import Metadata, Media class TwitterApiExtractor(Extractor): - valid_url = re.compile(r"(?:twitter|x).com\/(?:\#!\/)?(\w+)\/status(?:es)?\/(\d+)") + + valid_url: re.Pattern = re.compile(r"(?:twitter|x).com\/(?:\#!\/)?(\w+)\/status(?:es)?\/(\d+)") def setup(self, config: dict) -> None: super().setup(config) diff --git a/tests/extractors/test_extractor_base.py b/tests/extractors/test_extractor_base.py index 24689b4..6e77ec3 100644 --- a/tests/extractors/test_extractor_base.py +++ b/tests/extractors/test_extractor_base.py @@ -1,8 +1,11 @@ +from typing import Type + import pytest from auto_archiver.core.metadata import Metadata from auto_archiver.core.extractor import Extractor + class TestExtractorBase(object): extractor_module: str = None @@ -13,7 +16,7 @@ class TestExtractorBase(object): assert self.extractor_module is not None, "self.extractor_module must be set on the subclass" assert self.config is not None, "self.config must be a dict set on the subclass" - self.extractor: Extractor = setup_module(self.extractor_module, self.config) + self.extractor: Type[Extractor] = setup_module(self.extractor_module, self.config) def assertValidResponseMetadata(self, test_response: Metadata, title: str, timestamp: str, status: str = ""): assert test_response is not False diff --git a/tests/test_orchestrator.py b/tests/test_orchestrator.py index 9e81df7..03cb521 100644 --- a/tests/test_orchestrator.py +++ b/tests/test_orchestrator.py @@ -4,6 +4,7 @@ from argparse import ArgumentParser from auto_archiver.core.orchestrator import ArchivingOrchestrator from auto_archiver.version import __version__ from auto_archiver.core.config import read_yaml, store_yaml +from auto_archiver.core.module import _LAZY_LOADED_MODULES TEST_ORCHESTRATION = "tests/data/test_orchestration.yaml" TEST_MODULES = "tests/data/test_modules/" @@ -29,6 +30,10 @@ def orchestrator(): if logger._core.handlers.get(1): logger.remove(1) + # delete out any loaded modules + _LAZY_LOADED_MODULES.clear() + + @pytest.fixture def basic_parser(orchestrator) -> ArgumentParser: return orchestrator.setup_basic_parser() From 527438826c65cf5340b1d3560e1f001b77017324 Mon Sep 17 00:00:00 2001 From: erinhmclark Date: Thu, 30 Jan 2025 13:04:51 +0000 Subject: [PATCH 07/62] Fix manifests for required configs. --- .../modules/api_db/__manifest__.py | 1 - .../modules/atlos_feeder/__manifest__.py | 1 - .../modules/gdrive_storage/__manifest__.py | 3 +- .../instagram_api_extractor/__manifest__.py | 3 +- .../instagram_extractor/__manifest__.py | 4 +- .../modules/vk_extractor/__manifest__.py | 6 +- .../modules/wayback_enricher/__init__.py | 1 - .../modules/wayback_enricher/__manifest__.py | 30 ---------- .../wayback_extractor_enricher/__init__.py | 1 + .../__manifest__.py | 56 +++++++++++++++++++ .../wayback_extractor_enricher.py} | 0 11 files changed, 62 insertions(+), 44 deletions(-) delete mode 100644 src/auto_archiver/modules/wayback_enricher/__init__.py delete mode 100644 src/auto_archiver/modules/wayback_enricher/__manifest__.py create mode 100644 src/auto_archiver/modules/wayback_extractor_enricher/__init__.py create mode 100644 src/auto_archiver/modules/wayback_extractor_enricher/__manifest__.py rename src/auto_archiver/modules/{wayback_enricher/wayback_enricher.py => wayback_extractor_enricher/wayback_extractor_enricher.py} (100%) diff --git a/src/auto_archiver/modules/api_db/__manifest__.py b/src/auto_archiver/modules/api_db/__manifest__.py index 3874496..698c2e4 100644 --- a/src/auto_archiver/modules/api_db/__manifest__.py +++ b/src/auto_archiver/modules/api_db/__manifest__.py @@ -8,7 +8,6 @@ }, "configs": { "api_endpoint": { - "default": None, "required": True, "help": "API endpoint where calls are made to", }, diff --git a/src/auto_archiver/modules/atlos_feeder/__manifest__.py b/src/auto_archiver/modules/atlos_feeder/__manifest__.py index 5ae3540..d59f420 100644 --- a/src/auto_archiver/modules/atlos_feeder/__manifest__.py +++ b/src/auto_archiver/modules/atlos_feeder/__manifest__.py @@ -7,7 +7,6 @@ }, "configs": { "api_token": { - "default": None, "type": "str", "required": True, "help": "An Atlos API token. For more information, see https://docs.atlos.org/technical/api/", diff --git a/src/auto_archiver/modules/gdrive_storage/__manifest__.py b/src/auto_archiver/modules/gdrive_storage/__manifest__.py index 2ca7e27..632e52b 100644 --- a/src/auto_archiver/modules/gdrive_storage/__manifest__.py +++ b/src/auto_archiver/modules/gdrive_storage/__manifest__.py @@ -22,8 +22,7 @@ "help": "how to name stored files: 'random' creates a random string; 'static' uses a replicable strategy such as a hash.", "choices": ["random", "static"], }, - "root_folder_id": {"default": None, - # "required": True, + "root_folder_id": {"required": True, "help": "root google drive folder ID to use as storage, found in URL: 'https://drive.google.com/drive/folders/FOLDER_ID'"}, "oauth_token": {"default": None, "help": "JSON filename with Google Drive OAuth token: check auto-archiver repository scripts folder for create_update_gdrive_oauth_token.py. NOTE: storage used will count towards owner of GDrive folder, therefore it is best to use oauth_token_filename over service_account."}, diff --git a/src/auto_archiver/modules/instagram_api_extractor/__manifest__.py b/src/auto_archiver/modules/instagram_api_extractor/__manifest__.py index a958a99..2d8f1d9 100644 --- a/src/auto_archiver/modules/instagram_api_extractor/__manifest__.py +++ b/src/auto_archiver/modules/instagram_api_extractor/__manifest__.py @@ -12,8 +12,7 @@ "configs": { "access_token": {"default": None, "help": "a valid instagrapi-api token"}, - "api_endpoint": {"default": None, - # "required": True, + "api_endpoint": {"required": True, "help": "API endpoint to use"}, "full_profile": { "default": False, diff --git a/src/auto_archiver/modules/instagram_extractor/__manifest__.py b/src/auto_archiver/modules/instagram_extractor/__manifest__.py index d8e4a9b..05cae19 100644 --- a/src/auto_archiver/modules/instagram_extractor/__manifest__.py +++ b/src/auto_archiver/modules/instagram_extractor/__manifest__.py @@ -9,11 +9,9 @@ }, "requires_setup": True, "configs": { - "username": {"default": None, - "required": True, + "username": {"required": True, "help": "a valid Instagram username"}, "password": { - "default": None, "required": True, "help": "the corresponding Instagram account password", }, diff --git a/src/auto_archiver/modules/vk_extractor/__manifest__.py b/src/auto_archiver/modules/vk_extractor/__manifest__.py index 033fe50..61e454e 100644 --- a/src/auto_archiver/modules/vk_extractor/__manifest__.py +++ b/src/auto_archiver/modules/vk_extractor/__manifest__.py @@ -7,11 +7,9 @@ "python": ["loguru", "vk_url_scraper"], }, "configs": { - "username": {"default": None, - "required": True, + "username": {"required": True, "help": "valid VKontakte username"}, - "password": {"default": None, - "required": True, + "password": {"required": True, "help": "valid VKontakte password"}, "session_file": { "default": "secrets/vk_config.v2.json", diff --git a/src/auto_archiver/modules/wayback_enricher/__init__.py b/src/auto_archiver/modules/wayback_enricher/__init__.py deleted file mode 100644 index 9782831..0000000 --- a/src/auto_archiver/modules/wayback_enricher/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from .wayback_enricher import WaybackExtractorEnricher \ No newline at end of file diff --git a/src/auto_archiver/modules/wayback_enricher/__manifest__.py b/src/auto_archiver/modules/wayback_enricher/__manifest__.py deleted file mode 100644 index 5d1fe25..0000000 --- a/src/auto_archiver/modules/wayback_enricher/__manifest__.py +++ /dev/null @@ -1,30 +0,0 @@ -{ - "name": "Wayback Machine Enricher", - "type": ["enricher", "archiver"], - "requires_setup": True, - "dependencies": { - "python": ["loguru", "requests"], - }, - "entry_point": "wayback_enricher::WaybackExtractorEnricher", - "configs": { - "timeout": {"default": 15, "help": "seconds to wait for successful archive confirmation from wayback, if more than this passes the result contains the job_id so the status can later be checked manually."}, - "if_not_archived_within": {"default": None, "help": "only tell wayback to archive if no archive is available before the number of seconds specified, use None to ignore this option. For more information: https://docs.google.com/document/d/1Nsv52MvSjbLb2PCpHlat0gkzw0EvtSgpKHu4mk0MnrA"}, - "key": {"default": None, "required": True, "help": "wayback API key. to get credentials visit https://archive.org/account/s3.php"}, - "secret": {"default": None, "required": True, "help": "wayback API secret. to get credentials visit https://archive.org/account/s3.php"}, - "proxy_http": {"default": None, "help": "http proxy to use for wayback requests, eg http://proxy-user:password@proxy-ip:port"}, - "proxy_https": {"default": None, "help": "https proxy to use for wayback requests, eg https://proxy-user:password@proxy-ip:port"}, - }, - "description": """ - Submits the current URL to the Wayback Machine for archiving and returns either a job ID or the completed archive URL. - - ### Features - - Archives URLs using the Internet Archive's Wayback Machine API. - - Supports conditional archiving based on the existence of prior archives within a specified time range. - - Provides proxies for HTTP and HTTPS requests. - - Fetches and confirms the archive URL or provides a job ID for later status checks. - - ### Notes - - Requires a valid Wayback Machine API key and secret. - - Handles rate-limiting by Wayback Machine and retries status checks with exponential backoff. - """ -} diff --git a/src/auto_archiver/modules/wayback_extractor_enricher/__init__.py b/src/auto_archiver/modules/wayback_extractor_enricher/__init__.py new file mode 100644 index 0000000..b69332d --- /dev/null +++ b/src/auto_archiver/modules/wayback_extractor_enricher/__init__.py @@ -0,0 +1 @@ +from .wayback_extractor_enricher import WaybackExtractorEnricher \ No newline at end of file diff --git a/src/auto_archiver/modules/wayback_extractor_enricher/__manifest__.py b/src/auto_archiver/modules/wayback_extractor_enricher/__manifest__.py new file mode 100644 index 0000000..baecc14 --- /dev/null +++ b/src/auto_archiver/modules/wayback_extractor_enricher/__manifest__.py @@ -0,0 +1,56 @@ +{ + "name": "Wayback Machine Enricher", + "type": ["enricher", "archiver"], + "entry_point": "wayback_extractor_enricher::WaybackExtractorEnricher", + "requires_setup": True, + "dependencies": { + "python": ["loguru", "requests"], + }, + "configs": { + "timeout": { + "default": 15, + "help": "seconds to wait for successful archive confirmation from wayback, if more than this passes the result contains the job_id so the status can later be checked manually.", + }, + "if_not_archived_within": { + "default": None, + "help": "only tell wayback to archive if no archive is available before the number of seconds specified, use None to ignore this option. For more information: https://docs.google.com/document/d/1Nsv52MvSjbLb2PCpHlat0gkzw0EvtSgpKHu4mk0MnrA", + }, + "key": { + "required": True, + "help": "wayback API key. to get credentials visit https://archive.org/account/s3.php", + }, + "secret": { + "required": True, + "help": "wayback API secret. to get credentials visit https://archive.org/account/s3.php", + }, + "proxy_http": { + "default": None, + "help": "http proxy to use for wayback requests, eg http://proxy-user:password@proxy-ip:port", + }, + "proxy_https": { + "default": None, + "help": "https proxy to use for wayback requests, eg https://proxy-user:password@proxy-ip:port", + }, + }, + "description": """ + Submits the current URL to the Wayback Machine for archiving and returns either a job ID or the completed archive URL. + + ### Features + - Archives URLs using the Internet Archive's Wayback Machine API. + - Supports conditional archiving based on the existence of prior archives within a specified time range. + - Provides proxies for HTTP and HTTPS requests. + - Fetches and confirms the archive URL or provides a job ID for later status checks. + + ### Notes + - Requires a valid Wayback Machine API key and secret. + - Handles rate-limiting by Wayback Machine and retries status checks with exponential backoff. + + ### Steps to Get an Wayback API Key: + - Sign up for an account at [Internet Archive](https://archive.org/account/signup). + - Log in to your account. + - Navigte to your [account settings](https://archive.org/account). + - or: https://archive.org/developers/tutorial-get-ia-credentials.html + - Under Wayback Machine API Keys, generate a new key. + - Note down your API key and secret, as they will be required for authentication. + """, +} diff --git a/src/auto_archiver/modules/wayback_enricher/wayback_enricher.py b/src/auto_archiver/modules/wayback_extractor_enricher/wayback_extractor_enricher.py similarity index 100% rename from src/auto_archiver/modules/wayback_enricher/wayback_enricher.py rename to src/auto_archiver/modules/wayback_extractor_enricher/wayback_extractor_enricher.py From 953011f36851b887715788679e66c052484d573f Mon Sep 17 00:00:00 2001 From: Patrick Robertson Date: Thu, 30 Jan 2025 14:39:52 +0100 Subject: [PATCH 08/62] Don't make modules 'dataclasses' --- src/auto_archiver/core/database.py | 5 +---- src/auto_archiver/core/enricher.py | 4 +--- src/auto_archiver/core/feeder.py | 3 --- src/auto_archiver/core/formatter.py | 2 -- src/auto_archiver/core/storage.py | 8 +++----- .../modules/html_formatter/html_formatter.py | 2 -- .../modules/mute_formatter/mute_formatter.py | 2 -- 7 files changed, 5 insertions(+), 21 deletions(-) diff --git a/src/auto_archiver/core/database.py b/src/auto_archiver/core/database.py index f7deaef..0eb5d81 100644 --- a/src/auto_archiver/core/database.py +++ b/src/auto_archiver/core/database.py @@ -1,12 +1,9 @@ from __future__ import annotations -from dataclasses import dataclass -from abc import abstractmethod, ABC +from abc import abstractmethod from typing import Union from auto_archiver.core import Metadata, BaseModule - -@dataclass class Database(BaseModule): def started(self, item: Metadata) -> None: diff --git a/src/auto_archiver/core/enricher.py b/src/auto_archiver/core/enricher.py index fe0d05f..0e50fa9 100644 --- a/src/auto_archiver/core/enricher.py +++ b/src/auto_archiver/core/enricher.py @@ -9,11 +9,9 @@ the archiving step and before storage or formatting. Enrichers are optional but highly useful for making the archived data more powerful. """ from __future__ import annotations -from dataclasses import dataclass -from abc import abstractmethod, ABC +from abc import abstractmethod from auto_archiver.core import Metadata, BaseModule -@dataclass class Enricher(BaseModule): """Base classes and utilities for enrichers in the Auto-Archiver system.""" diff --git a/src/auto_archiver/core/feeder.py b/src/auto_archiver/core/feeder.py index e539f5f..352cfd9 100644 --- a/src/auto_archiver/core/feeder.py +++ b/src/auto_archiver/core/feeder.py @@ -1,11 +1,8 @@ from __future__ import annotations -from dataclasses import dataclass from abc import abstractmethod from auto_archiver.core import Metadata from auto_archiver.core import BaseModule - -@dataclass class Feeder(BaseModule): @abstractmethod diff --git a/src/auto_archiver/core/formatter.py b/src/auto_archiver/core/formatter.py index beb0c0d..cf27cb3 100644 --- a/src/auto_archiver/core/formatter.py +++ b/src/auto_archiver/core/formatter.py @@ -1,10 +1,8 @@ from __future__ import annotations -from dataclasses import dataclass from abc import abstractmethod from auto_archiver.core import Metadata, Media, BaseModule -@dataclass class Formatter(BaseModule): @abstractmethod diff --git a/src/auto_archiver/core/storage.py b/src/auto_archiver/core/storage.py index 5274204..b40c5cc 100644 --- a/src/auto_archiver/core/storage.py +++ b/src/auto_archiver/core/storage.py @@ -1,18 +1,16 @@ from __future__ import annotations from abc import abstractmethod -from dataclasses import dataclass from typing import IO, Optional import os +from loguru import logger +from slugify import slugify + from auto_archiver.utils.misc import random_str from auto_archiver.core import Media, BaseModule, ArchivingContext, Metadata from auto_archiver.modules.hash_enricher.hash_enricher import HashEnricher -from loguru import logger -from slugify import slugify - -@dataclass class Storage(BaseModule): def store(self, media: Media, url: str, metadata: Optional[Metadata]=None) -> None: diff --git a/src/auto_archiver/modules/html_formatter/html_formatter.py b/src/auto_archiver/modules/html_formatter/html_formatter.py index 8f006e0..bfc2efa 100644 --- a/src/auto_archiver/modules/html_formatter/html_formatter.py +++ b/src/auto_archiver/modules/html_formatter/html_formatter.py @@ -1,5 +1,4 @@ from __future__ import annotations -from dataclasses import dataclass import mimetypes, os, pathlib from jinja2 import Environment, FileSystemLoader from urllib.parse import quote @@ -14,7 +13,6 @@ from auto_archiver.modules.hash_enricher import HashEnricher from auto_archiver.utils.misc import random_str from auto_archiver.core.module import get_module -@dataclass class HtmlFormatter(Formatter): environment: Environment = None template: any = None diff --git a/src/auto_archiver/modules/mute_formatter/mute_formatter.py b/src/auto_archiver/modules/mute_formatter/mute_formatter.py index 1c7cca2..129ddcb 100644 --- a/src/auto_archiver/modules/mute_formatter/mute_formatter.py +++ b/src/auto_archiver/modules/mute_formatter/mute_formatter.py @@ -1,11 +1,9 @@ from __future__ import annotations -from dataclasses import dataclass from auto_archiver.core import Metadata, Media from auto_archiver.core import Formatter -@dataclass class MuteFormatter(Formatter): def format(self, item: Metadata) -> Media: return None From d6b4b7a932b7c8840265890583b79dc7e5038b47 Mon Sep 17 00:00:00 2001 From: Patrick Robertson Date: Thu, 30 Jan 2025 16:43:09 +0100 Subject: [PATCH 09/62] Further cleanup * Removes (partly) the ArchivingOrchestrator * Removes the cli_feeder module, and makes it the 'default', allowing you to pass URLs directly on the command line, without having to use the cumbersome --cli_feeder.urls. Just do auto-archiver https://my.url.com * More unit tests * Improved error handling --- src/auto_archiver/__main__.py | 2 +- src/auto_archiver/core/base_module.py | 100 +++++++++ src/auto_archiver/core/config.py | 33 ++- src/auto_archiver/core/context.py | 10 +- src/auto_archiver/core/extractor.py | 3 +- src/auto_archiver/core/module.py | 54 +---- src/auto_archiver/core/orchestrator.py | 200 ++++++++++++++---- .../enrichers/screenshot_enricher.py | 40 ++++ src/auto_archiver/feeders/csv_feeder.py | 38 ++++ .../modules/atlos_feeder/atlos_feeder.py | 2 - .../modules/cli_feeder/__init__.py | 1 - .../modules/cli_feeder/__manifest__.py | 27 --- .../modules/cli_feeder/cli_feeder.py | 15 -- .../modules/csv_feeder/__manifest__.py | 1 - .../modules/csv_feeder/csv_feeder.py | 4 +- .../generic_extractor/generic_extractor.py | 6 +- .../modules/html_formatter/html_formatter.py | 4 +- .../screenshot_enricher.py | 6 +- .../modules/ssl_enricher/ssl_enricher.py | 2 +- .../telethon_extractor/telethon_extractor.py | 4 +- .../thumbnail_enricher/thumbnail_enricher.py | 2 +- .../timestamping_enricher.py | 8 +- .../modules/vk_extractor/vk_extractor.py | 4 +- .../modules/wacz_enricher/wacz_enricher.py | 6 +- tests/__init__.py | 3 +- tests/conftest.py | 6 + tests/test_orchestrator.py | 27 ++- 27 files changed, 417 insertions(+), 191 deletions(-) create mode 100644 src/auto_archiver/core/base_module.py create mode 100644 src/auto_archiver/enrichers/screenshot_enricher.py create mode 100644 src/auto_archiver/feeders/csv_feeder.py delete mode 100644 src/auto_archiver/modules/cli_feeder/__init__.py delete mode 100644 src/auto_archiver/modules/cli_feeder/__manifest__.py delete mode 100644 src/auto_archiver/modules/cli_feeder/cli_feeder.py diff --git a/src/auto_archiver/__main__.py b/src/auto_archiver/__main__.py index 0e2f54f..0023a59 100644 --- a/src/auto_archiver/__main__.py +++ b/src/auto_archiver/__main__.py @@ -3,7 +3,7 @@ from auto_archiver.core.orchestrator import ArchivingOrchestrator import sys def main(): - ArchivingOrchestrator().run(sys.argv) + ArchivingOrchestrator().run(sys.argv[1:]) if __name__ == "__main__": main() diff --git a/src/auto_archiver/core/base_module.py b/src/auto_archiver/core/base_module.py new file mode 100644 index 0000000..a9a904f --- /dev/null +++ b/src/auto_archiver/core/base_module.py @@ -0,0 +1,100 @@ + + +from urllib.parse import urlparse +from typing import Mapping, Any +from abc import ABC +from copy import deepcopy, copy +from tempfile import TemporaryDirectory + +from loguru import logger + +class BaseModule(ABC): + + """ + Base module class. All modules should inherit from this class. + + The exact methods a class implements will depend on the type of module it is, + however all modules have a .setup(config: dict) method to run any setup code + (e.g. logging in to a site, spinning up a browser etc.) + + See BaseModule.MODULE_TYPES for the types of modules you can create, noting that + a subclass can be of multiple types. For example, a module that extracts data from + a website and stores it in a database would be both an 'extractor' and a 'database' module. + + Each module is a python package, and should have a __manifest__.py file in the + same directory as the module file. The __manifest__.py specifies the module information + like name, author, version, dependencies etc. See BaseModule._DEFAULT_MANIFEST for the + default manifest structure. + + """ + + MODULE_TYPES = [ + 'feeder', + 'extractor', + 'enricher', + 'database', + 'storage', + 'formatter' + ] + + _DEFAULT_MANIFEST = { + 'name': '', # the display name of the module + 'author': 'Bellingcat', # creator of the module, leave this as Bellingcat or set your own name! + 'type': [], # the type of the module, can be one or more of BaseModule.MODULE_TYPES + 'requires_setup': True, # whether or not this module requires additional setup such as setting API Keys or installing additional softare + 'description': '', # a description of the module + 'dependencies': {}, # external dependencies, e.g. python packages or binaries, in dictionary format + 'entry_point': '', # the entry point for the module, in the format 'module_name::ClassName'. This can be left blank to use the default entry point of module_name::ModuleName + 'version': '1.0', # the version of the module + 'configs': {} # any configuration options this module has, these will be exposed to the user in the config file or via the command line +} + + config: Mapping[str, Any] + authentication: Mapping[str, Mapping[str, str]] + name: str + + # this is set by the orchestrator prior to archiving + tmp_dir: TemporaryDirectory = None + + def setup(self, config: dict): + + authentication = config.get('authentication', {}) + # extract out contatenated sites + for key, val in copy(authentication).items(): + if "," in key: + for site in key.split(","): + authentication[site] = val + del authentication[key] + + # this is important. Each instance is given its own deepcopied config, so modules cannot + # change values to affect other modules + config = deepcopy(config) + authentication = deepcopy(config.pop('authentication', {})) + + self.authentication = authentication + self.config = config + for key, val in config.get(self.name, {}).items(): + setattr(self, key, val) + + def repr(self): + return f"Module<'{self.display_name}' (config: {self.config[self.name]})>" + + def auth_for_site(self, site: str) -> dict: + # TODO: think about if/how we can deal with sites that have multiple domains (main one is x.com/twitter.com) + # for now, just hard code those. + + # SECURITY: parse the domain using urllib + site = urlparse(site).netloc + # add the 'www' version of the site to the list of sites to check + for to_try in [site, f"www.{site}"]: + if to_try in self.authentication: + return self.authentication[to_try] + + # do a fuzzy string match just to print a warning - don't use it since it's insecure + for key in self.authentication.keys(): + if key in site or site in key: + logger.warning(f"Could not find exact authentication information for site '{site}'. \ + did find information for '{key}' which is close, is this what you meant? \ + If so, edit your authentication settings to make sure it exactly matches.") + + return {} \ No newline at end of file diff --git a/src/auto_archiver/core/config.py b/src/auto_archiver/core/config.py index ca8ed25..2d462e4 100644 --- a/src/auto_archiver/core/config.py +++ b/src/auto_archiver/core/config.py @@ -15,8 +15,14 @@ from .module import BaseModule from typing import Any, List, Type, Tuple -yaml = YAML() +yaml: YAML = YAML() +b = yaml.load(""" + # This is a comment + site.com,site2.com: + key: value + key2: value2 + """) EMPTY_CONFIG = yaml.load(""" # Auto Archiver Configuration # Steps are the modules that will be run in the order they are defined @@ -25,6 +31,24 @@ steps:""" + "".join([f"\n {module}s: []" for module in BaseModule.MODULE_TYPES """ # Global configuration + +# Authentication +# a dictionary of authentication information that can be used by extractors to login to website. +# you can use a comma separated list for multiple domains on the same line (common usecase: x.com,twitter.com) +# Common login 'types' are username/password, cookie, api key/token. +# Some Examples: +# facebook.com: +# username: "my_username" +# password: "my_password" +# or for a site that uses an API key: +# twitter.com,x.com: +# api_key +# api_secret +# youtube.com: +# cookie: "login_cookie=value ; other_cookie=123" # multiple 'key=value' pairs should be separated by ; + +authentication: {} + # These are the global configurations that are used by the modules logging: @@ -136,12 +160,9 @@ def read_yaml(yaml_filename: str) -> CommentedMap: # TODO: make this tidier/find a way to notify of which keys should not be stored -def store_yaml(config: CommentedMap, yaml_filename: str, do_not_store_keys: List[Tuple[str, str]] = []) -> None: +def store_yaml(config: CommentedMap, yaml_filename: str) -> None: config_to_save = deepcopy(config) - for key1, key2 in do_not_store_keys: - if key1 in config_to_save and key2 in config_to_save[key1]: - del config_to_save[key1][key2] - + config.pop('urls', None) with open(yaml_filename, "w", encoding="utf-8") as outf: yaml.dump(config_to_save, outf) \ No newline at end of file diff --git a/src/auto_archiver/core/context.py b/src/auto_archiver/core/context.py index 9a21b5c..0db5359 100644 --- a/src/auto_archiver/core/context.py +++ b/src/auto_archiver/core/context.py @@ -53,12 +53,4 @@ class ArchivingContext: if full_reset: ac.keep_on_reset = set() ac.configs = {k: v for k, v in ac.configs.items() if k in ac.keep_on_reset} - # ---- custom getters/setters for widely used context values - - @staticmethod - def set_tmp_dir(tmp_dir: str): - ArchivingContext.get_instance().configs["tmp_dir"] = tmp_dir - - @staticmethod - def get_tmp_dir() -> str: - return ArchivingContext.get_instance().configs.get("tmp_dir") + # ---- custom getters/setters for widely used context values \ No newline at end of file diff --git a/src/auto_archiver/core/extractor.py b/src/auto_archiver/core/extractor.py index ed261eb..b0d80bc 100644 --- a/src/auto_archiver/core/extractor.py +++ b/src/auto_archiver/core/extractor.py @@ -12,7 +12,6 @@ from dataclasses import dataclass import mimetypes import os import mimetypes - import requests from loguru import logger from retrying import retry @@ -71,7 +70,7 @@ class Extractor(BaseModule): to_filename = url.split('/')[-1].split('?')[0] if len(to_filename) > 64: to_filename = to_filename[-64:] - to_filename = os.path.join(ArchivingContext.get_tmp_dir(), to_filename) + to_filename = os.path.join(self.tmp_dir, to_filename) if verbose: logger.debug(f"downloading {url[0:50]=} {to_filename=}") headers = { 'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/81.0.4044.138 Safari/537.36' diff --git a/src/auto_archiver/core/module.py b/src/auto_archiver/core/module.py index 4542b88..501f238 100644 --- a/src/auto_archiver/core/module.py +++ b/src/auto_archiver/core/module.py @@ -7,7 +7,6 @@ from __future__ import annotations from dataclasses import dataclass from typing import List -from abc import ABC import shutil import ast import copy @@ -17,63 +16,12 @@ import os from os.path import join, dirname from loguru import logger import auto_archiver +from .base_module import BaseModule _LAZY_LOADED_MODULES = {} MANIFEST_FILE = "__manifest__.py" -class BaseModule(ABC): - - """ - Base module class. All modules should inherit from this class. - - The exact methods a class implements will depend on the type of module it is, - however all modules have a .setup(config: dict) method to run any setup code - (e.g. logging in to a site, spinning up a browser etc.) - - See BaseModule.MODULE_TYPES for the types of modules you can create, noting that - a subclass can be of multiple types. For example, a module that extracts data from - a website and stores it in a database would be both an 'extractor' and a 'database' module. - - Each module is a python package, and should have a __manifest__.py file in the - same directory as the module file. The __manifest__.py specifies the module information - like name, author, version, dependencies etc. See BaseModule._DEFAULT_MANIFEST for the - default manifest structure. - - """ - - MODULE_TYPES = [ - 'feeder', - 'extractor', - 'enricher', - 'database', - 'storage', - 'formatter' - ] - - _DEFAULT_MANIFEST = { - 'name': '', # the display name of the module - 'author': 'Bellingcat', # creator of the module, leave this as Bellingcat or set your own name! - 'type': [], # the type of the module, can be one or more of BaseModule.MODULE_TYPES - 'requires_setup': True, # whether or not this module requires additional setup such as setting API Keys or installing additional softare - 'description': '', # a description of the module - 'dependencies': {}, # external dependencies, e.g. python packages or binaries, in dictionary format - 'entry_point': '', # the entry point for the module, in the format 'module_name::ClassName'. This can be left blank to use the default entry point of module_name::ModuleName - 'version': '1.0', # the version of the module - 'configs': {} # any configuration options this module has, these will be exposed to the user in the config file or via the command line -} - - config: dict - name: str - - def setup(self, config: dict): - self.config = config - for key, val in config.get(self.name, {}).items(): - setattr(self, key, val) - - def repr(self): - return f"Module<'{self.display_name}' (config: {self.config[self.name]})>" - def setup_paths(paths: list[str]) -> None: """ diff --git a/src/auto_archiver/core/orchestrator.py b/src/auto_archiver/core/orchestrator.py index ba46492..ad11849 100644 --- a/src/auto_archiver/core/orchestrator.py +++ b/src/auto_archiver/core/orchestrator.py @@ -5,12 +5,15 @@ """ from __future__ import annotations -from typing import Generator, Union, List +from typing import Generator, Union, List, Type from urllib.parse import urlparse from ipaddress import ip_address import argparse import os import sys +import json +from tempfile import TemporaryDirectory +import traceback from rich_argparse import RichHelpFormatter @@ -18,17 +21,46 @@ from .context import ArchivingContext from .metadata import Metadata from ..version import __version__ -from .config import read_yaml, store_yaml, to_dot_notation, merge_dicts, EMPTY_CONFIG, DefaultValidatingParser +from .config import yaml, read_yaml, store_yaml, to_dot_notation, merge_dicts, EMPTY_CONFIG, DefaultValidatingParser from .module import available_modules, LazyBaseModule, get_module, setup_paths -from . import validators +from . import validators, Feeder, Extractor, Database, Storage, Formatter, Enricher from .module import BaseModule -import tempfile, traceback from loguru import logger DEFAULT_CONFIG_FILE = "orchestration.yaml" +class JsonParseAction(argparse.Action): + def __call__(self, parser, namespace, values, option_string=None): + try: + setattr(namespace, self.dest, json.loads(values)) + except json.JSONDecodeError as e: + raise argparse.ArgumentTypeError(f"Invalid JSON input for argument '{self.dest}': {e}") + + +class AuthenticationJsonParseAction(JsonParseAction): + def __call__(self, parser, namespace, values, option_string=None): + super().__call__(parser, namespace, values, option_string) + auth_dict = getattr(namespace, self.dest) + if isinstance(auth_dict, str): + # if it's a string + try: + with open(auth_dict, 'r') as f: + try: + auth_dict = json.load(f) + except json.JSONDecodeError: + # maybe it's yaml, try that + auth_dict = yaml.load(f) + except: + pass + + if not isinstance(auth_dict, dict): + raise argparse.ArgumentTypeError("Authentication must be a dictionary of site names and their authentication methods") + for site, auth in auth_dict.items(): + if not isinstance(site, str) or not isinstance(auth, dict): + raise argparse.ArgumentTypeError("Authentication must be a dictionary of site names and their authentication methods") + setattr(namespace, self.dest, auth_dict) class UniqueAppendAction(argparse.Action): def __call__(self, parser, namespace, values, option_string=None): if not hasattr(namespace, self.dest): @@ -38,9 +70,7 @@ class UniqueAppendAction(argparse.Action): getattr(namespace, self.dest).append(value) class ArchivingOrchestrator: - - _do_not_store_keys = [] - + def setup_basic_parser(self): parser = argparse.ArgumentParser( prog="auto-archiver", @@ -52,7 +82,7 @@ class ArchivingOrchestrator: epilog="Check the code at https://github.com/bellingcat/auto-archiver", formatter_class=RichHelpFormatter, ) - parser.add_argument('--help', '-h', action='store_true', dest='help', help='show this help message and exit') + parser.add_argument('--help', '-h', action='store_true', dest='help', help='show a full help message and exit') parser.add_argument('--version', action='version', version=__version__) parser.add_argument('--config', action='store', dest="config_file", help='the filename of the YAML configuration file (defaults to \'config.yaml\')', default=DEFAULT_CONFIG_FILE) parser.add_argument('--mode', action='store', dest='mode', type=str, choices=['simple', 'full'], help='the mode to run the archiver in', default='simple') @@ -80,7 +110,6 @@ class ArchivingOrchestrator: # only load the modules enabled in config # TODO: if some steps are empty (e.g. 'feeders' is empty), should we default to the 'simple' ones? Or only if they are ALL empty? enabled_modules = [] - # first loads the modules from the config file, then from the command line for config in [yaml_config['steps'], basic_config.__dict__]: for module_type in BaseModule.MODULE_TYPES: @@ -120,7 +149,7 @@ class ArchivingOrchestrator: if (self.config != yaml_config and basic_config.store) or not os.path.isfile(basic_config.config_file): logger.info(f"Storing configuration file to {basic_config.config_file}") - store_yaml(self.config, basic_config.config_file, self._do_not_store_keys) + store_yaml(self.config, basic_config.config_file) return self.config @@ -128,18 +157,29 @@ class ArchivingOrchestrator: if not parser: parser = self.parser - parser.add_argument('--feeders', dest='steps.feeders', nargs='+', help='the feeders to use', action=UniqueAppendAction) + + # allow passing URLs directly on the command line + parser.add_argument('urls', nargs='*', default=[], help='URL(s) to archive, either a single URL or a list of urls, should not come from config.yaml') + + parser.add_argument('--feeders', dest='steps.feeders', nargs='+', default=['cli_feeder'], help='the feeders to use', action=UniqueAppendAction) parser.add_argument('--enrichers', dest='steps.enrichers', nargs='+', help='the enrichers to use', action=UniqueAppendAction) parser.add_argument('--extractors', dest='steps.extractors', nargs='+', help='the extractors to use', action=UniqueAppendAction) parser.add_argument('--databases', dest='steps.databases', nargs='+', help='the databases to use', action=UniqueAppendAction) parser.add_argument('--storages', dest='steps.storages', nargs='+', help='the storages to use', action=UniqueAppendAction) parser.add_argument('--formatters', dest='steps.formatters', nargs='+', help='the formatter to use', action=UniqueAppendAction) + parser.add_argument('--authentication', dest='authentication', help='A dictionary of sites and their authentication methods \ + (token, username etc.) that extractors can use to log into \ + a website. If passing this on the command line, use a JSON string. \ + You may also pass a path to a valid JSON/YAML file which will be parsed.',\ + default={}, + action=AuthenticationJsonParseAction) # logging arguments parser.add_argument('--logging.level', action='store', dest='logging.level', choices=['INFO', 'DEBUG', 'ERROR', 'WARNING'], help='the logging level to use', default='INFO') parser.add_argument('--logging.file', action='store', dest='logging.file', help='the logging file to write to', default=None) parser.add_argument('--logging.rotation', action='store', dest='logging.rotation', help='the logging rotation to use', default=None) + def add_module_args(self, modules: list[LazyBaseModule] = None, parser: argparse.ArgumentParser = None) -> None: if not modules: @@ -147,6 +187,7 @@ class ArchivingOrchestrator: module: LazyBaseModule for module in modules: + if not module.configs: # this module has no configs, don't show anything in the help # (TODO: do we want to show something about this module though, like a description?) @@ -155,12 +196,6 @@ class ArchivingOrchestrator: group = parser.add_argument_group(module.display_name or module.name, f"{module.description[:100]}...") for name, kwargs in module.configs.items(): - # TODO: go through all the manifests and make sure we're not breaking anything with removing cli_set - # in most cases it'll mean replacing it with 'type': 'str' or 'type': 'int' or something - do_not_store = kwargs.pop('do_not_store', False) - if do_not_store: - self._do_not_store_keys.append((module.name, name)) - if not kwargs.get('metavar', None): # make a nicer metavar, metavar is what's used in the help, e.g. --cli_feeder.urls [METAVAR] kwargs['metavar'] = name.upper() @@ -208,8 +243,7 @@ class ArchivingOrchestrator: step_items = [] modules_to_load = self.config['steps'][f"{module_type}s"] - assert modules_to_load, f"No {module_type}s were configured. Make sure to set at least one {module_type} \ - in your configuration file or on the command line (using --{module_type}s)" + assert modules_to_load, f"No {module_type}s were configured. Make sure to set at least one {module_type} in your configuration file or on the command line (using --{module_type}s)" def check_steps_ok(): if not len(step_items): @@ -223,12 +257,37 @@ class ArchivingOrchestrator: exit() for module in modules_to_load: + if module == 'cli_feeder': + urls = self.config['urls'] + if not urls: + logger.error("No URLs provided. Please provide at least one URL to archive, or set up a feeder.") + self.basic_parser.print_help() + exit() + # cli_feeder is a pseudo module, it just takes the command line args + def feed(self) -> Generator[Metadata]: + for url in urls: + logger.debug(f"Processing URL: '{url}'") + yield Metadata().set_url(url) + ArchivingContext.set("folder", "cli") + + pseudo_module = type('CLIFeeder', (Feeder,), { + 'name': 'cli_feeder', + 'display_name': 'CLI Feeder', + '__iter__': feed + + })() + + + pseudo_module.__iter__ = feed + step_items.append(pseudo_module) + continue + if module in invalid_modules: continue try: loaded_module: BaseModule = get_module(module, self.config) except (KeyboardInterrupt, Exception) as e: - logger.error(f"Error during setup of archivers: {e}\n{traceback.format_exc()}") + logger.error(f"Error during setup of modules: {e}\n{traceback.format_exc()}") if module_type == 'extractor' and loaded_module.name == module: loaded_module.cleanup() exit() @@ -285,13 +344,18 @@ class ArchivingOrchestrator: def cleanup(self)->None: logger.info("Cleaning up") - for e in self.config['steps']['extractors']: + for e in self.extractors: e.cleanup() def feed(self) -> Generator[Metadata]: - for feeder in self.config['steps']['feeders']: + + url_count = 0 + for feeder in self.feeders: for item in feeder: yield self.feed_item(item) + url_count += 1 + + logger.success(f"Processed {url_count} URL(s)") self.cleanup() def feed_item(self, item: Metadata) -> Metadata: @@ -300,22 +364,33 @@ class ArchivingOrchestrator: - catches keyboard interruptions to do a clean exit - catches any unexpected error, logs it, and does a clean exit """ + tmp_dir: TemporaryDirectory = None try: - ArchivingContext.reset() - with tempfile.TemporaryDirectory(dir="./") as tmp_dir: - ArchivingContext.set_tmp_dir(tmp_dir) - return self.archive(item) + tmp_dir = TemporaryDirectory(dir="./") + # set tmp_dir on all modules + for m in self.all_modules: + m.tmp_dir = tmp_dir.name + return self.archive(item) except KeyboardInterrupt: # catches keyboard interruptions to do a clean exit logger.warning(f"caught interrupt on {item=}") - for d in self.config['steps']['databases']: d.aborted(item) + for d in self.databases: + d.aborted(item) self.cleanup() exit() except Exception as e: logger.error(f'Got unexpected error on item {item}: {e}\n{traceback.format_exc()}') - for d in self.config['steps']['databases']: - if type(e) == AssertionError: d.failed(item, str(e)) - else: d.failed(item, reason="unexpected error") + for d in self.databases: + if type(e) == AssertionError: + d.failed(item, str(e)) + else: + d.failed(item, reason="unexpected error") + finally: + if tmp_dir: + # remove the tmp_dir from all modules + for m in self.all_modules: + m.tmp_dir = None + tmp_dir.cleanup() def archive(self, result: Metadata) -> Union[Metadata, None]: @@ -328,31 +403,38 @@ class ArchivingOrchestrator: 5. Store all downloaded/generated media 6. Call selected Formatter and store formatted if needed """ + original_url = result.get_url().strip() - self.assert_valid_url(original_url) + try: + self.assert_valid_url(original_url) + except AssertionError as e: + logger.error(f"Error archiving URL {original_url}: {e}") + raise e # 1 - sanitize - each archiver is responsible for cleaning/expanding its own URLs url = original_url - for a in self.config["steps"]["extractors"]: url = a.sanitize_url(url) + for a in self.extractors: + url = a.sanitize_url(url) + result.set_url(url) if original_url != url: result.set("original_url", original_url) # 2 - notify start to DBs, propagate already archived if feature enabled in DBs cached_result = None - for d in self.config["steps"]["databases"]: + for d in self.databases: d.started(result) if (local_result := d.fetch(result)): cached_result = (cached_result or Metadata()).merge(local_result) if cached_result: logger.debug("Found previously archived entry") - for d in self.config["steps"]["databases"]: + for d in self.databases: try: d.done(cached_result, cached=True) except Exception as e: logger.error(f"ERROR database {d.name}: {e}: {traceback.format_exc()}") return cached_result # 3 - call extractors until one succeeds - for a in self.config["steps"]["extractors"]: + for a in self.extractors: logger.info(f"Trying extractor {a.name} for {url}") try: result.merge(a.download(result)) @@ -361,7 +443,7 @@ class ArchivingOrchestrator: logger.error(f"ERROR archiver {a.name}: {e}: {traceback.format_exc()}") # 4 - call enrichers to work with archived content - for e in self.config["steps"]["enrichers"]: + for e in self.enrichers: try: e.enrich(result) except Exception as exc: logger.error(f"ERROR enricher {e.name}: {exc}: {traceback.format_exc()}") @@ -370,7 +452,7 @@ class ArchivingOrchestrator: result.store() # 6 - format and store formatted if needed - if final_media := self.config["steps"]["formatters"][0].format(result): + if final_media := self.formatters[0].format(result): final_media.store(url=url, metadata=result) result.set_final_media(final_media) @@ -378,7 +460,7 @@ class ArchivingOrchestrator: result.status = "nothing archived" # signal completion to databases and archivers - for d in self.config["steps"]["databases"]: + for d in self.databases: try: d.done(result) except Exception as e: logger.error(f"ERROR database {d.name}: {e}: {traceback.format_exc()}") @@ -403,4 +485,44 @@ class ArchivingOrchestrator: assert ip.is_global, f"Invalid IP used" assert not ip.is_reserved, f"Invalid IP used" assert not ip.is_link_local, f"Invalid IP used" - assert not ip.is_private, f"Invalid IP used" \ No newline at end of file + assert not ip.is_private, f"Invalid IP used" + + + # Helper Properties + + @property + def feeders(self) -> List[Type[Feeder]]: + return self._get_property('feeders') + + @property + def extractors(self) -> List[Type[Extractor]]: + return self._get_property('extractors') + + @property + def enrichers(self) -> List[Type[Enricher]]: + return self._get_property('enrichers') + + @property + def databases(self) -> List[Type[Database]]: + return self._get_property('databases') + + @property + def storages(self) -> List[Type[Storage]]: + return self._get_property('storages') + + @property + def formatters(self) -> List[Type[Formatter]]: + return self._get_property('formatters') + + @property + def all_modules(self) -> List[Type[BaseModule]]: + return self.feeders + self.extractors + self.enrichers + self.databases + self.storages + self.formatters + + def _get_property(self, prop): + try: + f = self.config['steps'][prop] + if not (isinstance(f[0], BaseModule) or isinstance(f[0], LazyBaseModule)): + raise TypeError + return f + except: + exit("Property called prior to full initialisation") \ No newline at end of file diff --git a/src/auto_archiver/enrichers/screenshot_enricher.py b/src/auto_archiver/enrichers/screenshot_enricher.py new file mode 100644 index 0000000..0d05d92 --- /dev/null +++ b/src/auto_archiver/enrichers/screenshot_enricher.py @@ -0,0 +1,40 @@ +from loguru import logger +import time, os +from selenium.common.exceptions import TimeoutException + + +from auto_archiver.core import Enricher +from ..utils import Webdriver, UrlUtil, random_str +from ..core import Media, Metadata + +class ScreenshotEnricher(Enricher): + name = "screenshot_enricher" + + @staticmethod + def configs() -> dict: + return { + "width": {"default": 1280, "help": "width of the screenshots"}, + "height": {"default": 720, "help": "height of the screenshots"}, + "timeout": {"default": 60, "help": "timeout for taking the screenshot"}, + "sleep_before_screenshot": {"default": 4, "help": "seconds to wait for the pages to load before taking screenshot"}, + "http_proxy": {"default": "", "help": "http proxy to use for the webdriver, eg http://proxy-user:password@proxy-ip:port"}, + } + + def enrich(self, to_enrich: Metadata) -> None: + url = to_enrich.get_url() + if UrlUtil.is_auth_wall(url): + logger.debug(f"[SKIP] SCREENSHOT since url is behind AUTH WALL: {url=}") + return + + logger.debug(f"Enriching screenshot for {url=}") + with Webdriver(self.width, self.height, self.timeout, 'facebook.com' in url, http_proxy=self.http_proxy) as driver: + try: + driver.get(url) + time.sleep(int(self.sleep_before_screenshot)) + screenshot_file = os.path.join(self.tmp_dir, f"screenshot_{random_str(8)}.png") + driver.save_screenshot(screenshot_file) + to_enrich.add_media(Media(filename=screenshot_file), id="screenshot") + except TimeoutException: + logger.info("TimeoutException loading page for screenshot") + except Exception as e: + logger.error(f"Got error while loading webdriver for screenshot enricher: {e}") diff --git a/src/auto_archiver/feeders/csv_feeder.py b/src/auto_archiver/feeders/csv_feeder.py new file mode 100644 index 0000000..e9da518 --- /dev/null +++ b/src/auto_archiver/feeders/csv_feeder.py @@ -0,0 +1,38 @@ +from loguru import logger +import csv + +from . import Feeder +from ..core import Metadata, ArchivingContext +from ..utils import url_or_none + +class CSVFeeder(Feeder): + + @staticmethod + def configs() -> dict: + return { + "files": { + "default": None, + "help": "Path to the input file(s) to read the URLs from, comma separated. \ + Input files should be formatted with one URL per line", + "cli_set": lambda cli_val, cur_val: list(set(cli_val.split(","))) + }, + "column": { + "default": None, + "help": "Column number or name to read the URLs from, 0-indexed", + } + } + + + def __iter__(self) -> Metadata: + url_column = self.column or 0 + for file in self.files: + with open(file, "r") as f: + reader = csv.reader(f) + first_row = next(reader) + if not(url_or_none(first_row[url_column])): + # it's a header row, skip it + for row in reader: + url = row[0] + logger.debug(f"Processing {url}") + yield Metadata().set_url(url) + ArchivingContext.set("folder", "cli") \ No newline at end of file diff --git a/src/auto_archiver/modules/atlos_feeder/atlos_feeder.py b/src/auto_archiver/modules/atlos_feeder/atlos_feeder.py index bbf06f6..8c8f9cb 100644 --- a/src/auto_archiver/modules/atlos_feeder/atlos_feeder.py +++ b/src/auto_archiver/modules/atlos_feeder/atlos_feeder.py @@ -40,5 +40,3 @@ class AtlosFeeder(Feeder): if len(data["results"]) == 0 or cursor is None: break - - logger.success(f"Processed {count} URL(s)") diff --git a/src/auto_archiver/modules/cli_feeder/__init__.py b/src/auto_archiver/modules/cli_feeder/__init__.py deleted file mode 100644 index 9c85787..0000000 --- a/src/auto_archiver/modules/cli_feeder/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from .cli_feeder import CLIFeeder \ No newline at end of file diff --git a/src/auto_archiver/modules/cli_feeder/__manifest__.py b/src/auto_archiver/modules/cli_feeder/__manifest__.py deleted file mode 100644 index cf5c1b7..0000000 --- a/src/auto_archiver/modules/cli_feeder/__manifest__.py +++ /dev/null @@ -1,27 +0,0 @@ -{ - "name": "CLI Feeder", - "type": ["feeder"], - "requires_setup": False, - "dependencies": { - "python": ["loguru"], - }, - 'entry_point': 'cli_feeder::CLIFeeder', - "configs": { - "urls": { - "help": "URL(s) to archive, either a single URL or a list of urls, should not come from config.yaml", - "nargs": "+", - "required": True, - "do_not_store": True, - "metavar": "INPUT URLS", - }, - }, - "description": """ - Processes URLs to archive passed via the command line and feeds them into the archiving pipeline. - - ### Features - - Takes a single URL or a list of URLs provided via the command line. - - Converts each URL into a `Metadata` object and yields it for processing. - - Ensures URLs are processed only if they are explicitly provided. - - """ -} diff --git a/src/auto_archiver/modules/cli_feeder/cli_feeder.py b/src/auto_archiver/modules/cli_feeder/cli_feeder.py deleted file mode 100644 index 62cb659..0000000 --- a/src/auto_archiver/modules/cli_feeder/cli_feeder.py +++ /dev/null @@ -1,15 +0,0 @@ -from loguru import logger - -from auto_archiver.core import Feeder -from auto_archiver.core import Metadata, ArchivingContext - - -class CLIFeeder(Feeder): - - def __iter__(self) -> Metadata: - for url in self.urls: - logger.debug(f"Processing URL: '{url}'") - yield Metadata().set_url(url) - ArchivingContext.set("folder", "cli") - - logger.success(f"Processed {len(self.urls)} URL(s)") diff --git a/src/auto_archiver/modules/csv_feeder/__manifest__.py b/src/auto_archiver/modules/csv_feeder/__manifest__.py index b062ee6..7249395 100644 --- a/src/auto_archiver/modules/csv_feeder/__manifest__.py +++ b/src/auto_archiver/modules/csv_feeder/__manifest__.py @@ -26,7 +26,6 @@ - Supports reading URLs from multiple input files, specified as a comma-separated list. - Allows specifying the column number or name to extract URLs from. - Skips header rows if the first value is not a valid URL. - - Integrates with the `ArchivingContext` to manage URL feeding. ### Setu N - Input files should be formatted with one URL per line. diff --git a/src/auto_archiver/modules/csv_feeder/csv_feeder.py b/src/auto_archiver/modules/csv_feeder/csv_feeder.py index ad0a035..1cd9022 100644 --- a/src/auto_archiver/modules/csv_feeder/csv_feeder.py +++ b/src/auto_archiver/modules/csv_feeder/csv_feeder.py @@ -20,6 +20,4 @@ class CSVFeeder(Feeder): url = row[0] logger.debug(f"Processing {url}") yield Metadata().set_url(url) - ArchivingContext.set("folder", "cli") - - logger.success(f"Processed {len(self.urls)} URL(s)") \ No newline at end of file + ArchivingContext.set("folder", "cli") \ No newline at end of file diff --git a/src/auto_archiver/modules/generic_extractor/generic_extractor.py b/src/auto_archiver/modules/generic_extractor/generic_extractor.py index e643c21..2879c05 100644 --- a/src/auto_archiver/modules/generic_extractor/generic_extractor.py +++ b/src/auto_archiver/modules/generic_extractor/generic_extractor.py @@ -270,7 +270,11 @@ class GenericExtractor(Extractor): logger.debug('Using Facebook cookie') yt_dlp.utils.std_headers['cookie'] = self.facebook_cookie - ydl_options = {'outtmpl': os.path.join(ArchivingContext.get_tmp_dir(), f'%(id)s.%(ext)s'), 'quiet': False, 'noplaylist': not self.allow_playlist , 'writesubtitles': self.subtitles, 'writeautomaticsub': self.subtitles, "live_from_start": self.live_from_start, "proxy": self.proxy, "max_downloads": self.max_downloads, "playlistend": self.max_downloads} + ydl_options = {'outtmpl': os.path.join(self.tmp_dir, f'%(id)s.%(ext)s'), + 'quiet': False, 'noplaylist': not self.allow_playlist , + 'writesubtitles': self.subtitles,'writeautomaticsub': self.subtitles, + "live_from_start": self.live_from_start, "proxy": self.proxy, + "max_downloads": self.max_downloads, "playlistend": self.max_downloads} if item.netloc in ['youtube.com', 'www.youtube.com']: if self.cookies_from_browser: diff --git a/src/auto_archiver/modules/html_formatter/html_formatter.py b/src/auto_archiver/modules/html_formatter/html_formatter.py index bfc2efa..4da82c8 100644 --- a/src/auto_archiver/modules/html_formatter/html_formatter.py +++ b/src/auto_archiver/modules/html_formatter/html_formatter.py @@ -7,7 +7,7 @@ import json import base64 from auto_archiver.version import __version__ -from auto_archiver.core import Metadata, Media, ArchivingContext +from auto_archiver.core import Metadata, Media from auto_archiver.core import Formatter from auto_archiver.modules.hash_enricher import HashEnricher from auto_archiver.utils.misc import random_str @@ -46,7 +46,7 @@ class HtmlFormatter(Formatter): version=__version__ ) - html_path = os.path.join(ArchivingContext.get_tmp_dir(), f"formatted{random_str(24)}.html") + html_path = os.path.join(self.tmp_dir, f"formatted{random_str(24)}.html") with open(html_path, mode="w", encoding="utf-8") as outf: outf.write(content) final_media = Media(filename=html_path, _mimetype="text/html") diff --git a/src/auto_archiver/modules/screenshot_enricher/screenshot_enricher.py b/src/auto_archiver/modules/screenshot_enricher/screenshot_enricher.py index be775ce..8e7639a 100644 --- a/src/auto_archiver/modules/screenshot_enricher/screenshot_enricher.py +++ b/src/auto_archiver/modules/screenshot_enricher/screenshot_enricher.py @@ -7,7 +7,7 @@ from selenium.common.exceptions import TimeoutException from auto_archiver.core import Enricher from auto_archiver.utils import Webdriver, UrlUtil, random_str -from auto_archiver.core import Media, Metadata, ArchivingContext +from auto_archiver.core import Media, Metadata class ScreenshotEnricher(Enricher): @@ -23,11 +23,11 @@ class ScreenshotEnricher(Enricher): try: driver.get(url) time.sleep(int(self.sleep_before_screenshot)) - screenshot_file = os.path.join(ArchivingContext.get_tmp_dir(), f"screenshot_{random_str(8)}.png") + screenshot_file = os.path.join(self.tmp_dir, f"screenshot_{random_str(8)}.png") driver.save_screenshot(screenshot_file) to_enrich.add_media(Media(filename=screenshot_file), id="screenshot") if self.save_to_pdf: - pdf_file = os.path.join(ArchivingContext.get_tmp_dir(), f"pdf_{random_str(8)}.pdf") + pdf_file = os.path.join(self.tmp_dir, f"pdf_{random_str(8)}.pdf") pdf = driver.print_page(driver.print_options) with open(pdf_file, "wb") as f: f.write(base64.b64decode(pdf)) diff --git a/src/auto_archiver/modules/ssl_enricher/ssl_enricher.py b/src/auto_archiver/modules/ssl_enricher/ssl_enricher.py index 52237ee..76784fa 100644 --- a/src/auto_archiver/modules/ssl_enricher/ssl_enricher.py +++ b/src/auto_archiver/modules/ssl_enricher/ssl_enricher.py @@ -23,6 +23,6 @@ class SSLEnricher(Enricher): logger.debug(f"fetching SSL certificate for {domain=} in {url=}") cert = ssl.get_server_certificate((domain, 443)) - cert_fn = os.path.join(ArchivingContext.get_tmp_dir(), f"{slugify(domain)}.pem") + cert_fn = os.path.join(self.tmp_dir, f"{slugify(domain)}.pem") with open(cert_fn, "w") as f: f.write(cert) to_enrich.add_media(Media(filename=cert_fn), id="ssl_certificate") diff --git a/src/auto_archiver/modules/telethon_extractor/telethon_extractor.py b/src/auto_archiver/modules/telethon_extractor/telethon_extractor.py index 8a08954..3e952e8 100644 --- a/src/auto_archiver/modules/telethon_extractor/telethon_extractor.py +++ b/src/auto_archiver/modules/telethon_extractor/telethon_extractor.py @@ -9,7 +9,7 @@ from tqdm import tqdm import re, time, json, os from auto_archiver.core import Extractor -from auto_archiver.core import Metadata, Media, ArchivingContext +from auto_archiver.core import Metadata, Media from auto_archiver.utils import random_str @@ -120,7 +120,7 @@ class TelethonArchiver(Extractor): media_posts = self._get_media_posts_in_group(chat, post) logger.debug(f'got {len(media_posts)=} for {url=}') - tmp_dir = ArchivingContext.get_tmp_dir() + tmp_dir = self.tmp_dir group_id = post.grouped_id if post.grouped_id is not None else post.id title = post.message diff --git a/src/auto_archiver/modules/thumbnail_enricher/thumbnail_enricher.py b/src/auto_archiver/modules/thumbnail_enricher/thumbnail_enricher.py index b27243b..429ba38 100644 --- a/src/auto_archiver/modules/thumbnail_enricher/thumbnail_enricher.py +++ b/src/auto_archiver/modules/thumbnail_enricher/thumbnail_enricher.py @@ -28,7 +28,7 @@ class ThumbnailEnricher(Enricher): logger.debug(f"generating thumbnails for {to_enrich.get_url()}") for m_id, m in enumerate(to_enrich.media[::]): if m.is_video(): - folder = os.path.join(ArchivingContext.get_tmp_dir(), random_str(24)) + folder = os.path.join(self.tmp_dir, random_str(24)) os.makedirs(folder, exist_ok=True) logger.debug(f"generating thumbnails for {m.filename}") duration = m.get("duration") diff --git a/src/auto_archiver/modules/timestamping_enricher/timestamping_enricher.py b/src/auto_archiver/modules/timestamping_enricher/timestamping_enricher.py index a7a0aee..078c1ba 100644 --- a/src/auto_archiver/modules/timestamping_enricher/timestamping_enricher.py +++ b/src/auto_archiver/modules/timestamping_enricher/timestamping_enricher.py @@ -9,9 +9,7 @@ from asn1crypto import pem import certifi from auto_archiver.core import Enricher -from auto_archiver.core import Metadata, ArchivingContext, Media -from auto_archiver.core import Extractor - +from auto_archiver.core import Metadata, Media class TimestampingEnricher(Enricher): """ @@ -33,7 +31,7 @@ class TimestampingEnricher(Enricher): logger.warning(f"No hashes found in {url=}") return - tmp_dir = ArchivingContext.get_tmp_dir() + tmp_dir = self.tmp_dir hashes_fn = os.path.join(tmp_dir, "hashes.txt") data_to_sign = "\n".join(hashes) @@ -93,7 +91,7 @@ class TimestampingEnricher(Enricher): cert_chain = [] for cert in path: - cert_fn = os.path.join(ArchivingContext.get_tmp_dir(), f"{str(cert.serial_number)[:20]}.crt") + cert_fn = os.path.join(self.tmp_dir, f"{str(cert.serial_number)[:20]}.crt") with open(cert_fn, "wb") as f: f.write(cert.dump()) cert_chain.append(Media(filename=cert_fn).set("subject", cert.subject.native["common_name"])) diff --git a/src/auto_archiver/modules/vk_extractor/vk_extractor.py b/src/auto_archiver/modules/vk_extractor/vk_extractor.py index 301fa89..2d09138 100644 --- a/src/auto_archiver/modules/vk_extractor/vk_extractor.py +++ b/src/auto_archiver/modules/vk_extractor/vk_extractor.py @@ -3,7 +3,7 @@ from vk_url_scraper import VkScraper from auto_archiver.utils.misc import dump_payload from auto_archiver.core import Extractor -from auto_archiver.core import Metadata, Media, ArchivingContext +from auto_archiver.core import Metadata, Media class VkExtractor(Extractor): @@ -35,7 +35,7 @@ class VkExtractor(Extractor): result.set_content(dump_payload(vk_scrapes)) - filenames = self.vks.download_media(vk_scrapes, ArchivingContext.get_tmp_dir()) + filenames = self.vks.download_media(vk_scrapes, self.tmp_dir) for filename in filenames: result.add_media(Media(filename)) diff --git a/src/auto_archiver/modules/wacz_enricher/wacz_enricher.py b/src/auto_archiver/modules/wacz_enricher/wacz_enricher.py index 8810b84..3f67b7c 100644 --- a/src/auto_archiver/modules/wacz_enricher/wacz_enricher.py +++ b/src/auto_archiver/modules/wacz_enricher/wacz_enricher.py @@ -5,7 +5,7 @@ from zipfile import ZipFile from loguru import logger from warcio.archiveiterator import ArchiveIterator -from auto_archiver.core import Media, Metadata, ArchivingContext +from auto_archiver.core import Media, Metadata from auto_archiver.core import Extractor, Enricher from auto_archiver.utils import UrlUtil, random_str @@ -51,7 +51,7 @@ class WaczExtractorEnricher(Enricher, Extractor): url = to_enrich.get_url() collection = random_str(8) - browsertrix_home_host = self.browsertrix_home_host or os.path.abspath(ArchivingContext.get_tmp_dir()) + browsertrix_home_host = self.browsertrix_home_host or os.path.abspath(self.tmp_dir) browsertrix_home_container = self.browsertrix_home_container or browsertrix_home_host cmd = [ @@ -154,7 +154,7 @@ class WaczExtractorEnricher(Enricher, Extractor): logger.info(f"WACZ extract_media or extract_screenshot flag is set, extracting media from {wacz_filename=}") # unzipping the .wacz - tmp_dir = ArchivingContext.get_tmp_dir() + tmp_dir = self.tmp_dir unzipped_dir = os.path.join(tmp_dir, "unzipped") with ZipFile(wacz_filename, 'r') as z_obj: z_obj.extractall(path=unzipped_dir) diff --git a/tests/__init__.py b/tests/__init__.py index 3d66aff..31f38cb 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -2,5 +2,4 @@ import tempfile from auto_archiver.core.context import ArchivingContext -ArchivingContext.reset(full_reset=True) -ArchivingContext.set_tmp_dir(tempfile.gettempdir()) \ No newline at end of file +ArchivingContext.reset(full_reset=True) \ No newline at end of file diff --git a/tests/conftest.py b/tests/conftest.py index af0fd6d..3bd382b 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -2,6 +2,7 @@ pytest conftest file, for shared fixtures and configuration """ +from tempfile import TemporaryDirectory from typing import Dict, Tuple import hashlib import pytest @@ -25,8 +26,13 @@ def setup_module(request): m = get_module(module_name, {module_name: config}) + # add the tmp_dir to the module + tmp_dir = TemporaryDirectory() + m.tmp_dir = tmp_dir + def cleanup(): _LAZY_LOADED_MODULES.pop(module_name) + tmp_dir.cleanup() request.addfinalizer(cleanup) return m diff --git a/tests/test_orchestrator.py b/tests/test_orchestrator.py index 03cb521..68417aa 100644 --- a/tests/test_orchestrator.py +++ b/tests/test_orchestrator.py @@ -1,6 +1,6 @@ import pytest import sys -from argparse import ArgumentParser +from argparse import ArgumentParser, ArgumentTypeError from auto_archiver.core.orchestrator import ArchivingOrchestrator from auto_archiver.version import __version__ from auto_archiver.core.config import read_yaml, store_yaml @@ -113,16 +113,23 @@ def test_get_required_values_from_config(orchestrator, test_args, tmp_path): # run the orchestrator orchestrator.run(["--config", tmp_file, "--module_paths", TEST_MODULES]) + assert orchestrator.config is not None - # should run OK, since there are no missing required fields +def test_load_authentication_string(orchestrator, test_args): - # basic_args = basic_parser.parse_known_args(test_args) - # test_yaml = read_yaml(TEST_ORCHESTRATION) - # test_yaml['example_module'] = {'required_field': 'some_value'} + orchestrator.run(test_args + ["--authentication", '{"facebook.com": {"username": "my_username", "password": "my_password"}}']) + assert orchestrator.config['authentication'] == {"facebook.com": {"username": "my_username", "password": "my_password"}} - # # monkey patch the example_module to have a 'configs' setting of 'my_var' with required=True - # # load the module first - # m = get_module_lazy("example_module") +def test_load_authentication_string_concat_site(orchestrator, test_args): + + orchestrator.run(test_args + ["--authentication", '{"x.com,twitter.com": {"api_key": "my_key"}}']) + assert orchestrator.config['authentication'] == {"x.com": {"api_key": "my_key"}, + "twitter.com": {"api_key": "my_key"}} - # orchestrator.setup_complete_parser(basic_args, test_yaml, unused_args=[]) - # assert orchestrator.config is not None \ No newline at end of file +def test_load_invalid_authentication_string(orchestrator, test_args): + with pytest.raises(ArgumentTypeError): + orchestrator.run(test_args + ["--authentication", "{\''invalid_json"]) + +def test_load_authentication_invalid_dict(orchestrator, test_args): + with pytest.raises(ArgumentTypeError): + orchestrator.run(test_args + ["--authentication", "[true, false]"]) \ No newline at end of file From d76063c3f3c287230b4fee06267c64699f6f802a Mon Sep 17 00:00:00 2001 From: Patrick Robertson Date: Thu, 30 Jan 2025 16:46:53 +0100 Subject: [PATCH 10/62] Fix unit tests --- tests/conftest.py | 2 +- tests/test_modules.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index 3bd382b..f909bfb 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -28,7 +28,7 @@ def setup_module(request): # add the tmp_dir to the module tmp_dir = TemporaryDirectory() - m.tmp_dir = tmp_dir + m.tmp_dir = tmp_dir.name def cleanup(): _LAZY_LOADED_MODULES.pop(module_name) diff --git a/tests/test_modules.py b/tests/test_modules.py index a4c0ec8..854edb5 100644 --- a/tests/test_modules.py +++ b/tests/test_modules.py @@ -66,7 +66,7 @@ def test_load_module(example_module): # check that the vlaue is set on the module itself assert loaded_module.csv_file == "db.csv" -@pytest.mark.parametrize("module_name", ["cli_feeder", "local_storage", "generic_extractor", "html_formatter", "csv_db"]) +@pytest.mark.parametrize("module_name", ["local_storage", "generic_extractor", "html_formatter", "csv_db"]) def test_load_modules(module_name): # test that specific modules can be loaded module = get_module_lazy(module_name) @@ -84,7 +84,7 @@ def test_load_modules(module_name): assert loaded_module.name in loaded_module.config.keys() -@pytest.mark.parametrize("module_name", ["cli_feeder", "local_storage", "generic_extractor", "html_formatter", "csv_db"]) +@pytest.mark.parametrize("module_name", ["local_storage", "generic_extractor", "html_formatter", "csv_db"]) def test_lazy_base_module(module_name): lazy_module = get_module_lazy(module_name) From c25d5cae84de7521aea517920463ab97aec8506d Mon Sep 17 00:00:00 2001 From: Patrick Robertson Date: Thu, 30 Jan 2025 17:50:54 +0100 Subject: [PATCH 11/62] Remove ArchivingContext completely Context for a specific url/item is now passed around via the metadata (metadata.set_context('key', 'val') and metadata.get_context('key', default='something') The only other thing that was passed around in ArchivingContext was the storage info, which is already accessible now via self.config --- src/auto_archiver/core/__init__.py | 1 - src/auto_archiver/core/base_module.py | 12 ++-- src/auto_archiver/core/context.py | 56 ------------------- src/auto_archiver/core/extractor.py | 2 +- src/auto_archiver/core/media.py | 10 ++-- src/auto_archiver/core/metadata.py | 15 +++-- src/auto_archiver/core/module.py | 2 +- src/auto_archiver/core/orchestrator.py | 12 ++-- src/auto_archiver/core/storage.py | 23 ++++---- src/auto_archiver/feeders/csv_feeder.py | 5 +- .../modules/csv_feeder/csv_feeder.py | 5 +- .../generic_extractor/generic_extractor.py | 2 +- .../modules/gsheet_db/gsheet_db.py | 5 +- .../modules/gsheet_feeder/gsheet_feeder.py | 12 ++-- .../modules/hash_enricher/hash_enricher.py | 2 +- .../instagram_tbot_extractor.py | 4 +- .../modules/ssl_enricher/ssl_enricher.py | 2 +- .../whisper_enricher/whisper_enricher.py | 6 +- tests/__init__.py | 5 -- 19 files changed, 59 insertions(+), 122 deletions(-) delete mode 100644 src/auto_archiver/core/context.py diff --git a/src/auto_archiver/core/__init__.py b/src/auto_archiver/core/__init__.py index 858bdfd..ae4c41c 100644 --- a/src/auto_archiver/core/__init__.py +++ b/src/auto_archiver/core/__init__.py @@ -4,7 +4,6 @@ from .metadata import Metadata from .media import Media from .module import BaseModule -from .context import ArchivingContext # cannot import ArchivingOrchestrator/Config to avoid circular dep # from .orchestrator import ArchivingOrchestrator diff --git a/src/auto_archiver/core/base_module.py b/src/auto_archiver/core/base_module.py index a9a904f..2c1e8a3 100644 --- a/src/auto_archiver/core/base_module.py +++ b/src/auto_archiver/core/base_module.py @@ -56,6 +56,10 @@ class BaseModule(ABC): # this is set by the orchestrator prior to archiving tmp_dir: TemporaryDirectory = None + @property + def storages(self) -> list: + return self.config.get('storages', []) + def setup(self, config: dict): authentication = config.get('authentication', {}) @@ -75,9 +79,6 @@ class BaseModule(ABC): self.config = config for key, val in config.get(self.name, {}).items(): setattr(self, key, val) - - def repr(self): - return f"Module<'{self.display_name}' (config: {self.config[self.name]})>" def auth_for_site(self, site: str) -> dict: # TODO: think about if/how we can deal with sites that have multiple domains (main one is x.com/twitter.com) @@ -97,4 +98,7 @@ class BaseModule(ABC): did find information for '{key}' which is close, is this what you meant? \ If so, edit your authentication settings to make sure it exactly matches.") - return {} \ No newline at end of file + return {} + + def repr(self): + return f"Module<'{self.display_name}' (config: {self.config[self.name]})>" \ No newline at end of file diff --git a/src/auto_archiver/core/context.py b/src/auto_archiver/core/context.py deleted file mode 100644 index 0db5359..0000000 --- a/src/auto_archiver/core/context.py +++ /dev/null @@ -1,56 +0,0 @@ -""" ArchivingContext provides a global context for managing configurations and temporary data during the archiving process. - -This singleton class allows for: -- Storing and retrieving key-value pairs that are accessible throughout the application lifecycle. -- Marking certain values to persist across resets using `keep_on_reset`. -- Managing temporary directories and other shared data used during the archiving process. - -### Key Features: -- Creates a single global instance. -- Reset functionality allows for clearing configurations, with options for partial or full resets. -- Custom getters and setters for commonly used context values like temporary directories. - -""" - -class ArchivingContext: - """ - Singleton context class for managing global configurations and temporary data. - - ArchivingContext._get_instance() to retrieve it if needed - otherwise just - ArchivingContext.set(key, value) - and - ArchivingContext.get(key, default) - - When reset is called, all values are cleared EXCEPT if they were .set(keep_on_reset=True) - reset(full_reset=True) will recreate everything including the keep_on_reset status - """ - _instance = None - - def __init__(self): - self.configs = {} - self.keep_on_reset = set() - - @staticmethod - def get_instance(): - if ArchivingContext._instance is None: - ArchivingContext._instance = ArchivingContext() - return ArchivingContext._instance - - @staticmethod - def set(key, value, keep_on_reset: bool = False): - ac = ArchivingContext.get_instance() - ac.configs[key] = value - if keep_on_reset: ac.keep_on_reset.add(key) - - @staticmethod - def get(key: str, default=None): - return ArchivingContext.get_instance().configs.get(key, default) - - @staticmethod - def reset(full_reset: bool = False): - ac = ArchivingContext.get_instance() - if full_reset: ac.keep_on_reset = set() - ac.configs = {k: v for k, v in ac.configs.items() if k in ac.keep_on_reset} - - # ---- custom getters/setters for widely used context values \ No newline at end of file diff --git a/src/auto_archiver/core/extractor.py b/src/auto_archiver/core/extractor.py index b0d80bc..98f1370 100644 --- a/src/auto_archiver/core/extractor.py +++ b/src/auto_archiver/core/extractor.py @@ -17,7 +17,7 @@ from loguru import logger from retrying import retry import re -from ..core import Metadata, ArchivingContext, BaseModule +from ..core import Metadata, BaseModule class Extractor(BaseModule): diff --git a/src/auto_archiver/core/media.py b/src/auto_archiver/core/media.py index e5026af..2cb6fc9 100644 --- a/src/auto_archiver/core/media.py +++ b/src/auto_archiver/core/media.py @@ -11,8 +11,6 @@ from dataclasses import dataclass, field from dataclasses_json import dataclass_json, config import mimetypes -from .context import ArchivingContext - from loguru import logger @@ -36,12 +34,11 @@ class Media: _mimetype: str = None # eg: image/jpeg _stored: bool = field(default=False, repr=False, metadata=config(exclude=lambda _: True)) # always exclude - def store(self: Media, override_storages: List = None, url: str = "url-not-available", metadata: Any = None): + def store(self: Media, metadata: Any, url: str = "url-not-available", storages: List[Any] = None) -> None: # 'Any' typing for metadata to avoid circular imports. Stores the media # into the provided/available storages [Storage] repeats the process for # its properties, in case they have inner media themselves for now it # only goes down 1 level but it's easy to make it recursive if needed. - storages = override_storages or ArchivingContext.get("storages") if not len(storages): logger.warning(f"No storages found in local context or provided directly for {self.filename}.") return @@ -66,8 +63,9 @@ class Media: for inner_media in prop_media.all_inner_media(include_self=True): yield inner_media - def is_stored(self) -> bool: - return len(self.urls) > 0 and len(self.urls) == len(ArchivingContext.get("storages")) + def is_stored(self, in_storage) -> bool: + # checks if the media is already stored in the given storage + return len(self.urls) > 0 and any([u for u in self.urls if in_storage.get_cdn_url() in u]) def set(self, key: str, value: Any) -> Media: self.properties[key] = value diff --git a/src/auto_archiver/core/metadata.py b/src/auto_archiver/core/metadata.py index 04683dd..d20ea5e 100644 --- a/src/auto_archiver/core/metadata.py +++ b/src/auto_archiver/core/metadata.py @@ -20,8 +20,6 @@ from dateutil.parser import parse as parse_dt from loguru import logger from .media import Media -from .context import ArchivingContext - @dataclass_json # annotation order matters @dataclass @@ -32,6 +30,7 @@ class Metadata: def __post_init__(self): self.set("_processed_at", datetime.datetime.now(datetime.timezone.utc)) + self._context = {} def merge(self: Metadata, right: Metadata, overwrite_left=True) -> Metadata: """ @@ -57,12 +56,11 @@ class Metadata: return right.merge(self) return self - def store(self: Metadata, override_storages: List = None): + def store(self, storages=[]): # calls .store for all contained media. storages [Storage] self.remove_duplicate_media_by_hash() - storages = override_storages or ArchivingContext.get("storages") for media in self.media: - media.store(override_storages=storages, url=self.get_url(), metadata=self) + media.store(url=self.get_url(), metadata=self, storages=storages) def set(self, key: str, val: Any) -> Metadata: self.metadata[key] = val @@ -206,3 +204,10 @@ class Metadata: if len(r.media) > len(most_complete.media): most_complete = r elif len(r.media) == len(most_complete.media) and len(r.metadata) > len(most_complete.metadata): most_complete = r return most_complete + + def set_context(self, key: str, val: Any) -> Metadata: + self._context[key] = val + return self + + def get_context(self, key: str, default: Any = None) -> Any: + return self._context.get(key, default) \ No newline at end of file diff --git a/src/auto_archiver/core/module.py b/src/auto_archiver/core/module.py index 501f238..dec67e1 100644 --- a/src/auto_archiver/core/module.py +++ b/src/auto_archiver/core/module.py @@ -43,7 +43,6 @@ def setup_paths(paths: list[str]) -> None: # sort based on the length of the path, so that the longest path is last in the list auto_archiver.modules.__path__ = sorted(auto_archiver.modules.__path__, key=len, reverse=True) - def get_module(module_name: str, config: dict) -> BaseModule: """ Gets and sets up a module using the provided config @@ -69,6 +68,7 @@ def get_module_lazy(module_name: str, suppress_warnings: bool = False) -> LazyBa return module def available_modules(with_manifest: bool=False, limit_to_modules: List[str]= [], suppress_warnings: bool = False) -> List[LazyBaseModule]: + # search through all valid 'modules' paths. Default is 'modules' in the current directory # see odoo/modules/module.py -> get_modules diff --git a/src/auto_archiver/core/orchestrator.py b/src/auto_archiver/core/orchestrator.py index ad11849..f046bfe 100644 --- a/src/auto_archiver/core/orchestrator.py +++ b/src/auto_archiver/core/orchestrator.py @@ -17,9 +17,8 @@ import traceback from rich_argparse import RichHelpFormatter -from .context import ArchivingContext -from .metadata import Metadata +from .metadata import Metadata, Media from ..version import __version__ from .config import yaml, read_yaml, store_yaml, to_dot_notation, merge_dicts, EMPTY_CONFIG, DefaultValidatingParser from .module import available_modules, LazyBaseModule, get_module, setup_paths @@ -268,7 +267,6 @@ class ArchivingOrchestrator: for url in urls: logger.debug(f"Processing URL: '{url}'") yield Metadata().set_url(url) - ArchivingContext.set("folder", "cli") pseudo_module = type('CLIFeeder', (Feeder,), { 'name': 'cli_feeder', @@ -297,9 +295,6 @@ class ArchivingOrchestrator: continue if loaded_module: step_items.append(loaded_module) - # TODO temp solution - if module_type == "storage": - ArchivingContext.set("storages", step_items, keep_on_reset=True) check_steps_ok() self.config['steps'][f"{module_type}s"] = step_items @@ -449,11 +444,12 @@ class ArchivingOrchestrator: logger.error(f"ERROR enricher {e.name}: {exc}: {traceback.format_exc()}") # 5 - store all downloaded/generated media - result.store() + result.store(storages=self.storages) # 6 - format and store formatted if needed + final_media: Media if final_media := self.formatters[0].format(result): - final_media.store(url=url, metadata=result) + final_media.store(url=url, metadata=result, storages=self.storages) result.set_final_media(final_media) if result.is_empty(): diff --git a/src/auto_archiver/core/storage.py b/src/auto_archiver/core/storage.py index b40c5cc..9373ff9 100644 --- a/src/auto_archiver/core/storage.py +++ b/src/auto_archiver/core/storage.py @@ -8,16 +8,16 @@ from slugify import slugify from auto_archiver.utils.misc import random_str -from auto_archiver.core import Media, BaseModule, ArchivingContext, Metadata +from auto_archiver.core import Media, BaseModule, Metadata from auto_archiver.modules.hash_enricher.hash_enricher import HashEnricher - +from auto_archiver.core.module import get_module class Storage(BaseModule): - def store(self, media: Media, url: str, metadata: Optional[Metadata]=None) -> None: - if media.is_stored(): + def store(self, media: Media, url: str, metadata: Metadata=None) -> None: + if media.is_stored(in_storage=self): logger.debug(f"{media.key} already stored, skipping") return - self.set_key(media, url) + self.set_key(media, url, metadata) self.upload(media, metadata=metadata) media.add_url(self.get_cdn_url(media)) @@ -32,30 +32,31 @@ class Storage(BaseModule): with open(media.filename, 'rb') as f: return self.uploadf(f, media, **kwargs) - def set_key(self, media: Media, url) -> None: + def set_key(self, media: Media, url, metadata: Metadata) -> None: """takes the media and optionally item info and generates a key""" if media.key is not None and len(media.key) > 0: return - folder = ArchivingContext.get("folder", "") + folder = metadata.folder filename, ext = os.path.splitext(media.filename) # Handle path_generator logic - path_generator = ArchivingContext.get("path_generator", "url") + path_generator = self.config.get("path_generator", "url") if path_generator == "flat": path = "" filename = slugify(filename) # Ensure filename is slugified elif path_generator == "url": path = slugify(url) elif path_generator == "random": - path = ArchivingContext.get("random_path", random_str(24), True) + path = self.config.get("random_path", random_str(24), True) else: raise ValueError(f"Invalid path_generator: {path_generator}") # Handle filename_generator logic - filename_generator = ArchivingContext.get("filename_generator", "random") + filename_generator = self.config.get("filename_generator", "random") if filename_generator == "random": filename = random_str(24) elif filename_generator == "static": - he = HashEnricher({"hash_enricher": {"algorithm": ArchivingContext.get("hash_enricher.algorithm"), "chunksize": 1.6e7}}) + # load the hash_enricher module + he = get_module(HashEnricher, self.config) hd = he.calculate_hash(media.filename) filename = hd[:24] else: diff --git a/src/auto_archiver/feeders/csv_feeder.py b/src/auto_archiver/feeders/csv_feeder.py index e9da518..b1aedb7 100644 --- a/src/auto_archiver/feeders/csv_feeder.py +++ b/src/auto_archiver/feeders/csv_feeder.py @@ -2,7 +2,7 @@ from loguru import logger import csv from . import Feeder -from ..core import Metadata, ArchivingContext +from ..core import Metadata from ..utils import url_or_none class CSVFeeder(Feeder): @@ -34,5 +34,4 @@ class CSVFeeder(Feeder): for row in reader: url = row[0] logger.debug(f"Processing {url}") - yield Metadata().set_url(url) - ArchivingContext.set("folder", "cli") \ No newline at end of file + yield Metadata().set_url(url) \ No newline at end of file diff --git a/src/auto_archiver/modules/csv_feeder/csv_feeder.py b/src/auto_archiver/modules/csv_feeder/csv_feeder.py index 1cd9022..15dfa85 100644 --- a/src/auto_archiver/modules/csv_feeder/csv_feeder.py +++ b/src/auto_archiver/modules/csv_feeder/csv_feeder.py @@ -2,7 +2,7 @@ from loguru import logger import csv from auto_archiver.core import Feeder -from auto_archiver.core import Metadata, ArchivingContext +from auto_archiver.core import Metadata from auto_archiver.utils import url_or_none class CSVFeeder(Feeder): @@ -19,5 +19,4 @@ class CSVFeeder(Feeder): for row in reader: url = row[0] logger.debug(f"Processing {url}") - yield Metadata().set_url(url) - ArchivingContext.set("folder", "cli") \ No newline at end of file + yield Metadata().set_url(url) \ No newline at end of file diff --git a/src/auto_archiver/modules/generic_extractor/generic_extractor.py b/src/auto_archiver/modules/generic_extractor/generic_extractor.py index 2879c05..4838489 100644 --- a/src/auto_archiver/modules/generic_extractor/generic_extractor.py +++ b/src/auto_archiver/modules/generic_extractor/generic_extractor.py @@ -6,7 +6,7 @@ from yt_dlp.extractor.common import InfoExtractor from loguru import logger from auto_archiver.core.extractor import Extractor -from ...core import Metadata, Media, ArchivingContext +from ...core import Metadata, Media class GenericExtractor(Extractor): _dropins = {} diff --git a/src/auto_archiver/modules/gsheet_db/gsheet_db.py b/src/auto_archiver/modules/gsheet_db/gsheet_db.py index e7e8e5c..5e1ed1e 100644 --- a/src/auto_archiver/modules/gsheet_db/gsheet_db.py +++ b/src/auto_archiver/modules/gsheet_db/gsheet_db.py @@ -6,7 +6,7 @@ from urllib.parse import quote from loguru import logger from auto_archiver.core import Database -from auto_archiver.core import Metadata, Media, ArchivingContext +from auto_archiver.core import Metadata, Media from auto_archiver.modules.gsheet_feeder import GWorksheet @@ -93,8 +93,7 @@ class GsheetsDb(Database): logger.debug(f"Unable to update sheet: {e}") def _retrieve_gsheet(self, item: Metadata) -> Tuple[GWorksheet, int]: - # TODO: to make gsheet_db less coupled with gsheet_feeder's "gsheet" parameter, this method could 1st try to fetch "gsheet" from ArchivingContext and, if missing, manage its own singleton - not needed for now - if gsheet := ArchivingContext.get("gsheet"): + if gsheet := item.get_context("gsheet"): gw: GWorksheet = gsheet.get("worksheet") row: int = gsheet.get("row") elif self.sheet_id: diff --git a/src/auto_archiver/modules/gsheet_feeder/gsheet_feeder.py b/src/auto_archiver/modules/gsheet_feeder/gsheet_feeder.py index 235dd63..d129182 100644 --- a/src/auto_archiver/modules/gsheet_feeder/gsheet_feeder.py +++ b/src/auto_archiver/modules/gsheet_feeder/gsheet_feeder.py @@ -15,7 +15,7 @@ from loguru import logger from slugify import slugify from auto_archiver.core import Feeder -from auto_archiver.core import Metadata, ArchivingContext +from auto_archiver.core import Metadata from . import GWorksheet @@ -60,17 +60,15 @@ class GsheetsFeeder(Feeder): # All checks done - archival process starts here m = Metadata().set_url(url) - ArchivingContext.set("gsheet", {"row": row, "worksheet": gw}, keep_on_reset=True) if gw.get_cell_or_default(row, 'folder', "") is None: folder = '' else: folder = slugify(gw.get_cell_or_default(row, 'folder', "").strip()) - if len(folder): - if self.use_sheet_names_in_stored_paths: - ArchivingContext.set("folder", os.path.join(folder, slugify(self.sheet), slugify(wks.title)), True) - else: - ArchivingContext.set("folder", folder, True) + if len(folder) and self.use_sheet_names_in_stored_paths: + folder = os.path.join(folder, slugify(self.sheet), slugify(wks.title)) + m.set_context('folder', folder) + m.set_context('worksheet', {"row": row, "worksheet": gw}) yield m logger.success(f'Finished worksheet {wks.title}') diff --git a/src/auto_archiver/modules/hash_enricher/hash_enricher.py b/src/auto_archiver/modules/hash_enricher/hash_enricher.py index 94b5dce..58c6abe 100644 --- a/src/auto_archiver/modules/hash_enricher/hash_enricher.py +++ b/src/auto_archiver/modules/hash_enricher/hash_enricher.py @@ -11,7 +11,7 @@ import hashlib from loguru import logger from auto_archiver.core import Enricher -from auto_archiver.core import Metadata, ArchivingContext +from auto_archiver.core import Metadata class HashEnricher(Enricher): diff --git a/src/auto_archiver/modules/instagram_tbot_extractor/instagram_tbot_extractor.py b/src/auto_archiver/modules/instagram_tbot_extractor/instagram_tbot_extractor.py index 791b9c0..5b49484 100644 --- a/src/auto_archiver/modules/instagram_tbot_extractor/instagram_tbot_extractor.py +++ b/src/auto_archiver/modules/instagram_tbot_extractor/instagram_tbot_extractor.py @@ -16,7 +16,7 @@ from loguru import logger from telethon.sync import TelegramClient from auto_archiver.core import Extractor -from auto_archiver.core import Metadata, Media, ArchivingContext +from auto_archiver.core import Metadata, Media from auto_archiver.utils import random_str @@ -61,7 +61,7 @@ class InstagramTbotExtractor(Extractor): if not "instagram.com" in url: return False result = Metadata() - tmp_dir = ArchivingContext.get_tmp_dir() + tmp_dir = self.tmp_dir with self.client.start(): chat = self.client.get_entity("instagram_load_bot") since_id = self.client.send_message(entity=chat, message=url).id diff --git a/src/auto_archiver/modules/ssl_enricher/ssl_enricher.py b/src/auto_archiver/modules/ssl_enricher/ssl_enricher.py index 76784fa..b429163 100644 --- a/src/auto_archiver/modules/ssl_enricher/ssl_enricher.py +++ b/src/auto_archiver/modules/ssl_enricher/ssl_enricher.py @@ -4,7 +4,7 @@ from urllib.parse import urlparse from loguru import logger from auto_archiver.core import Enricher -from auto_archiver.core import Metadata, ArchivingContext, Media +from auto_archiver.core import Metadata, Media class SSLEnricher(Enricher): diff --git a/src/auto_archiver/modules/whisper_enricher/whisper_enricher.py b/src/auto_archiver/modules/whisper_enricher/whisper_enricher.py index b8fe634..8ca2131 100644 --- a/src/auto_archiver/modules/whisper_enricher/whisper_enricher.py +++ b/src/auto_archiver/modules/whisper_enricher/whisper_enricher.py @@ -3,7 +3,7 @@ import requests, time from loguru import logger from auto_archiver.core import Enricher -from auto_archiver.core import Metadata, Media, ArchivingContext +from auto_archiver.core import Metadata, Media from auto_archiver.modules.s3_storage import S3Storage from auto_archiver.core.module import get_module @@ -25,7 +25,7 @@ class WhisperEnricher(Enricher): job_results = {} for i, m in enumerate(to_enrich.media): if m.is_video() or m.is_audio(): - m.store(url=url, metadata=to_enrich) + m.store(url=url, metadata=to_enrich, storages=self.storages) try: job_id = self.submit_job(m) job_results[job_id] = False @@ -110,7 +110,7 @@ class WhisperEnricher(Enricher): def _get_s3_storage(self) -> S3Storage: try: - return next(s for s in ArchivingContext.get("storages") if s.__class__ == S3Storage) + return next(s for s in self.storages if s.__class__ == S3Storage) except: logger.warning("No S3Storage instance found in storages") return diff --git a/tests/__init__.py b/tests/__init__.py index 31f38cb..e69de29 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -1,5 +0,0 @@ -import tempfile - -from auto_archiver.core.context import ArchivingContext - -ArchivingContext.reset(full_reset=True) \ No newline at end of file From 9a8c94b641581b61e0696d24b5e4ed7bdb778e32 Mon Sep 17 00:00:00 2001 From: Patrick Robertson Date: Mon, 3 Feb 2025 16:02:17 +0100 Subject: [PATCH 12/62] Fix getting/setting folder context for metadata --- src/auto_archiver/core/storage.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/auto_archiver/core/storage.py b/src/auto_archiver/core/storage.py index 9373ff9..9f355f6 100644 --- a/src/auto_archiver/core/storage.py +++ b/src/auto_archiver/core/storage.py @@ -35,7 +35,7 @@ class Storage(BaseModule): def set_key(self, media: Media, url, metadata: Metadata) -> None: """takes the media and optionally item info and generates a key""" if media.key is not None and len(media.key) > 0: return - folder = metadata.folder + folder = metadata.get_context('folder', '') filename, ext = os.path.splitext(media.filename) # Handle path_generator logic From 9c9e9b370e73675da6ec9028b4a798b2ba81cf53 Mon Sep 17 00:00:00 2001 From: Patrick Robertson Date: Mon, 3 Feb 2025 16:02:38 +0100 Subject: [PATCH 13/62] Remove lingering reference to ArchivingContext --- .../modules/thumbnail_enricher/thumbnail_enricher.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/auto_archiver/modules/thumbnail_enricher/thumbnail_enricher.py b/src/auto_archiver/modules/thumbnail_enricher/thumbnail_enricher.py index 429ba38..e0ac937 100644 --- a/src/auto_archiver/modules/thumbnail_enricher/thumbnail_enricher.py +++ b/src/auto_archiver/modules/thumbnail_enricher/thumbnail_enricher.py @@ -10,7 +10,7 @@ import ffmpeg, os from loguru import logger from auto_archiver.core import Enricher -from auto_archiver.core import Media, Metadata, ArchivingContext +from auto_archiver.core import Media, Metadata from auto_archiver.utils.misc import random_str From 7a2be5a0da13713980ced0a34aed37cc0b891979 Mon Sep 17 00:00:00 2001 From: Patrick Robertson Date: Mon, 3 Feb 2025 16:03:07 +0100 Subject: [PATCH 14/62] Add cookie extraction to 'authentication' options, get generic_extractor working using this info --- src/auto_archiver/core/base_module.py | 57 +++++++++++++++---- src/auto_archiver/core/orchestrator.py | 3 +- .../generic_extractor/generic_extractor.py | 29 ++++++---- 3 files changed, 66 insertions(+), 23 deletions(-) diff --git a/src/auto_archiver/core/base_module.py b/src/auto_archiver/core/base_module.py index 2c1e8a3..d23643c 100644 --- a/src/auto_archiver/core/base_module.py +++ b/src/auto_archiver/core/base_module.py @@ -1,5 +1,4 @@ - from urllib.parse import urlparse from typing import Mapping, Any from abc import ABC @@ -80,25 +79,63 @@ class BaseModule(ABC): for key, val in config.get(self.name, {}).items(): setattr(self, key, val) - def auth_for_site(self, site: str) -> dict: + def auth_for_site(self, site: str, extract_cookies=True) -> Mapping[str, Any]: + """ + Returns the authentication information for a given site. This is used to authenticate + with a site before extracting data. The site should be the domain of the site, e.g. 'twitter.com' + + extract_cookies: bool - whether or not to extract cookies from the given browser and return the + cookie jar (disabling can speed up) processing if you don't actually need the cookies jar + + Currently, the dict can have keys of the following types: + - username: str - the username to use for login + - password: str - the password to use for login + - api_key: str - the API key to use for login + - api_secret: str - the API secret to use for login + - cookie: str - a cookie string to use for login (specific to this site) + - cookies_jar: YoutubeDLCookieJar | http.cookiejar.MozillaCookieJar - a cookie jar compatible with requests (e.g. `requests.get(cookies=cookie_jar)`) + """ # TODO: think about if/how we can deal with sites that have multiple domains (main one is x.com/twitter.com) - # for now, just hard code those. + # for now the user must enter them both, like "x.com,twitter.com" in their config. Maybe we just hard-code? # SECURITY: parse the domain using urllib site = urlparse(site).netloc # add the 'www' version of the site to the list of sites to check + authdict = {} + + for to_try in [site, f"www.{site}"]: if to_try in self.authentication: - return self.authentication[to_try] + authdict.update(self.authentication[to_try]) + break # do a fuzzy string match just to print a warning - don't use it since it's insecure - for key in self.authentication.keys(): - if key in site or site in key: - logger.warning(f"Could not find exact authentication information for site '{site}'. \ - did find information for '{key}' which is close, is this what you meant? \ - If so, edit your authentication settings to make sure it exactly matches.") + if not authdict: + for key in self.authentication.keys(): + if key in site or site in key: + logger.debug(f"Could not find exact authentication information for site '{site}'. \ + did find information for '{key}' which is close, is this what you meant? \ + If so, edit your authentication settings to make sure it exactly matches.") - return {} + + def get_ytdlp_cookiejar(args): + import yt_dlp + from yt_dlp import parse_options + + # parse_options returns a named tuple as follows, we only need the ydl_options part + # collections.namedtuple('ParsedOptions', ('parser', 'options', 'urls', 'ydl_opts')) + ytdlp_opts = getattr(parse_options(args), 'ydl_opts') + return yt_dlp.YoutubeDL(ytdlp_opts).cookiejar + + # get the cookies jar, prefer the browser cookies than the file + if 'cookies_from_browser' in self.authentication: + authdict['cookies_from_browser'] = self.authentication['cookies_from_browser'] + authdict['cookies_jar'] = get_ytdlp_cookiejar(['--cookies-from-browser', self.authentication['cookies_from_browser']]) + elif 'cookies_file' in self.authentication: + authdict['cookies_file'] = self.authentication['cookies_file'] + authdict['cookies_jar'] = get_ytdlp_cookiejar(['--cookies', self.authentication['cookies_file']]) + + return authdict def repr(self): return f"Module<'{self.display_name}' (config: {self.config[self.name]})>" \ No newline at end of file diff --git a/src/auto_archiver/core/orchestrator.py b/src/auto_archiver/core/orchestrator.py index f046bfe..85b3d61 100644 --- a/src/auto_archiver/core/orchestrator.py +++ b/src/auto_archiver/core/orchestrator.py @@ -259,8 +259,7 @@ class ArchivingOrchestrator: if module == 'cli_feeder': urls = self.config['urls'] if not urls: - logger.error("No URLs provided. Please provide at least one URL to archive, or set up a feeder.") - self.basic_parser.print_help() + logger.error("No URLs provided. Please provide at least one URL to archive, or set up a feeder. Use --help for more information.") exit() # cli_feeder is a pseudo module, it just takes the command line args def feed(self) -> Generator[Metadata]: diff --git a/src/auto_archiver/modules/generic_extractor/generic_extractor.py b/src/auto_archiver/modules/generic_extractor/generic_extractor.py index 4838489..bc884a6 100644 --- a/src/auto_archiver/modules/generic_extractor/generic_extractor.py +++ b/src/auto_archiver/modules/generic_extractor/generic_extractor.py @@ -266,23 +266,30 @@ class GenericExtractor(Extractor): def download(self, item: Metadata) -> Metadata: url = item.get_url() - if item.netloc in ['facebook.com', 'www.facebook.com'] and self.facebook_cookie: - logger.debug('Using Facebook cookie') - yt_dlp.utils.std_headers['cookie'] = self.facebook_cookie ydl_options = {'outtmpl': os.path.join(self.tmp_dir, f'%(id)s.%(ext)s'), 'quiet': False, 'noplaylist': not self.allow_playlist , 'writesubtitles': self.subtitles,'writeautomaticsub': self.subtitles, "live_from_start": self.live_from_start, "proxy": self.proxy, "max_downloads": self.max_downloads, "playlistend": self.max_downloads} - - if item.netloc in ['youtube.com', 'www.youtube.com']: - if self.cookies_from_browser: - logger.debug(f'Extracting cookies from browser {self.cookies_from_browser} for Youtube') - ydl_options['cookiesfrombrowser'] = (self.cookies_from_browser,) - elif self.cookie_file: - logger.debug(f'Using cookies from file {self.cookie_file}') - ydl_options['cookiefile'] = self.cookie_file + + # set up auth + auth = self.auth_for_site(url) + # order of importance: username/pasword -> api_key -> cookie -> cookie_from_browser -> cookies_file + if auth: + if 'username' in auth and 'password' in auth: + logger.debug(f'Using provided auth username and password for {url}') + ydl_options['username'] = auth['username'] + ydl_options['password'] = auth['password'] + elif 'cookie' in auth: + logger.debug(f'Using provided auth cookie for {url}') + yt_dlp.utils.std_headers['cookie'] = auth['cookie'] + elif 'cookie_from_browser' in auth: + logger.debug(f'Using extracted cookies from browser {self.cookies_from_browser} for {url}') + ydl_options['cookiesfrombrowser'] = auth['cookies_from_browser'] + elif 'cookies_file' in auth: + logger.debug(f'Using cookies from file {self.cookie_file} for {url}') + ydl_options['cookiesfile'] = auth['cookies_file'] ydl = yt_dlp.YoutubeDL(ydl_options) # allsubtitles and subtitleslangs not working as expected, so default lang is always "en" From 7ec328ab409064e4e81a443f84565195b1848655 Mon Sep 17 00:00:00 2001 From: Patrick Robertson Date: Mon, 3 Feb 2025 16:04:36 +0100 Subject: [PATCH 15/62] Remove cookie options from generic_extractor - it now uses 'authentication' global settings :D --- .../modules/generic_extractor/__manifest__.py | 14 +------------- 1 file changed, 1 insertion(+), 13 deletions(-) diff --git a/src/auto_archiver/modules/generic_extractor/__manifest__.py b/src/auto_archiver/modules/generic_extractor/__manifest__.py index d5f363f..caa3ae1 100644 --- a/src/auto_archiver/modules/generic_extractor/__manifest__.py +++ b/src/auto_archiver/modules/generic_extractor/__manifest__.py @@ -20,6 +20,7 @@ the broader archiving framework. - Retrieves metadata like titles, descriptions, upload dates, and durations. - Downloads subtitles and comments when enabled. - Configurable options for handling live streams, proxies, and more. +- Supports authentication of websites using the 'authentication' settings from your orchestration. ### Dropins - For websites supported by `yt-dlp` that also contain posts in addition to videos @@ -29,10 +30,6 @@ custom dropins can be created to handle additional websites and passed to the ar via the command line using the `--dropins` option (TODO!). """, "configs": { - "facebook_cookie": { - "default": None, - "help": "optional facebook cookie to have more access to content, from browser, looks like 'cookie: datr= xxxx'", - }, "subtitles": {"default": True, "help": "download subtitles if available", "type": "bool"}, "comments": { "default": False, @@ -67,14 +64,5 @@ via the command line using the `--dropins` option (TODO!). "default": "inf", "help": "Use to limit the number of videos to download when a channel or long page is being extracted. 'inf' means no limit.", }, - "cookies_from_browser": { - "default": None, - "type": "str", - "help": "optional browser for ytdl to extract cookies from, can be one of: brave, chrome, chromium, edge, firefox, opera, safari, vivaldi, whale", - }, - "cookie_file": { - "default": None, - "help": "optional cookie file to use for Youtube, see instructions here on how to export from your browser: https://github.com/yt-dlp/yt-dlp/wiki/FAQ#how-do-i-pass-cookies-to-yt-dlp", - }, }, } From c574b694ed0db50792b0719504486252848adfdd Mon Sep 17 00:00:00 2001 From: Patrick Robertson Date: Mon, 3 Feb 2025 17:25:59 +0100 Subject: [PATCH 16/62] Set up screenshot enricher to use authentication/cookies --- src/auto_archiver/core/base_module.py | 15 +- src/auto_archiver/core/orchestrator.py | 2 +- .../enrichers/screenshot_enricher.py | 2 +- .../generic_extractor/generic_extractor.py | 2 +- .../modules/generic_extractor/twitter.py | 2 +- .../screenshot_enricher.py | 6 +- .../modules/wacz_enricher/wacz_enricher.py | 2 +- .../wayback_extractor_enricher.py | 2 +- src/auto_archiver/utils/__init__.py | 1 - src/auto_archiver/utils/url.py | 129 +++++++++--------- src/auto_archiver/utils/webdriver.py | 86 +++++++++--- 11 files changed, 153 insertions(+), 96 deletions(-) diff --git a/src/auto_archiver/core/base_module.py b/src/auto_archiver/core/base_module.py index d23643c..fcfe9ea 100644 --- a/src/auto_archiver/core/base_module.py +++ b/src/auto_archiver/core/base_module.py @@ -4,6 +4,7 @@ from typing import Mapping, Any from abc import ABC from copy import deepcopy, copy from tempfile import TemporaryDirectory +from auto_archiver.utils import url as UrlUtil from loguru import logger @@ -78,7 +79,7 @@ class BaseModule(ABC): self.config = config for key, val in config.get(self.name, {}).items(): setattr(self, key, val) - + def auth_for_site(self, site: str, extract_cookies=True) -> Mapping[str, Any]: """ Returns the authentication information for a given site. This is used to authenticate @@ -98,8 +99,7 @@ class BaseModule(ABC): # TODO: think about if/how we can deal with sites that have multiple domains (main one is x.com/twitter.com) # for now the user must enter them both, like "x.com,twitter.com" in their config. Maybe we just hard-code? - # SECURITY: parse the domain using urllib - site = urlparse(site).netloc + site = UrlUtil.domain_for_url(site) # add the 'www' version of the site to the list of sites to check authdict = {} @@ -116,12 +116,11 @@ class BaseModule(ABC): logger.debug(f"Could not find exact authentication information for site '{site}'. \ did find information for '{key}' which is close, is this what you meant? \ If so, edit your authentication settings to make sure it exactly matches.") - def get_ytdlp_cookiejar(args): import yt_dlp from yt_dlp import parse_options - + logger.debug(f"Extracting cookies from settings: {args[1]}") # parse_options returns a named tuple as follows, we only need the ydl_options part # collections.namedtuple('ParsedOptions', ('parser', 'options', 'urls', 'ydl_opts')) ytdlp_opts = getattr(parse_options(args), 'ydl_opts') @@ -130,10 +129,12 @@ class BaseModule(ABC): # get the cookies jar, prefer the browser cookies than the file if 'cookies_from_browser' in self.authentication: authdict['cookies_from_browser'] = self.authentication['cookies_from_browser'] - authdict['cookies_jar'] = get_ytdlp_cookiejar(['--cookies-from-browser', self.authentication['cookies_from_browser']]) + if extract_cookies: + authdict['cookies_jar'] = get_ytdlp_cookiejar(['--cookies-from-browser', self.authentication['cookies_from_browser']]) elif 'cookies_file' in self.authentication: authdict['cookies_file'] = self.authentication['cookies_file'] - authdict['cookies_jar'] = get_ytdlp_cookiejar(['--cookies', self.authentication['cookies_file']]) + if extract_cookies: + authdict['cookies_jar'] = get_ytdlp_cookiejar(['--cookies', self.authentication['cookies_file']]) return authdict diff --git a/src/auto_archiver/core/orchestrator.py b/src/auto_archiver/core/orchestrator.py index 85b3d61..dbc8a33 100644 --- a/src/auto_archiver/core/orchestrator.py +++ b/src/auto_archiver/core/orchestrator.py @@ -174,7 +174,7 @@ class ArchivingOrchestrator: default={}, action=AuthenticationJsonParseAction) # logging arguments - parser.add_argument('--logging.level', action='store', dest='logging.level', choices=['INFO', 'DEBUG', 'ERROR', 'WARNING'], help='the logging level to use', default='INFO') + parser.add_argument('--logging.level', action='store', dest='logging.level', choices=['INFO', 'DEBUG', 'ERROR', 'WARNING'], help='the logging level to use', default='INFO', type=str.upper) parser.add_argument('--logging.file', action='store', dest='logging.file', help='the logging file to write to', default=None) parser.add_argument('--logging.rotation', action='store', dest='logging.rotation', help='the logging rotation to use', default=None) diff --git a/src/auto_archiver/enrichers/screenshot_enricher.py b/src/auto_archiver/enrichers/screenshot_enricher.py index 0d05d92..abb1e16 100644 --- a/src/auto_archiver/enrichers/screenshot_enricher.py +++ b/src/auto_archiver/enrichers/screenshot_enricher.py @@ -4,7 +4,7 @@ from selenium.common.exceptions import TimeoutException from auto_archiver.core import Enricher -from ..utils import Webdriver, UrlUtil, random_str +from ..utils import Webdriver, url as UrlUtil, random_str from ..core import Media, Metadata class ScreenshotEnricher(Enricher): diff --git a/src/auto_archiver/modules/generic_extractor/generic_extractor.py b/src/auto_archiver/modules/generic_extractor/generic_extractor.py index bc884a6..d1b1fb6 100644 --- a/src/auto_archiver/modules/generic_extractor/generic_extractor.py +++ b/src/auto_archiver/modules/generic_extractor/generic_extractor.py @@ -274,7 +274,7 @@ class GenericExtractor(Extractor): "max_downloads": self.max_downloads, "playlistend": self.max_downloads} # set up auth - auth = self.auth_for_site(url) + auth = self.auth_for_site(url, extract_cookies=False) # order of importance: username/pasword -> api_key -> cookie -> cookie_from_browser -> cookies_file if auth: if 'username' in auth and 'password' in auth: diff --git a/src/auto_archiver/modules/generic_extractor/twitter.py b/src/auto_archiver/modules/generic_extractor/twitter.py index 83c1f4f..3faed6b 100644 --- a/src/auto_archiver/modules/generic_extractor/twitter.py +++ b/src/auto_archiver/modules/generic_extractor/twitter.py @@ -5,7 +5,7 @@ from loguru import logger from slugify import slugify from auto_archiver.core.metadata import Metadata, Media -from auto_archiver.utils import UrlUtil +from auto_archiver.utils import url as UrlUtil from auto_archiver.core.extractor import Extractor from .dropin import GenericDropin, InfoExtractor diff --git a/src/auto_archiver/modules/screenshot_enricher/screenshot_enricher.py b/src/auto_archiver/modules/screenshot_enricher/screenshot_enricher.py index 8e7639a..e1da99d 100644 --- a/src/auto_archiver/modules/screenshot_enricher/screenshot_enricher.py +++ b/src/auto_archiver/modules/screenshot_enricher/screenshot_enricher.py @@ -6,7 +6,7 @@ from selenium.common.exceptions import TimeoutException from auto_archiver.core import Enricher -from auto_archiver.utils import Webdriver, UrlUtil, random_str +from auto_archiver.utils import Webdriver, url as UrlUtil, random_str from auto_archiver.core import Media, Metadata class ScreenshotEnricher(Enricher): @@ -19,7 +19,9 @@ class ScreenshotEnricher(Enricher): return logger.debug(f"Enriching screenshot for {url=}") - with Webdriver(self.width, self.height, self.timeout, 'facebook.com' in url, http_proxy=self.http_proxy, print_options=self.print_options) as driver: + auth = self.auth_for_site(url) + with Webdriver(self.width, self.height, self.timeout, facebook_accept_cookies='facebook.com' in url, + http_proxy=self.http_proxy, print_options=self.print_options, auth=auth) as driver: try: driver.get(url) time.sleep(int(self.sleep_before_screenshot)) diff --git a/src/auto_archiver/modules/wacz_enricher/wacz_enricher.py b/src/auto_archiver/modules/wacz_enricher/wacz_enricher.py index 3f67b7c..1586b75 100644 --- a/src/auto_archiver/modules/wacz_enricher/wacz_enricher.py +++ b/src/auto_archiver/modules/wacz_enricher/wacz_enricher.py @@ -7,7 +7,7 @@ from warcio.archiveiterator import ArchiveIterator from auto_archiver.core import Media, Metadata from auto_archiver.core import Extractor, Enricher -from auto_archiver.utils import UrlUtil, random_str +from auto_archiver.utils import url as UrlUtil, random_str class WaczExtractorEnricher(Enricher, Extractor): diff --git a/src/auto_archiver/modules/wayback_extractor_enricher/wayback_extractor_enricher.py b/src/auto_archiver/modules/wayback_extractor_enricher/wayback_extractor_enricher.py index 0e25440..1763b12 100644 --- a/src/auto_archiver/modules/wayback_extractor_enricher/wayback_extractor_enricher.py +++ b/src/auto_archiver/modules/wayback_extractor_enricher/wayback_extractor_enricher.py @@ -3,7 +3,7 @@ from loguru import logger import time, requests from auto_archiver.core import Extractor, Enricher -from auto_archiver.utils import UrlUtil +from auto_archiver.utils import url as UrlUtil from auto_archiver.core import Metadata class WaybackExtractorEnricher(Enricher, Extractor): diff --git a/src/auto_archiver/utils/__init__.py b/src/auto_archiver/utils/__init__.py index d2063d0..ed2d3bb 100644 --- a/src/auto_archiver/utils/__init__.py +++ b/src/auto_archiver/utils/__init__.py @@ -2,7 +2,6 @@ # we need to explicitly expose the available imports here from .misc import * from .webdriver import Webdriver -from .url import UrlUtil from .atlos import get_atlos_config_options # handy utils from ytdlp diff --git a/src/auto_archiver/utils/url.py b/src/auto_archiver/utils/url.py index 3b67514..40884da 100644 --- a/src/auto_archiver/utils/url.py +++ b/src/auto_archiver/utils/url.py @@ -1,83 +1,84 @@ import re from urllib.parse import urlparse, urlunparse -class UrlUtil: - AUTHWALL_URLS = [ - re.compile(r"https:\/\/t\.me(\/c)\/(.+)\/(\d+)"), # telegram private channels - re.compile(r"https:\/\/www\.instagram\.com"), # instagram - ] +AUTHWALL_URLS = [ + re.compile(r"https:\/\/t\.me(\/c)\/(.+)\/(\d+)"), # telegram private channels + re.compile(r"https:\/\/www\.instagram\.com"), # instagram +] - @staticmethod - def clean(url: str) -> str: return url +def domain_for_url(url: str) -> str: + """ + SECURITY: parse the domain using urllib to avoid any potential security issues + """ + return urlparse(url).netloc - @staticmethod - def is_auth_wall(url: str) -> bool: - """ - checks if URL is behind an authentication wall meaning steps like wayback, wacz, ... may not work - """ - for regex in UrlUtil.AUTHWALL_URLS: - if regex.match(url): - return True +def clean(url: str) -> str: + return url - return False +def is_auth_wall(url: str) -> bool: + """ + checks if URL is behind an authentication wall meaning steps like wayback, wacz, ... may not work + """ + for regex in AUTHWALL_URLS: + if regex.match(url): + return True - @staticmethod - def remove_get_parameters(url: str) -> str: - # http://example.com/file.mp4?t=1 -> http://example.com/file.mp4 - # useful for mimetypes to work - parsed_url = urlparse(url) - new_url = urlunparse(parsed_url._replace(query='')) - return new_url + return False - @staticmethod - def is_relevant_url(url: str) -> bool: - """ - Detect if a detected media URL is recurring and therefore irrelevant to a specific archive. Useful, for example, for the enumeration of the media files in WARC files which include profile pictures, favicons, etc. - """ - clean_url = UrlUtil.remove_get_parameters(url) +def remove_get_parameters(url: str) -> str: + # http://example.com/file.mp4?t=1 -> http://example.com/file.mp4 + # useful for mimetypes to work + parsed_url = urlparse(url) + new_url = urlunparse(parsed_url._replace(query='')) + return new_url - # favicons - if "favicon" in url: return False - # ifnore icons - if clean_url.endswith(".ico"): return False - # ignore SVGs - if UrlUtil.remove_get_parameters(url).endswith(".svg"): return False +def is_relevant_url(url: str) -> bool: + """ + Detect if a detected media URL is recurring and therefore irrelevant to a specific archive. Useful, for example, for the enumeration of the media files in WARC files which include profile pictures, favicons, etc. + """ + clean_url = remove_get_parameters(url) - # twitter profile pictures - if "twimg.com/profile_images" in url: return False - if "twimg.com" in url and "/default_profile_images" in url: return False + # favicons + if "favicon" in url: return False + # ifnore icons + if clean_url.endswith(".ico"): return False + # ignore SVGs + if remove_get_parameters(url).endswith(".svg"): return False - # instagram profile pictures - if "https://scontent.cdninstagram.com/" in url and "150x150" in url: return False - # instagram recurring images - if "https://static.cdninstagram.com/rsrc.php/" in url: return False + # twitter profile pictures + if "twimg.com/profile_images" in url: return False + if "twimg.com" in url and "/default_profile_images" in url: return False - # telegram - if "https://telegram.org/img/emoji/" in url: return False + # instagram profile pictures + if "https://scontent.cdninstagram.com/" in url and "150x150" in url: return False + # instagram recurring images + if "https://static.cdninstagram.com/rsrc.php/" in url: return False - # youtube - if "https://www.youtube.com/s/gaming/emoji/" in url: return False - if "https://yt3.ggpht.com" in url and "default-user=" in url: return False - if "https://www.youtube.com/s/search/audio/" in url: return False + # telegram + if "https://telegram.org/img/emoji/" in url: return False - # ok - if " https://ok.ru/res/i/" in url: return False + # youtube + if "https://www.youtube.com/s/gaming/emoji/" in url: return False + if "https://yt3.ggpht.com" in url and "default-user=" in url: return False + if "https://www.youtube.com/s/search/audio/" in url: return False - # vk - if "https://vk.com/emoji/" in url: return False - if "vk.com/images/" in url: return False - if "vk.com/images/reaction/" in url: return False + # ok + if " https://ok.ru/res/i/" in url: return False - # wikipedia - if "wikipedia.org/static" in url: return False + # vk + if "https://vk.com/emoji/" in url: return False + if "vk.com/images/" in url: return False + if "vk.com/images/reaction/" in url: return False - return True + # wikipedia + if "wikipedia.org/static" in url: return False - @staticmethod - def twitter_best_quality_url(url: str) -> str: - """ - some twitter image URLs point to a less-than best quality - this returns the URL pointing to the highest (original) quality - """ - return re.sub(r"name=(\w+)", "name=orig", url, 1) + return True + +def twitter_best_quality_url(url: str) -> str: + """ + some twitter image URLs point to a less-than best quality + this returns the URL pointing to the highest (original) quality + """ + return re.sub(r"name=(\w+)", "name=orig", url, 1) diff --git a/src/auto_archiver/utils/webdriver.py b/src/auto_archiver/utils/webdriver.py index cf84c35..efb1102 100644 --- a/src/auto_archiver/utils/webdriver.py +++ b/src/auto_archiver/utils/webdriver.py @@ -9,12 +9,72 @@ from loguru import logger from selenium.webdriver.common.by import By import time +#import domain_for_url +from urllib.parse import urlparse, urlunparse +from http.cookiejar import MozillaCookieJar +class CookieSettingDriver(webdriver.Firefox): + + facebook_accept_cookies: bool + cookies: str + cookiejar: MozillaCookieJar + + def __init__(self, cookies, cookiejar, facebook_accept_cookies, *args, **kwargs): + super(CookieSettingDriver, self).__init__(*args, **kwargs) + self.cookies = cookies + self.cookiejar = cookiejar + self.facebook_accept_cookies = facebook_accept_cookies + + def get(self, url: str): + if self.cookies or self.cookiejar: + # set up the driver to make it not 'cookie averse' (needs a context/URL) + # get the 'robots.txt' file which should be quick and easy + robots_url = urlunparse(urlparse(url)._replace(path='/robots.txt', query='', fragment='')) + super(CookieSettingDriver, self).get(robots_url) + + if self.cookies: + # an explicit cookie is set for this site, use that first + for cookie in self.cookies.split(";"): + for name, value in cookie.split("="): + self.driver.add_cookie({'name': name, 'value': value}) + elif self.cookiejar: + domain = urlparse(url).netloc.lstrip("www.") + for cookie in self.cookiejar: + if domain in cookie.domain: + try: + self.add_cookie({ + 'name': cookie.name, + 'value': cookie.value, + 'path': cookie.path, + 'domain': cookie.domain, + 'secure': bool(cookie.secure), + 'expiry': cookie.expires + }) + except Exception as e: + logger.warning(f"Failed to add cookie to webdriver: {e}") + + if self.facebook_accept_cookies: + try: + logger.debug(f'Trying fb click accept cookie popup.') + super(CookieSettingDriver, self).get("http://www.facebook.com") + essential_only = self.find_element(By.XPATH, "//span[contains(text(), 'Decline optional cookies')]") + essential_only.click() + logger.debug(f'fb click worked') + # linux server needs a sleep otherwise facebook cookie won't have worked and we'll get a popup on next page + time.sleep(2) + except Exception as e: + logger.warning(f'Failed on fb accept cookies.', e) + # now get the actual URL + super(CookieSettingDriver, self).get(url) + class Webdriver: - def __init__(self, width: int, height: int, timeout_seconds: int, facebook_accept_cookies: bool = False, http_proxy: str = "", print_options: dict = {}) -> webdriver: + def __init__(self, width: int, height: int, timeout_seconds: int, + facebook_accept_cookies: bool = False, http_proxy: str = "", + print_options: dict = {}, auth: dict = {}) -> webdriver: self.width = width self.height = height self.timeout_seconds = timeout_seconds + self.auth = auth self.facebook_accept_cookies = facebook_accept_cookies self.http_proxy = http_proxy # create and set print options @@ -23,32 +83,26 @@ class Webdriver: setattr(self.print_options, k, v) def __enter__(self) -> webdriver: + options = webdriver.FirefoxOptions() - options.add_argument("--headless") + # options.add_argument("--headless") options.add_argument(f'--proxy-server={self.http_proxy}') options.set_preference('network.protocol-handler.external.tg', False) + # if facebook cookie popup is present, force the browser to English since then it's easier to click the 'Decline optional cookies' option + if self.facebook_accept_cookies: + options.add_argument('--lang=en') + try: - self.driver = webdriver.Firefox(options=options) + self.driver = CookieSettingDriver(cookies=self.auth.get('cookies'), cookiejar=self.auth.get('cookies_jar'), + facebook_accept_cookies=self.facebook_accept_cookies, options=options) self.driver.set_window_size(self.width, self.height) self.driver.set_page_load_timeout(self.timeout_seconds) self.driver.print_options = self.print_options except TimeoutException as e: logger.error(f"failed to get new webdriver, possibly due to insufficient system resources or timeout settings: {e}") - if self.facebook_accept_cookies: - try: - logger.debug(f'Trying fb click accept cookie popup.') - self.driver.get("http://www.facebook.com") - foo = self.driver.find_element(By.XPATH, "//button[@data-cookiebanner='accept_only_essential_button']") - foo.click() - logger.debug(f'fb click worked') - # linux server needs a sleep otherwise facebook cookie won't have worked and we'll get a popup on next page - time.sleep(2) - except: - logger.warning(f'Failed on fb accept cookies.') - return self.driver - + def __exit__(self, exc_type, exc_val, exc_tb): self.driver.close() self.driver.quit() From 72b5ea9ab61d8ae367339cf06b380ed7de1323f2 Mon Sep 17 00:00:00 2001 From: Patrick Robertson Date: Mon, 3 Feb 2025 17:40:40 +0100 Subject: [PATCH 17/62] Restore headless arg --- src/auto_archiver/utils/webdriver.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/auto_archiver/utils/webdriver.py b/src/auto_archiver/utils/webdriver.py index efb1102..005f49d 100644 --- a/src/auto_archiver/utils/webdriver.py +++ b/src/auto_archiver/utils/webdriver.py @@ -85,7 +85,7 @@ class Webdriver: def __enter__(self) -> webdriver: options = webdriver.FirefoxOptions() - # options.add_argument("--headless") + options.add_argument("--headless") options.add_argument(f'--proxy-server={self.http_proxy}') options.set_preference('network.protocol-handler.external.tg', False) # if facebook cookie popup is present, force the browser to English since then it's easier to click the 'Decline optional cookies' option From a873e56b8726c20a3eb93c951451e7bc84133d9a Mon Sep 17 00:00:00 2001 From: Patrick Robertson Date: Tue, 4 Feb 2025 12:57:35 +0100 Subject: [PATCH 18/62] Remove old csv_feeder file - now inside a module --- src/auto_archiver/feeders/csv_feeder.py | 37 ------------------------- 1 file changed, 37 deletions(-) delete mode 100644 src/auto_archiver/feeders/csv_feeder.py diff --git a/src/auto_archiver/feeders/csv_feeder.py b/src/auto_archiver/feeders/csv_feeder.py deleted file mode 100644 index b1aedb7..0000000 --- a/src/auto_archiver/feeders/csv_feeder.py +++ /dev/null @@ -1,37 +0,0 @@ -from loguru import logger -import csv - -from . import Feeder -from ..core import Metadata -from ..utils import url_or_none - -class CSVFeeder(Feeder): - - @staticmethod - def configs() -> dict: - return { - "files": { - "default": None, - "help": "Path to the input file(s) to read the URLs from, comma separated. \ - Input files should be formatted with one URL per line", - "cli_set": lambda cli_val, cur_val: list(set(cli_val.split(","))) - }, - "column": { - "default": None, - "help": "Column number or name to read the URLs from, 0-indexed", - } - } - - - def __iter__(self) -> Metadata: - url_column = self.column or 0 - for file in self.files: - with open(file, "r") as f: - reader = csv.reader(f) - first_row = next(reader) - if not(url_or_none(first_row[url_column])): - # it's a header row, skip it - for row in reader: - url = row[0] - logger.debug(f"Processing {url}") - yield Metadata().set_url(url) \ No newline at end of file From b301f60ea3ba8d8b10658ab2e5a8f592bb3e2af4 Mon Sep 17 00:00:00 2001 From: Patrick Robertson Date: Tue, 4 Feb 2025 13:36:05 +0100 Subject: [PATCH 19/62] Fix using validators set in __manifest__.py E.g. you can use the validator 'is_file' to check if a config is a valid file --- src/auto_archiver/core/orchestrator.py | 4 ++-- src/auto_archiver/core/validators.py | 18 +++++++++++++++--- .../modules/csv_feeder/__manifest__.py | 3 +++ 3 files changed, 20 insertions(+), 5 deletions(-) diff --git a/src/auto_archiver/core/orchestrator.py b/src/auto_archiver/core/orchestrator.py index dbc8a33..8a634de 100644 --- a/src/auto_archiver/core/orchestrator.py +++ b/src/auto_archiver/core/orchestrator.py @@ -207,9 +207,9 @@ class ArchivingOrchestrator: should_store = kwargs.pop('should_store', False) kwargs['dest'] = f"{module.name}.{kwargs.pop('dest', name)}" try: + kwargs['type'] = getattr(validators, kwargs.get('type', '__invalid__')) + except AttributeError: kwargs['type'] = __builtins__.get(kwargs.get('type'), str) - except KeyError: - kwargs['type'] = getattr(validators, kwargs['type']) arg = group.add_argument(f"--{module.name}.{name}", **kwargs) arg.should_store = should_store diff --git a/src/auto_archiver/core/validators.py b/src/auto_archiver/core/validators.py index 681d564..b868ddf 100644 --- a/src/auto_archiver/core/validators.py +++ b/src/auto_archiver/core/validators.py @@ -1,7 +1,19 @@ -# used as validators for config values. +# used as validators for config values. Should raise an exception if the value is invalid. +from pathlib import Path +import argparse def example_validator(value): - return "example" in value + if "example" not in value: + raise argparse.ArgumentTypeError(f"{value} is not a valid value for this argument") + return value def positive_number(value): - return value > 0 \ No newline at end of file + if value < 0: + raise argparse.ArgumentTypeError(f"{value} is not a positive number") + return value + + +def valid_file(value): + if not Path(value).is_file(): + raise argparse.ArgumentTypeError(f"File '{value}' does not exist.") + return value \ No newline at end of file diff --git a/src/auto_archiver/modules/csv_feeder/__manifest__.py b/src/auto_archiver/modules/csv_feeder/__manifest__.py index 7249395..b6d7543 100644 --- a/src/auto_archiver/modules/csv_feeder/__manifest__.py +++ b/src/auto_archiver/modules/csv_feeder/__manifest__.py @@ -13,6 +13,9 @@ "default": None, "help": "Path to the input file(s) to read the URLs from, comma separated. \ Input files should be formatted with one URL per line", + "required": True, + "type": "valid_file", + "nargs": "+", }, "column": { "default": None, From 78e6418249fbf5806e6d0ac110cad81fe526cf8c Mon Sep 17 00:00:00 2001 From: Patrick Robertson Date: Tue, 4 Feb 2025 13:37:17 +0100 Subject: [PATCH 20/62] Unit tests for csv feeder + fix some bugs --- .../modules/csv_feeder/csv_feeder.py | 24 ++++++-- tests/data/csv_no_headers.csv | 2 + tests/data/csv_with_headers.csv | 3 + tests/feeders/test_csv_feeder.py | 57 +++++++++++++++++++ 4 files changed, 82 insertions(+), 4 deletions(-) create mode 100644 tests/data/csv_no_headers.csv create mode 100644 tests/data/csv_with_headers.csv create mode 100644 tests/feeders/test_csv_feeder.py diff --git a/src/auto_archiver/modules/csv_feeder/csv_feeder.py b/src/auto_archiver/modules/csv_feeder/csv_feeder.py index 15dfa85..c3f6eea 100644 --- a/src/auto_archiver/modules/csv_feeder/csv_feeder.py +++ b/src/auto_archiver/modules/csv_feeder/csv_feeder.py @@ -7,16 +7,32 @@ from auto_archiver.utils import url_or_none class CSVFeeder(Feeder): + column = None + + def __iter__(self) -> Metadata: - url_column = self.column or 0 for file in self.files: with open(file, "r") as f: reader = csv.reader(f) first_row = next(reader) - if not(url_or_none(first_row[url_column])): - # it's a header row, skip it + url_column = self.column or 0 + if isinstance(url_column, str): + try: + url_column = first_row.index(url_column) + except ValueError: + logger.error(f"Column {url_column} not found in header row: {first_row}. Did you set the 'column' config correctly?") + return + elif not(url_or_none(first_row[url_column])): + # it's a header row, but we've been given a column number already logger.debug(f"Skipping header row: {first_row}") + else: + # first row isn't a header row, rewind the file + f.seek(0) + for row in reader: - url = row[0] + if not url_or_none(row[url_column]): + logger.warning(f"Not a valid URL in row: {row}, skipping") + continue + url = row[url_column] logger.debug(f"Processing {url}") yield Metadata().set_url(url) \ No newline at end of file diff --git a/tests/data/csv_no_headers.csv b/tests/data/csv_no_headers.csv new file mode 100644 index 0000000..cd66b33 --- /dev/null +++ b/tests/data/csv_no_headers.csv @@ -0,0 +1,2 @@ +https://example.com/1/,data 1 +https://example.com/2/,data 2 \ No newline at end of file diff --git a/tests/data/csv_with_headers.csv b/tests/data/csv_with_headers.csv new file mode 100644 index 0000000..c3e296d --- /dev/null +++ b/tests/data/csv_with_headers.csv @@ -0,0 +1,3 @@ +webpages,other data +https://example.com/1/,data 1 +https://example.com/2/,data 2 \ No newline at end of file diff --git a/tests/feeders/test_csv_feeder.py b/tests/feeders/test_csv_feeder.py new file mode 100644 index 0000000..546c3a7 --- /dev/null +++ b/tests/feeders/test_csv_feeder.py @@ -0,0 +1,57 @@ +import pytest + +@pytest.fixture +def headerless_csv_file(): + return "tests/data/csv_no_headers.csv" + +@pytest.fixture +def header_csv_file(): + return "tests/data/csv_with_headers.csv" + +@pytest.fixture +def header_csv_file_non_default_column(): + return "tests/data/csv_with_headers_non_default_column.csv" + + +def test_csv_feeder_no_headers(headerless_csv_file, setup_module): + from auto_archiver.modules.csv_feeder.csv_feeder import CSVFeeder + + feeder = setup_module(CSVFeeder, {"files": [headerless_csv_file]}) + + urls = list(feeder) + assert len(urls) == 2 + assert urls[0].get_url() == "https://example.com/1/" + assert urls[1].get_url() == "https://example.com/2/" + +def test_csv_feeder_with_headers(header_csv_file, setup_module): + from auto_archiver.modules.csv_feeder.csv_feeder import CSVFeeder + + feeder = setup_module(CSVFeeder, {"files": [header_csv_file]}) + + urls = list(feeder) + assert len(urls) == 2 + assert urls[0].get_url() == "https://example.com/1/" + assert urls[1].get_url() == "https://example.com/2/" + +def test_csv_feeder_wrong_column(header_csv_file, setup_module, caplog): + from auto_archiver.modules.csv_feeder.csv_feeder import CSVFeeder + + + with caplog.at_level("WARNING"): + feeder = setup_module(CSVFeeder, {"files": [header_csv_file], "column": 1}) + urls = list(feeder) + + assert len(urls) == 0 + assert "Not a valid URL in row" in caplog.text + assert len(caplog.records) == 2 + + +def test_csv_feeder_column_by_name(header_csv_file, setup_module): + from auto_archiver.modules.csv_feeder.csv_feeder import CSVFeeder + + feeder = setup_module(CSVFeeder, {"files": [header_csv_file], "column": "webpages"}) + + urls = list(feeder) + assert len(urls) == 2 + assert urls[0].get_url() == "https://example.com/1/" + assert urls[1].get_url() == "https://example.com/2/" \ No newline at end of file From 034197a81f83210d6a4350010d8432e823226231 Mon Sep 17 00:00:00 2001 From: Patrick Robertson Date: Tue, 4 Feb 2025 13:40:07 +0100 Subject: [PATCH 21/62] Fix typos in csv feeder docs (in manifest) --- src/auto_archiver/modules/csv_feeder/__manifest__.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/auto_archiver/modules/csv_feeder/__manifest__.py b/src/auto_archiver/modules/csv_feeder/__manifest__.py index b6d7543..6d4c7bf 100644 --- a/src/auto_archiver/modules/csv_feeder/__manifest__.py +++ b/src/auto_archiver/modules/csv_feeder/__manifest__.py @@ -30,7 +30,8 @@ - Allows specifying the column number or name to extract URLs from. - Skips header rows if the first value is not a valid URL. - ### Setu N - - Input files should be formatted with one URL per line. + ### Setup + - Input files should be formatted with one URL per line, with or without a header row. + - If you have a header row, you can specify the column number or name to read URLs from using the 'column' config option. """ } From 0633e17998807e6d3c4c564b103003e93df39b98 Mon Sep 17 00:00:00 2001 From: Patrick Robertson Date: Tue, 4 Feb 2025 14:18:46 +0100 Subject: [PATCH 22/62] Close the facebook 'login' window if it's there - to allow for proper screenshots --- src/auto_archiver/utils/webdriver.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/src/auto_archiver/utils/webdriver.py b/src/auto_archiver/utils/webdriver.py index 005f49d..db26d04 100644 --- a/src/auto_archiver/utils/webdriver.py +++ b/src/auto_archiver/utils/webdriver.py @@ -66,6 +66,13 @@ class CookieSettingDriver(webdriver.Firefox): logger.warning(f'Failed on fb accept cookies.', e) # now get the actual URL super(CookieSettingDriver, self).get(url) + if self.facebook_accept_cookies: + # try and click the 'close' button on the 'login' window to close it + close_button = self.find_element(By.XPATH, "//div[@role='dialog']//div[@aria-label='Close']") + if close_button: + close_button.click() + + class Webdriver: def __init__(self, width: int, height: int, timeout_seconds: int, From 91ca325fd510c0c1989b6b31050571331a074c3b Mon Sep 17 00:00:00 2001 From: Patrick Robertson Date: Tue, 4 Feb 2025 17:46:46 +0100 Subject: [PATCH 23/62] Update yt-dlp to latest version + remove code no longer needed from bluesky dropin --- poetry.lock | 8 ++++---- pyproject.toml | 2 +- .../modules/generic_extractor/bluesky.py | 13 +------------ .../modules/generic_extractor/facebook.py | 17 +++++++++++++++++ 4 files changed, 23 insertions(+), 17 deletions(-) create mode 100644 src/auto_archiver/modules/generic_extractor/facebook.py diff --git a/poetry.lock b/poetry.lock index 088fc70..8fb48ec 100644 --- a/poetry.lock +++ b/poetry.lock @@ -3077,14 +3077,14 @@ h11 = ">=0.9.0,<1" [[package]] name = "yt-dlp" -version = "2025.1.12" +version = "2025.1.26" description = "A feature-rich command-line audio/video downloader" optional = false python-versions = ">=3.9" groups = ["main"] files = [ - {file = "yt_dlp-2025.1.12-py3-none-any.whl", hash = "sha256:f7ea19afb64f8e457a1b9598ddb67f8deaa313bf1d57abd5612db9272ab10795"}, - {file = "yt_dlp-2025.1.12.tar.gz", hash = "sha256:8e7e246e2a5a2cff0a9c13db46844a37a547680702012058c94ec18fce0ca25a"}, + {file = "yt_dlp-2025.1.26-py3-none-any.whl", hash = "sha256:3e76bd896b9f96601021ca192ca0fbdd195e3c3dcc28302a3a34c9bc4979da7b"}, + {file = "yt_dlp-2025.1.26.tar.gz", hash = "sha256:1c9738266921ad43c568ad01ac3362fb7c7af549276fbec92bd72f140da16240"}, ] [package.extras] @@ -3100,4 +3100,4 @@ test = ["pytest (>=8.1,<9.0)", "pytest-rerunfailures (>=14.0,<15.0)"] [metadata] lock-version = "2.1" python-versions = ">=3.10,<3.13" -content-hash = "5a54c84ba388db7b77d1c28973b710fc99aa3822a2860b30acaf5b02ba1927bd" +content-hash = "9ca114395e73af8982abbccc25b385bbca62e50ba7cca8239e52e5c1227cb4b0" diff --git a/pyproject.toml b/pyproject.toml index 3cd47e7..f1be273 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -46,7 +46,7 @@ dependencies = [ "cryptography (>=41.0.0,<42.0.0)", "boto3 (>=1.28.0,<2.0.0)", "dataclasses-json (>=0.0.0)", - "yt-dlp (==2025.1.12)", + "yt-dlp (>=2025.1.26,<2026.0.0)", "numpy (==2.1.3)", "vk-url-scraper (>=0.0.0)", "requests[socks] (>=0.0.0)", diff --git a/src/auto_archiver/modules/generic_extractor/bluesky.py b/src/auto_archiver/modules/generic_extractor/bluesky.py index 1f92fd8..f2086b0 100644 --- a/src/auto_archiver/modules/generic_extractor/bluesky.py +++ b/src/auto_archiver/modules/generic_extractor/bluesky.py @@ -23,19 +23,8 @@ class Bluesky(GenericDropin): def extract_post(self, url: str, ie_instance: InfoExtractor) -> dict: # TODO: If/when this PR (https://github.com/yt-dlp/yt-dlp/pull/12098) is merged on ytdlp, remove the comments and delete the code below - # handle, video_id = ie_instance._match_valid_url(url).group('handle', 'id') - # return ie_instance._extract_post(handle=handle, post_id=video_id) - handle, video_id = ie_instance._match_valid_url(url).group('handle', 'id') - return ie_instance._download_json( - 'https://public.api.bsky.app/xrpc/app.bsky.feed.getPostThread', - video_id, query={ - 'uri': f'at://{handle}/app.bsky.feed.post/{video_id}', - 'depth': 0, - 'parentHeight': 0, - })['thread']['post'] - - + return ie_instance._extract_post(handle=handle, post_id=video_id) def _download_bsky_embeds(self, post: dict, archiver: Extractor) -> list[Media]: """ diff --git a/src/auto_archiver/modules/generic_extractor/facebook.py b/src/auto_archiver/modules/generic_extractor/facebook.py new file mode 100644 index 0000000..352d44e --- /dev/null +++ b/src/auto_archiver/modules/generic_extractor/facebook.py @@ -0,0 +1,17 @@ +from .dropin import GenericDropin + + +class Facebook(GenericDropin): + def extract_post(self, url: str, ie_instance): + video_id = ie_instance._match_valid_url(url).group('id') + ie_instance._download_webpage( + url.replace('://m.facebook.com/', '://www.facebook.com/'), video_id) + webpage = ie_instance._download_webpage(url, ie_instance._match_valid_url(url).group('id')) + + post_data = ie_instance._extract_from_url.extract_metadata(webpage) + return post_data + + def create_metadata(self, post: dict, ie_instance, archiver, url): + metadata = archiver.create_metadata(url) + metadata.set_title(post.get('title')).set_content(post.get('description')).set_post_data(post) + return metadata \ No newline at end of file From 48abb5e66b989f7b97dad5b3ea2479b22c90e2d0 Mon Sep 17 00:00:00 2001 From: Patrick Robertson Date: Tue, 4 Feb 2025 18:16:03 +0100 Subject: [PATCH 24/62] Remove dangling screenshot_enricher file. Moved to modules/screenshot_enricher --- .../enrichers/screenshot_enricher.py | 40 ------------------- 1 file changed, 40 deletions(-) delete mode 100644 src/auto_archiver/enrichers/screenshot_enricher.py diff --git a/src/auto_archiver/enrichers/screenshot_enricher.py b/src/auto_archiver/enrichers/screenshot_enricher.py deleted file mode 100644 index abb1e16..0000000 --- a/src/auto_archiver/enrichers/screenshot_enricher.py +++ /dev/null @@ -1,40 +0,0 @@ -from loguru import logger -import time, os -from selenium.common.exceptions import TimeoutException - - -from auto_archiver.core import Enricher -from ..utils import Webdriver, url as UrlUtil, random_str -from ..core import Media, Metadata - -class ScreenshotEnricher(Enricher): - name = "screenshot_enricher" - - @staticmethod - def configs() -> dict: - return { - "width": {"default": 1280, "help": "width of the screenshots"}, - "height": {"default": 720, "help": "height of the screenshots"}, - "timeout": {"default": 60, "help": "timeout for taking the screenshot"}, - "sleep_before_screenshot": {"default": 4, "help": "seconds to wait for the pages to load before taking screenshot"}, - "http_proxy": {"default": "", "help": "http proxy to use for the webdriver, eg http://proxy-user:password@proxy-ip:port"}, - } - - def enrich(self, to_enrich: Metadata) -> None: - url = to_enrich.get_url() - if UrlUtil.is_auth_wall(url): - logger.debug(f"[SKIP] SCREENSHOT since url is behind AUTH WALL: {url=}") - return - - logger.debug(f"Enriching screenshot for {url=}") - with Webdriver(self.width, self.height, self.timeout, 'facebook.com' in url, http_proxy=self.http_proxy) as driver: - try: - driver.get(url) - time.sleep(int(self.sleep_before_screenshot)) - screenshot_file = os.path.join(self.tmp_dir, f"screenshot_{random_str(8)}.png") - driver.save_screenshot(screenshot_file) - to_enrich.add_media(Media(filename=screenshot_file), id="screenshot") - except TimeoutException: - logger.info("TimeoutException loading page for screenshot") - except Exception as e: - logger.error(f"Got error while loading webdriver for screenshot enricher: {e}") From 52542812dcbd171f1606a4f7502becb1101bd570 Mon Sep 17 00:00:00 2001 From: erinhmclark Date: Wed, 5 Feb 2025 16:42:58 +0000 Subject: [PATCH 25/62] Merge tests from version with context. --- .../modules/gsheet_db/gsheet_db.py | 15 +- .../instagram_tbot_extractor.py | 80 ++++-- .../modules/telethon_extractor/__init__.py | 2 +- .../telethon_extractor/telethon_extractor.py | 2 +- tests/conftest.py | 19 +- tests/databases/test_gsheet_db.py | 140 +++++++++ .../test_instagram_api_extractor.py | 108 +++++++ .../test_instagram_tbot_extractor.py | 111 ++++++++ tests/feeders/test_gsheet_feeder.py | 268 ++++++++++++++++++ tests/feeders/test_gworksheet.py | 144 ++++++++++ tests/storages/test_S3_storage.py | 100 +++++++ tests/storages/test_gdrive_storage.py | 43 +++ tests/storages/test_storage_base.py | 23 ++ 13 files changed, 1022 insertions(+), 33 deletions(-) create mode 100644 tests/databases/test_gsheet_db.py create mode 100644 tests/extractors/test_instagram_api_extractor.py create mode 100644 tests/extractors/test_instagram_tbot_extractor.py create mode 100644 tests/feeders/test_gsheet_feeder.py create mode 100644 tests/feeders/test_gworksheet.py create mode 100644 tests/storages/test_S3_storage.py create mode 100644 tests/storages/test_gdrive_storage.py create mode 100644 tests/storages/test_storage_base.py diff --git a/src/auto_archiver/modules/gsheet_db/gsheet_db.py b/src/auto_archiver/modules/gsheet_db/gsheet_db.py index 5e1ed1e..644015e 100644 --- a/src/auto_archiver/modules/gsheet_db/gsheet_db.py +++ b/src/auto_archiver/modules/gsheet_db/gsheet_db.py @@ -12,10 +12,11 @@ from auto_archiver.modules.gsheet_feeder import GWorksheet class GsheetsDb(Database): """ - NB: only works if GsheetFeeder is used. - could be updated in the future to support non-GsheetFeeder metadata + NB: only works if GsheetFeeder is used. + could be updated in the future to support non-GsheetFeeder metadata """ + def started(self, item: Metadata) -> None: logger.warning(f"STARTED {item}") gw, row = self._retrieve_gsheet(item) @@ -57,7 +58,7 @@ class GsheetsDb(Database): media: Media = item.get_final_media() if hasattr(media, "urls"): batch_if_valid('archive', "\n".join(media.urls)) - batch_if_valid('date', True, datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=datetime.timezone.utc).isoformat()) + batch_if_valid('date', True, self._get_current_datetime_iso()) batch_if_valid('title', item.get_title()) batch_if_valid('text', item.get("content", "")) batch_if_valid('timestamp', item.get_timestamp()) @@ -85,6 +86,12 @@ class GsheetsDb(Database): gw.batch_set_cell(cell_updates) + @staticmethod + def _get_current_datetime_iso() -> str: + """Helper method to generate the current datetime in ISO format.""" + return datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=datetime.timezone.utc).isoformat() + + def _safe_status_update(self, item: Metadata, new_status: str) -> None: try: gw, row = self._retrieve_gsheet(item) @@ -93,9 +100,11 @@ class GsheetsDb(Database): logger.debug(f"Unable to update sheet: {e}") def _retrieve_gsheet(self, item: Metadata) -> Tuple[GWorksheet, int]: + if gsheet := item.get_context("gsheet"): gw: GWorksheet = gsheet.get("worksheet") row: int = gsheet.get("row") + # todo doesn't exist, should be passed from elif self.sheet_id: print(self.sheet_id) diff --git a/src/auto_archiver/modules/instagram_tbot_extractor/instagram_tbot_extractor.py b/src/auto_archiver/modules/instagram_tbot_extractor/instagram_tbot_extractor.py index 5b49484..5660cd2 100644 --- a/src/auto_archiver/modules/instagram_tbot_extractor/instagram_tbot_extractor.py +++ b/src/auto_archiver/modules/instagram_tbot_extractor/instagram_tbot_extractor.py @@ -34,19 +34,30 @@ class InstagramTbotExtractor(Extractor): """ super().setup(configs) logger.info(f"SETUP {self.name} checking login...") + self._prepare_session_file() + self._initialize_telegram_client() - # make a copy of the session that is used exclusively with this archiver instance + def _prepare_session_file(self): + """ + Creates a copy of the session file for exclusive use with this archiver instance. + Ensures that a valid session file exists before proceeding. + """ new_session_file = os.path.join("secrets/", f"instabot-{time.strftime('%Y-%m-%d')}{random_str(8)}.session") if not os.path.exists(f"{self.session_file}.session"): - raise FileNotFoundError(f"session file {self.session_file}.session not found, " - f"to set this up run the setup script in scripts/telegram_setup.py") + raise FileNotFoundError(f"Session file {self.session_file}.session not found.") shutil.copy(self.session_file + ".session", new_session_file) self.session_file = new_session_file.replace(".session", "") + def _initialize_telegram_client(self): + """Initializes the Telegram client.""" try: self.client = TelegramClient(self.session_file, self.api_id, self.api_hash) except OperationalError as e: - logger.error(f"Unable to access the {self.session_file} session, please make sure you don't use the same session file here and in telethon_extractor. if you do then disable at least one of the archivers for the 1st time you setup telethon session: {e}") + logger.error( + f"Unable to access the {self.session_file} session. " + "Ensure that you don't use the same session file here and in telethon_extractor. " + "If you do, disable at least one of the archivers for the first-time setup of the telethon session: {e}" + ) with self.client.start(): logger.success(f"SETUP {self.name} login works.") @@ -63,32 +74,49 @@ class InstagramTbotExtractor(Extractor): result = Metadata() tmp_dir = self.tmp_dir with self.client.start(): - chat = self.client.get_entity("instagram_load_bot") - since_id = self.client.send_message(entity=chat, message=url).id - attempts = 0 - seen_media = [] - message = "" - time.sleep(3) - # media is added before text by the bot so it can be used as a stop-logic mechanism - while attempts < (self.timeout - 3) and (not message or not len(seen_media)): - attempts += 1 - time.sleep(1) - for post in self.client.iter_messages(chat, min_id=since_id): - since_id = max(since_id, post.id) - if post.media and post.id not in seen_media: - filename_dest = os.path.join(tmp_dir, f'{chat.id}_{post.id}') - media = self.client.download_media(post.media, filename_dest) - if media: - result.add_media(Media(media)) - seen_media.append(post.id) - if post.message: message += post.message + chat, since_id = self._send_url_to_bot(url) + message = self._process_messages(chat, since_id, tmp_dir, result) - if "You must enter a URL to a post" in message: + if "You must enter a URL to a post" in message: logger.debug(f"invalid link {url=} for {self.name}: {message}") return False - + # # TODO: It currently returns this as a success - is that intentional? + # if "Media not found or unavailable" in message: + # logger.debug(f"invalid link {url=} for {self.name}: {message}") + # return False + if message: result.set_content(message).set_title(message[:128]) - return result.success("insta-via-bot") + + def _send_url_to_bot(self, url: str): + """ + Sends the URL to the 'instagram_load_bot' and returns (chat, since_id). + """ + chat = self.client.get_entity("instagram_load_bot") + since_message = self.client.send_message(entity=chat, message=url) + return chat, since_message.id + + def _process_messages(self, chat, since_id, tmp_dir, result): + attempts = 0 + seen_media = [] + message = "" + time.sleep(3) + # media is added before text by the bot so it can be used as a stop-logic mechanism + while attempts < (self.timeout - 3) and (not message or not len(seen_media)): + attempts += 1 + time.sleep(1) + for post in self.client.iter_messages(chat, min_id=since_id): + since_id = max(since_id, post.id) + # Skip known filler message: + if post.message == 'The bot receives information through https://hikerapi.com/p/hJqpppqi': + continue + if post.media and post.id not in seen_media: + filename_dest = os.path.join(tmp_dir, f'{chat.id}_{post.id}') + media = self.client.download_media(post.media, filename_dest) + if media: + result.add_media(Media(media)) + seen_media.append(post.id) + if post.message: message += post.message + return message.strip() \ No newline at end of file diff --git a/src/auto_archiver/modules/telethon_extractor/__init__.py b/src/auto_archiver/modules/telethon_extractor/__init__.py index a837fdf..2eaa57c 100644 --- a/src/auto_archiver/modules/telethon_extractor/__init__.py +++ b/src/auto_archiver/modules/telethon_extractor/__init__.py @@ -1 +1 @@ -from .telethon_extractor import TelethonArchiver \ No newline at end of file +from .telethon_extractor import TelethonExtractor \ No newline at end of file diff --git a/src/auto_archiver/modules/telethon_extractor/telethon_extractor.py b/src/auto_archiver/modules/telethon_extractor/telethon_extractor.py index 3e952e8..0147ff2 100644 --- a/src/auto_archiver/modules/telethon_extractor/telethon_extractor.py +++ b/src/auto_archiver/modules/telethon_extractor/telethon_extractor.py @@ -13,7 +13,7 @@ from auto_archiver.core import Metadata, Media from auto_archiver.utils import random_str -class TelethonArchiver(Extractor): +class TelethonExtractor(Extractor): valid_url = re.compile(r"https:\/\/t\.me(\/c){0,1}\/(.+)\/(\d+)") invite_pattern = re.compile(r"t.me(\/joinchat){0,1}\/\+?(.+)") diff --git a/tests/conftest.py b/tests/conftest.py index f909bfb..8675fbc 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,7 +1,8 @@ """ pytest conftest file, for shared fixtures and configuration """ - +import os +import pickle from tempfile import TemporaryDirectory from typing import Dict, Tuple import hashlib @@ -113,4 +114,18 @@ def pytest_runtest_setup(item): test_name = _test_failed_incremental[cls_name].get((), None) # if name found, test has failed for the combination of class name & test name if test_name is not None: - pytest.xfail(f"previous test failed ({test_name})") \ No newline at end of file + pytest.xfail(f"previous test failed ({test_name})") + + + +@pytest.fixture() +def unpickle(): + """ + Returns a helper function that unpickles a file + ** gets the file from the test_files directory: tests/data/test_files ** + """ + def _unpickle(path): + test_data_dir = os.path.join(os.path.dirname(__file__), "data", "test_files") + with open(os.path.join(test_data_dir, path), "rb") as f: + return pickle.load(f) + return _unpickle \ No newline at end of file diff --git a/tests/databases/test_gsheet_db.py b/tests/databases/test_gsheet_db.py new file mode 100644 index 0000000..bdc2811 --- /dev/null +++ b/tests/databases/test_gsheet_db.py @@ -0,0 +1,140 @@ +from datetime import datetime, timezone +from unittest.mock import MagicMock, patch + +import pytest + +from auto_archiver.core import Metadata, Media +from auto_archiver.modules.gsheet_db import GsheetsDb +from auto_archiver.modules.gsheet_feeder import GWorksheet + + +@pytest.fixture +def mock_gworksheet(): + mock_gworksheet = MagicMock(spec=GWorksheet) + mock_gworksheet.col_exists.return_value = True + mock_gworksheet.get_cell.return_value = "" + mock_gworksheet.get_row.return_value = {} + return mock_gworksheet + + +@pytest.fixture +def mock_metadata(): + metadata: Metadata = MagicMock(spec=Metadata) + metadata.get_url.return_value = "http://example.com" + metadata.status = "done" + metadata.get_title.return_value = "Example Title" + metadata.get.return_value = "Example Content" + metadata.get_timestamp.return_value = "2025-01-01T00:00:00Z" + metadata.get_final_media.return_value = MagicMock(spec=Media) + metadata.get_all_media.return_value = [] + metadata.get_media_by_id.return_value = None + metadata.get_first_image.return_value = None + return metadata + +@pytest.fixture +def metadata(): + metadata = Metadata() + metadata.add_media(Media(filename="screenshot", urls=["http://example.com/screenshot.png"])) + metadata.add_media(Media(filename="browsertrix", urls=["http://example.com/browsertrix.wacz"])) + metadata.add_media(Media(filename="thumbnail", urls=["http://example.com/thumbnail.png"])) + metadata.set_url("http://example.com") + metadata.set_title("Example Title") + metadata.set_content("Example Content") + metadata.success("my-archiver") + metadata.set("timestamp", "2025-01-01T00:00:00Z") + metadata.set("date", "2025-02-04T18:22:24.909112+00:00") + return metadata + + +@pytest.fixture +def mock_media(): + """Fixture for a mock Media object.""" + mock_media = MagicMock(spec=Media) + mock_media.urls = ["http://example.com/media"] + mock_media.get.return_value = "not-calculated" + return mock_media + +@pytest.fixture +def gsheets_db(mock_gworksheet, setup_module): + db = setup_module("gsheet_db", { + "allow_worksheets": "set()", + "block_worksheets": "set()", + "use_sheet_names_in_stored_paths": "True", + }) + db._retrieve_gsheet = MagicMock(return_value=(mock_gworksheet, 1)) + return db + + +@pytest.fixture +def fixed_timestamp(): + """Fixture for a fixed timestamp.""" + return datetime(2025, 1, 1, tzinfo=timezone.utc) + + +@pytest.fixture +def expected_calls(mock_media, fixed_timestamp): + """Fixture for the expected cell updates.""" + return [ + (1, 'status', 'my-archiver: success'), + (1, 'archive', 'http://example.com/screenshot.png'), + (1, 'date', '2025-02-01T00:00:00+00:00'), + (1, 'title', 'Example Title'), + (1, 'text', 'Example Content'), + (1, 'timestamp', '2025-01-01T00:00:00+00:00'), + (1, 'hash', 'not-calculated'), + # (1, 'screenshot', 'http://example.com/screenshot.png'), + # (1, 'thumbnail', '=IMAGE("http://example.com/thumbnail.png")'), + # (1, 'wacz', 'http://example.com/browsertrix.wacz'), + # (1, 'replaywebpage', 'https://replayweb.page/?source=http%3A%2F%2Fexample.com%2Fbrowsertrix.wacz#view=pages&url=') + ] + +def test_retrieve_gsheet(gsheets_db, metadata, mock_gworksheet): + gw, row = gsheets_db._retrieve_gsheet(metadata) + assert gw == mock_gworksheet + assert row == 1 + + +def test_started(gsheets_db, mock_metadata, mock_gworksheet): + gsheets_db.started(mock_metadata) + mock_gworksheet.set_cell.assert_called_once_with(1, 'status', 'Archive in progress') + +def test_failed(gsheets_db, mock_metadata, mock_gworksheet): + reason = "Test failure" + gsheets_db.failed(mock_metadata, reason) + mock_gworksheet.set_cell.assert_called_once_with(1, 'status', f'Archive failed {reason}') + +def test_aborted(gsheets_db, mock_metadata, mock_gworksheet): + gsheets_db.aborted(mock_metadata) + mock_gworksheet.set_cell.assert_called_once_with(1, 'status', '') + + +def test_done(gsheets_db, metadata, mock_gworksheet, expected_calls): + with patch.object(gsheets_db, '_get_current_datetime_iso', return_value='2025-02-01T00:00:00+00:00'): + gsheets_db.done(metadata) + mock_gworksheet.batch_set_cell.assert_called_once_with(expected_calls) + + +def test_done_cached(gsheets_db, metadata, mock_gworksheet): + with patch.object(gsheets_db, '_get_current_datetime_iso', return_value='2025-02-01T00:00:00+00:00'): + gsheets_db.done(metadata, cached=True) + + # Verify the status message includes "[cached]" + call_args = mock_gworksheet.batch_set_cell.call_args[0][0] + assert any(call[2].startswith("[cached]") for call in call_args) + + +def test_done_missing_media(gsheets_db, metadata, mock_gworksheet): + # clear media from metadata + metadata.media = [] + with patch.object(gsheets_db, '_get_current_datetime_iso', return_value='2025-02-01T00:00:00+00:00'): + gsheets_db.done(metadata) + # Verify nothing media-related gets updated + call_args = mock_gworksheet.batch_set_cell.call_args[0][0] + media_fields = {'archive', 'screenshot', 'thumbnail', 'wacz', 'replaywebpage'} + assert all(call[1] not in media_fields for call in call_args) + +def test_safe_status_update(gsheets_db, metadata, mock_gworksheet): + gsheets_db._safe_status_update(metadata, "Test status") + mock_gworksheet.set_cell.assert_called_once_with(1, 'status', 'Test status') + + diff --git a/tests/extractors/test_instagram_api_extractor.py b/tests/extractors/test_instagram_api_extractor.py new file mode 100644 index 0000000..7a19233 --- /dev/null +++ b/tests/extractors/test_instagram_api_extractor.py @@ -0,0 +1,108 @@ +from datetime import datetime +from typing import Type + +import pytest +from unittest.mock import patch, MagicMock + +from auto_archiver.core import Metadata +from auto_archiver.modules.instagram_api_extractor.instagram_api_extractor import InstagramAPIExtractor +from .test_extractor_base import TestExtractorBase + + +@pytest.fixture +def mock_user_response(): + return { + "user": { + "pk": "123", + "username": "test_user", + "full_name": "Test User", + "profile_pic_url_hd": "http://example.com/profile.jpg", + "profile_pic_url": "http://example.com/profile_lowres.jpg" + } + } + +@pytest.fixture +def mock_post_response(): + return { + "id": "post_123", + "code": "abc123", + "caption_text": "Test Caption", + "taken_at": datetime.now().timestamp(), + "video_url": "http://example.com/video.mp4", + "thumbnail_url": "http://example.com/thumbnail.jpg" + } + +@pytest.fixture +def mock_story_response(): + return [{ + "id": "story_123", + "taken_at": datetime.now().timestamp(), + "video_url": "http://example.com/story.mp4" + }] + +@pytest.fixture +def mock_highlight_response(): + return { + "response": { + "reels": { + "highlight:123": { + "id": "123", + "title": "Test Highlight", + "items": [{ + "id": "item_123", + "taken_at": datetime.now().timestamp(), + "video_url": "http://example.com/highlight.mp4" + }] + } + } + } + } + + +# @pytest.mark.incremental +class TestInstagramAPIExtractor(TestExtractorBase): + """ + Test suite for InstagramAPIExtractor. + """ + + extractor_module = "instagram_api_extractor" + extractor: InstagramAPIExtractor + + config = { + "access_token": "test_access_token", + "api_endpoint": "https://api.instagram.com/v1", + # "full_profile": False, + # "full_profile_max_posts": 0, + # "minimize_json_output": True, + } + + @pytest.mark.parametrize("url,expected", [ + ("https://instagram.com/user", [("", "user", "")]), + ("https://instagr.am/p/post_id", []), + ("https://youtube.com", []), + ("https://www.instagram.com/reel/reel_id", [("reel", "reel_id", "")]), + ("https://instagram.com/stories/highlights/123", [("stories/highlights", "123", "")]), + ("https://instagram.com/stories/user/123", [("stories", "user", "123")]), + ]) + def test_url_parsing(self, url, expected): + assert self.extractor.valid_url.findall(url) == expected + + def test_initialize(self): + self.extractor.initialise() + assert self.extractor.api_endpoint[-1] != "/" + + @pytest.mark.parametrize("input_dict,expected", [ + ({"x": 0, "valid": "data"}, {"valid": "data"}), + ({"nested": {"y": None, "valid": [{}]}}, {"nested": {"valid": [{}]}}), + ]) + def test_cleanup_dict(self, input_dict, expected): + assert self.extractor.cleanup_dict(input_dict) == expected + + def test_download_post(self): + # test with context=reel + # test with context=post + # test with multiple images + # test gets text (metadata title) + + + pass \ No newline at end of file diff --git a/tests/extractors/test_instagram_tbot_extractor.py b/tests/extractors/test_instagram_tbot_extractor.py new file mode 100644 index 0000000..4fe80be --- /dev/null +++ b/tests/extractors/test_instagram_tbot_extractor.py @@ -0,0 +1,111 @@ +import os +import pickle +from typing import Type +from unittest.mock import patch, MagicMock + +import pytest + +from auto_archiver.core.extractor import Extractor +from auto_archiver.modules.instagram_tbot_extractor import InstagramTbotExtractor + + +TESTFILES = os.path.join(os.path.dirname(__file__), "testfiles") + + +@pytest.fixture +def test_session_file(tmpdir): + """Fixture to create a test session file.""" + session_file = os.path.join(tmpdir, "test_session.session") + with open(session_file, "w") as f: + f.write("mock_session_data") + return session_file.replace(".session", "") + + +@pytest.mark.incremental +class TestInstagramTbotExtractor(object): + """ + Test suite for InstagramTbotExtractor. + """ + + extractor_module = "instagram_tbot_extractor" + extractor: InstagramTbotExtractor + config = { + "api_id": 12345, + "api_hash": "test_api_hash", + # "session_file" + } + + @pytest.fixture(autouse=True) + def setup_extractor(self, setup_module): + assert self.extractor_module is not None, "self.extractor_module must be set on the subclass" + assert self.config is not None, "self.config must be a dict set on the subclass" + extractor: Type[Extractor] = setup_module(self.extractor_module, self.config) + return extractor + + @pytest.fixture + def mock_telegram_client(self): + """Fixture to mock TelegramClient interactions.""" + with patch("auto_archiver.modules.instagram_tbot_extractor._initialize_telegram_client") as mock_client: + instance = MagicMock() + mock_client.return_value = instance + yield instance + + + # @pytest.fixture + # def mock_session_file(self, temp_session_file): + # """Patch the extractor’s session file setup to use a temporary path.""" + # with patch.object(InstagramTbotExtractor, "session_file", temp_session_file): + # with patch.object(InstagramTbotExtractor, "_prepare_session_file", return_value=None): + # yield # Mocks are applied for the duration of the test + + @pytest.fixture + def metadata_sample(self): + """Loads a Metadata object from a pickle file.""" + with open(os.path.join(TESTFILES, "metadata_item.pkl"), "rb") as f: + return pickle.load(f) + + + @pytest.mark.download + @pytest.mark.parametrize("url, expected_status, bot_responses", [ + ("https://www.instagram.com/p/C4QgLbrIKXG", "insta-via-bot: success", [MagicMock(id=101, media=None, message="Are you new to Bellingcat? - The way we share our investigations is different. 💭\nWe want you to read our story but also learn ou")]), + ("https://www.instagram.com/reel/DEVLK8qoIbg/", "insta-via-bot: success", [MagicMock(id=101, media=None, message="Our volunteer community is at the centre of many incredible Bellingcat investigations and tools. Stephanie Ladel is one such vol")]), + # todo tbot not working for stories :( + ("https://www.instagram.com/stories/bellingcatofficial/3556336382743057476/", False, [MagicMock(id=101, media=None, message="Media not found or unavailable")]), + ("https://www.youtube.com/watch?v=ymCMy8OffHM", False, []), + ("https://www.instagram.com/p/INVALID", False, [MagicMock(id=101, media=None, message="You must enter a URL to a post")]), + ]) + def test_download(self, url, expected_status, bot_responses, metadata_sample): + """Test the `download()` method with various Instagram URLs.""" + metadata_sample.set_url(url) + self.extractor.initialise() + result = self.extractor.download(metadata_sample) + if expected_status: + assert result.is_success() + assert result.status == expected_status + assert result.metadata.get("title") in [msg.message[:128] for msg in bot_responses if msg.message] + else: + assert result is False + # self.extractor.cleanup() + + # @patch.object(InstagramTbotExtractor, '_send_url_to_bot') + # @patch.object(InstagramTbotExtractor, '_process_messages') + # def test_download_invalid_link_returns_false( + # self, mock_process, mock_send, extractor, metadata_instagram + # ): + # # Setup Mocks + # # _send_url_to_bot -> simulate it returns (chat=MagicMock, since_id=100) + # mock_chat = MagicMock() + # mock_send.return_value = (mock_chat, 100) + # # _process_messages -> simulate it returns the text "You must enter a URL to a post" + # mock_process.return_value = "You must enter a URL to a post" + # result = extractor.download(metadata_instagram) + # assert result is False, "Should return False if message includes 'You must enter a URL to a post'" + + + + + # Test story +# Test expired story +# Test requires login/ access (?) +# Test post +# Test multiple images? \ No newline at end of file diff --git a/tests/feeders/test_gsheet_feeder.py b/tests/feeders/test_gsheet_feeder.py new file mode 100644 index 0000000..dbd2416 --- /dev/null +++ b/tests/feeders/test_gsheet_feeder.py @@ -0,0 +1,268 @@ +from typing import Type + +import gspread +import pytest +from unittest.mock import patch, MagicMock +from auto_archiver.modules.gsheet_feeder import GsheetsFeeder +from auto_archiver.core import Metadata, Feeder, ArchivingContext + + +def test_initialise_without_sheet_and_sheet_id(setup_module): + """Ensure initialise() raises AssertionError if neither sheet nor sheet_id is set. + (shouldn't really be asserting in there) + """ + with patch("gspread.service_account"): + feeder = setup_module("gsheet_feeder", + {"service_account": "dummy.json", + "sheet": None, + "sheet_id": None}) + with pytest.raises(AssertionError): + feeder.initialise() + + +@pytest.fixture +def gsheet_feeder(setup_module) -> GsheetsFeeder: + feeder = setup_module("gsheet_feeder", + {"service_account": "dummy.json", + "sheet": "test-auto-archiver", + "sheet_id": None, + "header": 1, + "columns": { + "url": "link", + "status": "archive status", + "folder": "destination folder", + "archive": "archive location", + "date": "archive date", + "thumbnail": "thumbnail", + "timestamp": "upload timestamp", + "title": "upload title", + "text": "text content", + "screenshot": "screenshot", + "hash": "hash", + "pdq_hash": "perceptual hashes", + "wacz": "wacz", + "replaywebpage": "replaywebpage", + }, + "allow_worksheets": set(), + "block_worksheets": set(), + "use_sheet_names_in_stored_paths": True, + } + ) + feeder.gsheets_client = MagicMock() + return feeder + + +@pytest.fixture() +def worksheet(unpickle): + # Load the worksheet data from the pickle file + # only works for simple usage, cant reauthenticate but give structure + return unpickle("test_worksheet.pickle") + + +class TestWorksheet(): + """ + mimics the bits we need from gworksheet + """ + + class SheetSheet: + title = "TestSheet" + + rows = [ + { "row": 2, "url": "http://example.com", "status": "", "folder": "" }, + { "row": 3, "url": "http://example.com", "status": "", "folder": "" }, + { "row": 4, "url": "", "status": "", "folder": "" }, + { "row": 5, "url": "https://another.com", "status": None, "folder": "" }, + { "row": 6, "url": "https://another.com", "status": "success", "folder": "some_folder" }, + ] + + def __init__(self): + self.wks = self.SheetSheet() + + def count_rows(self): + if not self.rows: + return 0 + return max(r["row"] for r in self.rows) + + def get_cell(self, row, col_name, fresh=False): + matching = next((r for r in self.rows if r["row"] == row), {}) + return matching.get(col_name, "") + + def get_cell_or_default(self, row, col_name, default): + matching = next((r for r in self.rows if r["row"] == row), {}) + return matching.get(col_name, default) + +def test__process_rows(gsheet_feeder: GsheetsFeeder): + testworksheet = TestWorksheet() + metadata_items = list(gsheet_feeder._process_rows(testworksheet)) + assert len(metadata_items) == 3 + assert isinstance(metadata_items[0], Metadata) + assert metadata_items[0].get("url") == "http://example.com" + +def test__set_metadata(gsheet_feeder: GsheetsFeeder, worksheet): + gsheet_feeder._set_context(worksheet, 1) + assert ArchivingContext.get("gsheet") == {"row": 1, "worksheet": worksheet} + + +@pytest.mark.skip(reason="Not recognising folder column") +def test__set_metadata_with_folder_pickled(gsheet_feeder: GsheetsFeeder, worksheet): + gsheet_feeder._set_context(worksheet, 7) + assert ArchivingContext.get("gsheet") == {"row": 1, "worksheet": worksheet} + + +def test__set_metadata_with_folder(gsheet_feeder: GsheetsFeeder): + testworksheet = TestWorksheet() + testworksheet.wks.title = "TestSheet" + gsheet_feeder._set_context(testworksheet, 6) + assert ArchivingContext.get("gsheet") == {"row": 6, "worksheet": testworksheet} + assert ArchivingContext.get("folder") == "some-folder/test-auto-archiver/testsheet" + + +@pytest.mark.usefixtures("setup_module") +@pytest.mark.parametrize("sheet, sheet_id, expected_method, expected_arg, description", [ + ("TestSheet", None, "open", "TestSheet", "opening by sheet name"), + (None, "ABC123", "open_by_key", "ABC123", "opening by sheet ID") +]) +def test_open_sheet_with_name_or_id(setup_module, sheet, sheet_id, expected_method, expected_arg, description): + """Ensure open_sheet() correctly opens by name or ID based on configuration.""" + with patch("gspread.service_account") as mock_service_account: + mock_client = MagicMock() + mock_service_account.return_value = mock_client + mock_client.open.return_value = "MockSheet" + mock_client.open_by_key.return_value = "MockSheet" + + # Setup module with parameterized values + feeder = setup_module("gsheet_feeder", { + "service_account": "dummy.json", + "sheet": sheet, + "sheet_id": sheet_id + }) + feeder.initialise() + sheet_result = feeder.open_sheet() + # Validate the correct method was called + getattr(mock_client, expected_method).assert_called_once_with(expected_arg), f"Failed: {description}" + assert sheet_result == "MockSheet", f"Failed: {description}" + + +@pytest.mark.usefixtures("setup_module") +def test_open_sheet_with_sheet_id(setup_module): + """Ensure open_sheet() correctly opens a sheet by ID.""" + with patch("gspread.service_account") as mock_service_account: + mock_client = MagicMock() + mock_service_account.return_value = mock_client + mock_client.open_by_key.return_value = "MockSheet" + feeder = setup_module("gsheet_feeder", + {"service_account": "dummy.json", + "sheet": None, + "sheet_id": "ABC123"}) + feeder.initialise() + sheet = feeder.open_sheet() + mock_client.open_by_key.assert_called_once_with("ABC123") + assert sheet == "MockSheet" + + +def test_should_process_sheet(setup_module): + gdb = setup_module("gsheet_feeder", {"service_account": "dummy.json", + "sheet": "TestSheet", + "sheet_id": None, + "allow_worksheets": {"TestSheet", "Sheet2"}, + "block_worksheets": {"Sheet3"}} + ) + assert gdb.should_process_sheet("TestSheet") == True + assert gdb.should_process_sheet("Sheet3") == False + # False if allow_worksheets is set + assert gdb.should_process_sheet("AnotherSheet") == False + + + +@pytest.mark.skip +class TestGSheetsFeederReal: + + """ Testing GSheetsFeeder class """ + module_name: str = 'gsheet_feeder' + feeder: GsheetsFeeder + config: dict = { + # TODO: Create test creds + "service_account": "secrets/service_account.json", + "sheet": "test-auto-archiver", + "sheet_id": None, + "header": 1, + "columns": { + "url": "link", + "status": "archive status", + "folder": "destination folder", + "archive": "archive location", + "date": "archive date", + "thumbnail": "thumbnail", + "timestamp": "upload timestamp", + "title": "upload title", + "text": "text content", + "screenshot": "screenshot", + "hash": "hash", + "pdq_hash": "perceptual hashes", + "wacz": "wacz", + "replaywebpage": "replaywebpage", + }, + "allow_worksheets": set(), + "block_worksheets": set(), + "use_sheet_names_in_stored_paths": True, + } + + @pytest.fixture(autouse=True) + def setup_feeder(self, setup_module): + assert ( + self.module_name is not None + ), "self.module_name must be set on the subclass" + assert self.config is not None, "self.config must be a dict set on the subclass" + self.feeder: Type[Feeder] = setup_module( + self.module_name, self.config + ) + + def reset_test_sheet(self): + """Clears test sheet and re-adds headers to ensure consistent test results.""" + client = gspread.service_account(self.config["service_account"]) + sheet = client.open(self.config["sheet"]) + worksheet = sheet.get_worksheet(0) + worksheet.clear() + worksheet.append_row(["Link", "Archive Status"]) + + def test_initialise(self): + self.feeder.initialise() + assert hasattr(self.feeder, "gsheets_client") + + @pytest.mark.download + def test_open_sheet_real_connection(self): + """Ensure open_sheet() connects to a real Google Sheets instance.""" + self.feeder.initialise() + sheet = self.feeder.open_sheet() + assert sheet is not None, "open_sheet() should return a valid sheet instance" + assert hasattr(sheet, "worksheets"), "Returned object should have worksheets method" + + @pytest.mark.download + def test_iter_yields_metadata_real_data(self): + """Ensure __iter__() yields Metadata objects for real test sheet data.""" + self.reset_test_sheet() + client = gspread.service_account(self.config["service_account"]) + sheet = client.open(self.config["sheet"]) + worksheet = sheet.get_worksheet(0) + # Insert test rows as a temp method + # Next we will refactor the feeder for better testing + test_rows = [ + ["https://example.com", ""], + ["", ""], + ["https://example.com", "done"], + ] + worksheet.append_rows(test_rows) + self.feeder.initialise() + metadata_list = list(self.feeder) + + # Validate that only the first row is processed + assert len(metadata_list) == 1 + assert metadata_list[0].metadata.get("url") == "https://example.com" + + + +# TODO + +# Test two sheets +# test two sheets with different columns +# test folder implementation diff --git a/tests/feeders/test_gworksheet.py b/tests/feeders/test_gworksheet.py new file mode 100644 index 0000000..e6f5cc6 --- /dev/null +++ b/tests/feeders/test_gworksheet.py @@ -0,0 +1,144 @@ +import pytest +from unittest.mock import MagicMock + +from auto_archiver.modules.gsheet_feeder import GWorksheet + + +class TestGWorksheet: + @pytest.fixture + def mock_worksheet(self): + mock_ws = MagicMock() + mock_ws.get_values.return_value = [ + ["Link", "Archive Status", "Archive Location", "Archive Date"], + ["url1", "archived", "filepath1", "2023-01-01"], + ["url2", "pending", "filepath2", "2023-01-02"], + ] + return mock_ws + + @pytest.fixture + def gworksheet(self, mock_worksheet): + return GWorksheet(mock_worksheet) + + # Test initialization and basic properties + def test_initialization_sets_headers(self, gworksheet): + assert gworksheet.headers == ["link", "archive status", "archive location", "archive date"] + + def test_count_rows_returns_correct_value(self, gworksheet): + # inc header row + assert gworksheet.count_rows() == 3 + + # Test column validation and lookup + @pytest.mark.parametrize( + "col,expected_index", + [ + ("url", 0), + ("status", 1), + ("archive", 2), + ("date", 3), + ], + ) + def test_col_index_returns_correct_index(self, gworksheet, col, expected_index): + assert gworksheet._col_index(col) == expected_index + + def test_check_col_exists_raises_for_invalid_column(self, gworksheet): + with pytest.raises(Exception, match="Column invalid_col"): + gworksheet._check_col_exists("invalid_col") + + # Test data retrieval + @pytest.mark.parametrize( + "row,expected", + [ + (1, ["Link", "Archive Status", "Archive Location", "Archive Date"]), + (2, ["url1", "archived", "filepath1", "2023-01-01"]), + (3, ["url2", "pending", "filepath2", "2023-01-02"]), + ], + ) + def test_get_row_returns_correct_data(self, gworksheet, row, expected): + assert gworksheet.get_row(row) == expected + + @pytest.mark.parametrize( + "row,col,expected", + [ + (2, "url", "url1"), + (2, "status", "archived"), + (3, "date", "2023-01-02"), + ], + ) + def test_get_cell_returns_correct_value(self, gworksheet, row, col, expected): + assert gworksheet.get_cell(row, col) == expected + + def test_get_cell_handles_fresh_data(self, mock_worksheet, gworksheet): + mock_worksheet.cell.return_value.value = "fresh_value" + result = gworksheet.get_cell(2, "url", fresh=True) + assert result == "fresh_value" + mock_worksheet.cell.assert_called_once_with(2, 1) + + # Test edge cases and error handling + @pytest.mark.parametrize( + "when_empty,expected", + [ + (True, "default"), + (False, ""), + ], + ) + def test_get_cell_or_default_handles_empty_values( + self, mock_worksheet, when_empty, expected + ): + mock_worksheet.get_values.return_value[1][0] = "" # Empty URL cell + g = GWorksheet(mock_worksheet) + assert ( + g.get_cell_or_default( + 2, "url", default="default", when_empty_use_default=when_empty + ) + == expected + ) + + def test_get_cell_or_default_handles_missing_columns(self, gworksheet): + assert ( + gworksheet.get_cell_or_default(1, "invalid_col", default="safe") == "safe" + ) + + # Test write operations + def test_set_cell_updates_correct_position(self, mock_worksheet, gworksheet): + gworksheet.set_cell(2, "url", "new_url") + mock_worksheet.update_cell.assert_called_once_with(2, 1, "new_url") + + def test_batch_set_cell_formats_requests_correctly( + self, mock_worksheet, gworksheet + ): + updates = [(2, "url", "new_url"), (3, "status", "processed")] + gworksheet.batch_set_cell(updates) + expected_batch = [ + {"range": "A2", "values": [["new_url"]]}, + {"range": "B3", "values": [["processed"]]}, + ] + mock_worksheet.batch_update.assert_called_once_with( + expected_batch, value_input_option="USER_ENTERED" + ) + + def test_batch_set_cell_truncates_long_values(self, mock_worksheet, gworksheet): + long_value = "x" * 50000 + gworksheet.batch_set_cell([(1, "url", long_value)]) + submitted_value = mock_worksheet.batch_update.call_args[0][0][0]["values"][0][0] + assert len(submitted_value) == 49999 + + # Test coordinate conversion + @pytest.mark.parametrize( + "row,col,expected", + [ + (1, "url", "A1"), + (2, "status", "B2"), + (3, "archive", "C3"), + (4, "date", "D4"), + ], + ) + def test_to_a1_conversion(self, gworksheet, row, col, expected): + assert gworksheet.to_a1(row, col) == expected + + # Test empty worksheet + def test_empty_worksheet_initialization(self): + mock_ws = MagicMock() + mock_ws.get_values.return_value = [] + g = GWorksheet(mock_ws) + assert g.headers == [] + assert g.count_rows() == 0 diff --git a/tests/storages/test_S3_storage.py b/tests/storages/test_S3_storage.py new file mode 100644 index 0000000..df1c1f1 --- /dev/null +++ b/tests/storages/test_S3_storage.py @@ -0,0 +1,100 @@ +from typing import Type +import pytest +from unittest.mock import MagicMock, patch, mock_open +from auto_archiver.core import Media +from auto_archiver.modules.s3_storage import s3_storage +from tests.storages.test_storage_base import TestStorageBase + + +class TestGDriveStorage: + """ + Test suite for GDriveStorage. + """ + module_name: str = "s3_storage" + storage: Type[s3_storage] + s3: MagicMock + config: dict = { + "path_generator": "flat", + "filename_generator": "static", + "bucket": "test-bucket", + "region": "test-region", + "key": "test-key", + "secret": "test-secret", + "random_no_duplicate": False, + "endpoint_url": "https://{region}.example.com", + "cdn_url": "https://cdn.example.com/{key}", + "private": False, + } + + @patch('boto3.client') + @pytest.fixture(autouse=True) + def setup_storage(self, setup_module): + self.storage = setup_module(self.module_name, self.config) + self.storage.initialise() + + @patch('boto3.client') + def test_client_initialization(self, mock_boto_client, setup_module): + """Test that S3 client is initialized with correct parameters""" + self.storage.initialise() + mock_boto_client.assert_called_once_with( + 's3', + region_name='test-region', + endpoint_url='https://test-region.example.com', + aws_access_key_id='test-key', + aws_secret_access_key='test-secret' + ) + + def test_get_cdn_url_generation(self): + """Test CDN URL formatting """ + media = Media("test.txt") + media.key = "path/to/file.txt" + url = self.storage.get_cdn_url(media) + assert url == "https://cdn.example.com/path/to/file.txt" + media.key = "another/path.jpg" + assert self.storage.get_cdn_url(media) == "https://cdn.example.com/another/path.jpg" + + + @patch.object(s3_storage.S3Storage, 'file_in_folder') + def test_skips_upload_when_duplicate_exists(self, mock_file_in_folder): + """Test that upload skips when file_in_folder finds existing object""" + # Setup test-specific configuration + self.storage.random_no_duplicate = True + mock_file_in_folder.return_value = "existing_folder/existing_file.txt" + # Create test media with calculated hash + media = Media("test.txt") + media.key = "original_path.txt" + + # Mock hash calculation + with patch.object(self.storage, 'calculate_hash') as mock_calculate_hash: + mock_calculate_hash.return_value = "testhash123" + # Verify upload + assert self.storage.is_upload_needed(media) is False + assert media.key == "existing_folder/existing_file.txt" + assert media.get("previously archived") is True + + with patch.object(self.storage.s3, 'upload_fileobj') as mock_upload: + result = self.storage.uploadf(None, media) + mock_upload.assert_not_called() + assert result is True + + @patch.object(s3_storage.S3Storage, 'is_upload_needed') + def test_uploads_with_correct_parameters(self, mock_upload_needed): + media = Media("test.txt") + mock_upload_needed.return_value = True + media.mimetype = 'image/png' + mock_file = MagicMock() + + with patch.object(self.storage.s3, 'upload_fileobj') as mock_upload: + self.storage.uploadf(mock_file, media) + + # Verify core upload parameters + mock_upload.assert_called_once_with( + mock_file, + Bucket='test-bucket', + # Key='original_key.txt', + Key=None, + ExtraArgs={ + 'ACL': 'public-read', + 'ContentType': 'image/png' + } + ) \ No newline at end of file diff --git a/tests/storages/test_gdrive_storage.py b/tests/storages/test_gdrive_storage.py new file mode 100644 index 0000000..b7417ad --- /dev/null +++ b/tests/storages/test_gdrive_storage.py @@ -0,0 +1,43 @@ +from typing import Type +import pytest +from unittest.mock import MagicMock, patch +from auto_archiver.core import Media +from auto_archiver.modules.gdrive_storage import GDriveStorage +from auto_archiver.core.metadata import Metadata +from tests.storages.test_storage_base import TestStorageBase + + +class TestGDriveStorage(TestStorageBase): + """ + Test suite for GDriveStorage. + """ + + module_name: str = "gdrive_storage" + storage: Type[GDriveStorage] + config: dict = {'path_generator': 'url', + 'filename_generator': 'static', + 'root_folder_id': "fake_root_folder_id", + 'oauth_token': None, + 'service_account': 'fake_service_account.json' + } + + @pytest.mark.skip(reason="Requires real credentials") + @pytest.mark.download + def test_initialize_with_real_credentials(self): + """ + Test that the Google Drive service can be initialized with real credentials. + """ + self.storage.service_account = 'secrets/service_account.json' # Path to real credentials + self.storage.initialise() + assert self.storage.service is not None + + + def test_initialize_fails_with_non_existent_creds(self): + """ + Test that the Google Drive service raises a FileNotFoundError when the service account file does not exist. + """ + # Act and Assert + with pytest.raises(FileNotFoundError) as exc_info: + self.storage.initialise() + assert "No such file or directory" in str(exc_info.value) + diff --git a/tests/storages/test_storage_base.py b/tests/storages/test_storage_base.py new file mode 100644 index 0000000..50d8846 --- /dev/null +++ b/tests/storages/test_storage_base.py @@ -0,0 +1,23 @@ +from typing import Type + +import pytest + +from auto_archiver.core.context import ArchivingContext +from auto_archiver.core.metadata import Metadata +from auto_archiver.core.storage import Storage + + +class TestStorageBase(object): + + module_name: str = None + config: dict = None + + @pytest.fixture(autouse=True) + def setup_storage(self, setup_module): + assert ( + self.module_name is not None + ), "self.module_name must be set on the subclass" + assert self.config is not None, "self.config must be a dict set on the subclass" + self.storage: Type[Storage] = setup_module( + self.module_name, self.config + ) From 6ab8fd2ee49d5ec4c1bc5f7e22cddf732a0e9371 Mon Sep 17 00:00:00 2001 From: Patrick Robertson Date: Wed, 5 Feb 2025 20:39:53 +0100 Subject: [PATCH 26/62] Tidy up setting modules as Orchestrator attributes on startup. MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Don't override the values in config['steps'] – the config should be left as is --- src/auto_archiver/core/orchestrator.py | 52 ++++++-------------------- tests/test_orchestrator.py | 1 - 2 files changed, 11 insertions(+), 42 deletions(-) diff --git a/src/auto_archiver/core/orchestrator.py b/src/auto_archiver/core/orchestrator.py index 8a634de..5ac091c 100644 --- a/src/auto_archiver/core/orchestrator.py +++ b/src/auto_archiver/core/orchestrator.py @@ -69,6 +69,13 @@ class UniqueAppendAction(argparse.Action): getattr(namespace, self.dest).append(value) class ArchivingOrchestrator: + + feeders: List[Type[Feeder]] + extractors: List[Type[Extractor]] + enrichers: List[Type[Enricher]] + databases: List[Type[Database]] + storages: List[Type[Storage]] + formatters: List[Type[Formatter]] def setup_basic_parser(self): parser = argparse.ArgumentParser( @@ -296,11 +303,7 @@ class ArchivingOrchestrator: step_items.append(loaded_module) check_steps_ok() - self.config['steps'][f"{module_type}s"] = step_items - - - assert len(step_items) > 0, f"No {module_type}s were loaded. Please check your configuration file and try again." - self.config['steps'][f"{module_type}s"] = step_items + setattr(self, f"{module_type}s", step_items) def load_config(self, config_file: str) -> dict: if not os.path.exists(config_file) and config_file != DEFAULT_CONFIG_FILE: @@ -331,9 +334,9 @@ class ArchivingOrchestrator: # log out the modules that were loaded for module_type in BaseModule.MODULE_TYPES: - logger.info(f"{module_type.upper()}S: " + ", ".join(m.display_name for m in self.config['steps'][f"{module_type}s"])) + logger.info(f"{module_type.upper()}S: " + ", ".join(m.display_name for m in getattr(self, f"{module_type}s"))) - for item in self.feed(): + for _ in self.feed(): pass def cleanup(self)->None: @@ -484,40 +487,7 @@ class ArchivingOrchestrator: # Helper Properties - - @property - def feeders(self) -> List[Type[Feeder]]: - return self._get_property('feeders') - - @property - def extractors(self) -> List[Type[Extractor]]: - return self._get_property('extractors') - - @property - def enrichers(self) -> List[Type[Enricher]]: - return self._get_property('enrichers') - - @property - def databases(self) -> List[Type[Database]]: - return self._get_property('databases') - - @property - def storages(self) -> List[Type[Storage]]: - return self._get_property('storages') - - @property - def formatters(self) -> List[Type[Formatter]]: - return self._get_property('formatters') @property def all_modules(self) -> List[Type[BaseModule]]: - return self.feeders + self.extractors + self.enrichers + self.databases + self.storages + self.formatters - - def _get_property(self, prop): - try: - f = self.config['steps'][prop] - if not (isinstance(f[0], BaseModule) or isinstance(f[0], LazyBaseModule)): - raise TypeError - return f - except: - exit("Property called prior to full initialisation") \ No newline at end of file + return self.feeders + self.extractors + self.enrichers + self.databases + self.storages + self.formatters \ No newline at end of file diff --git a/tests/test_orchestrator.py b/tests/test_orchestrator.py index 68417aa..5ba57d0 100644 --- a/tests/test_orchestrator.py +++ b/tests/test_orchestrator.py @@ -89,7 +89,6 @@ def test_add_custom_modules_path_invalid(orchestrator, caplog, test_args): orchestrator.run(test_args + # we still need to load the real path to get the example_module ["--module_paths", "tests/data/invalid_test_modules/"]) - # assert False assert caplog.records[0].message == "Path 'tests/data/invalid_test_modules/' does not exist. Skipping..." From a506f2a88f6e073a2c45367a0f4cb8bab1f79bb7 Mon Sep 17 00:00:00 2001 From: Patrick Robertson Date: Thu, 6 Feb 2025 10:19:28 +0100 Subject: [PATCH 27/62] Clarify that an extractor's method can also return False if no valid data was found --- src/auto_archiver/core/extractor.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/auto_archiver/core/extractor.py b/src/auto_archiver/core/extractor.py index 98f1370..57320df 100644 --- a/src/auto_archiver/core/extractor.py +++ b/src/auto_archiver/core/extractor.py @@ -95,5 +95,11 @@ class Extractor(BaseModule): logger.warning(f"Failed to fetch the Media URL: {e}") @abstractmethod - def download(self, item: Metadata) -> Metadata: + def download(self, item: Metadata) -> Metadata | False: + """ + Downloads the media from the given URL and returns a Metadata object with the downloaded media. + + If the URL is not supported or the download fails, this method should return False. + + """ pass \ No newline at end of file From 5b0bad832f0bcf787979f18c5b8027f10b95b0a6 Mon Sep 17 00:00:00 2001 From: erinhmclark Date: Thu, 6 Feb 2025 10:11:56 +0000 Subject: [PATCH 28/62] Updated test, test metadata --- .../modules/gsheet_db/gsheet_db.py | 1 - .../modules/gsheet_feeder/gsheet_feeder.py | 59 ++++--- .../test_instagram_api_extractor.py | 89 +++++++++- tests/feeders/test_gsheet_feeder.py | 10 +- tests/test_metadata.py | 161 ++++++++++++++++++ 5 files changed, 284 insertions(+), 36 deletions(-) create mode 100644 tests/test_metadata.py diff --git a/src/auto_archiver/modules/gsheet_db/gsheet_db.py b/src/auto_archiver/modules/gsheet_db/gsheet_db.py index 644015e..3bb27b7 100644 --- a/src/auto_archiver/modules/gsheet_db/gsheet_db.py +++ b/src/auto_archiver/modules/gsheet_db/gsheet_db.py @@ -104,7 +104,6 @@ class GsheetsDb(Database): if gsheet := item.get_context("gsheet"): gw: GWorksheet = gsheet.get("worksheet") row: int = gsheet.get("row") - # todo doesn't exist, should be passed from elif self.sheet_id: print(self.sheet_id) diff --git a/src/auto_archiver/modules/gsheet_feeder/gsheet_feeder.py b/src/auto_archiver/modules/gsheet_feeder/gsheet_feeder.py index d129182..a51574e 100644 --- a/src/auto_archiver/modules/gsheet_feeder/gsheet_feeder.py +++ b/src/auto_archiver/modules/gsheet_feeder/gsheet_feeder.py @@ -37,41 +37,48 @@ class GsheetsFeeder(Feeder): def __iter__(self) -> Metadata: sh = self.open_sheet() - for ii, wks in enumerate(sh.worksheets()): - if not self.should_process_sheet(wks.title): - logger.debug(f"SKIPPED worksheet '{wks.title}' due to allow/block rules") + for ii, worksheet in enumerate(sh.worksheets()): + if not self.should_process_sheet(worksheet.title): + logger.debug(f"SKIPPED worksheet '{worksheet.title}' due to allow/block rules") continue - - logger.info(f'Opening worksheet {ii=}: {wks.title=} header={self.header}') - gw = GWorksheet(wks, header_row=self.header, columns=self.columns) - + logger.info(f'Opening worksheet {ii=}: {worksheet.title=} header={self.header}') + gw = GWorksheet(worksheet, header_row=self.header, columns=self.columns) if len(missing_cols := self.missing_required_columns(gw)): logger.warning(f"SKIPPED worksheet '{wks.title}' due to missing required column(s) for {missing_cols}") continue - for row in range(1 + self.header, gw.count_rows() + 1): - url = gw.get_cell(row, 'url').strip() - if not len(url): continue + # process and yield metadata here: + yield from self._process_rows(gw) + logger.success(f'Finished worksheet {worksheet.title}') - original_status = gw.get_cell(row, 'status') - status = gw.get_cell(row, 'status', fresh=original_status in ['', None]) - # TODO: custom status parser(?) aka should_retry_from_status - if status not in ['', None]: continue + def _process_rows(self, gw: GWorksheet) -> Metadata: + for row in range(1 + self.header, gw.count_rows() + 1): + url = gw.get_cell(row, 'url').strip() + if not len(url): continue + original_status = gw.get_cell(row, 'status') + status = gw.get_cell(row, 'status', fresh=original_status in ['', None]) + # TODO: custom status parser(?) aka should_retry_from_status + if status not in ['', None]: continue - # All checks done - archival process starts here - m = Metadata().set_url(url) - if gw.get_cell_or_default(row, 'folder', "") is None: - folder = '' - else: - folder = slugify(gw.get_cell_or_default(row, 'folder', "").strip()) - if len(folder) and self.use_sheet_names_in_stored_paths: - folder = os.path.join(folder, slugify(self.sheet), slugify(wks.title)) + # All checks done - archival process starts here + m = Metadata().set_url(url) + self._set_context(m, gw, row) + yield m - m.set_context('folder', folder) - m.set_context('worksheet', {"row": row, "worksheet": gw}) - yield m + def _set_context(self, m: Metadata, gw: GWorksheet, row: int) -> Metadata: + # TODO: Check folder value not being recognised + m.set_context("gsheet", {"row": row, "worksheet": gw}) + + if gw.get_cell_or_default(row, 'folder', "") is None: + folder = '' + else: + folder = slugify(gw.get_cell_or_default(row, 'folder', "").strip()) + if len(folder): + if self.use_sheet_names_in_stored_paths: + m.set_context("folder", os.path.join(folder, slugify(self.sheet), slugify(gw.wks.title))) + else: + m.set_context("folder", folder) - logger.success(f'Finished worksheet {wks.title}') def should_process_sheet(self, sheet_name: str) -> bool: if len(self.allow_worksheets) and sheet_name not in self.allow_worksheets: diff --git a/tests/extractors/test_instagram_api_extractor.py b/tests/extractors/test_instagram_api_extractor.py index 7a19233..d3f7bd6 100644 --- a/tests/extractors/test_instagram_api_extractor.py +++ b/tests/extractors/test_instagram_api_extractor.py @@ -9,6 +9,7 @@ from auto_archiver.modules.instagram_api_extractor.instagram_api_extractor impor from .test_extractor_base import TestExtractorBase + @pytest.fixture def mock_user_response(): return { @@ -71,11 +72,18 @@ class TestInstagramAPIExtractor(TestExtractorBase): config = { "access_token": "test_access_token", "api_endpoint": "https://api.instagram.com/v1", - # "full_profile": False, + "full_profile": False, # "full_profile_max_posts": 0, # "minimize_json_output": True, } + @pytest.fixture + def metadata(self): + m = Metadata() + m.set_url("https://instagram.com/test_user") + m.set("netloc", "instagram.com") + return m + @pytest.mark.parametrize("url,expected", [ ("https://instagram.com/user", [("", "user", "")]), ("https://instagr.am/p/post_id", []), @@ -88,7 +96,6 @@ class TestInstagramAPIExtractor(TestExtractorBase): assert self.extractor.valid_url.findall(url) == expected def test_initialize(self): - self.extractor.initialise() assert self.extractor.api_endpoint[-1] != "/" @pytest.mark.parametrize("input_dict,expected", [ @@ -98,11 +105,85 @@ class TestInstagramAPIExtractor(TestExtractorBase): def test_cleanup_dict(self, input_dict, expected): assert self.extractor.cleanup_dict(input_dict) == expected - def test_download_post(self): + def test_download(self): + pass + + def test_download_post(self, metadata, mock_user_response): # test with context=reel # test with context=post # test with multiple images # test gets text (metadata title) + pass + def test_download_profile_basic(self, metadata, mock_user_response): + """Test basic profile download without full_profile""" + with patch.object(self.extractor, 'call_api') as mock_call, \ + patch.object(self.extractor, 'download_from_url') as mock_download: + # Mock API responses + mock_call.return_value = mock_user_response + mock_download.return_value = "profile.jpg" - pass \ No newline at end of file + result = self.extractor.download_profile(metadata, "test_user") + assert result.status == "insta profile: success" + assert result.get_title() == "Test User" + assert result.get("data") == self.extractor.cleanup_dict(mock_user_response["user"]) + # Verify profile picture download + mock_call.assert_called_once_with("v2/user/by/username", {"username": "test_user"}) + mock_download.assert_called_once_with("http://example.com/profile.jpg") + assert len(result.media) == 1 + assert result.media[0].filename == "profile.jpg" + + def test_download_profile_full(self, metadata, mock_user_response, mock_story_response): + """Test full profile download with stories/posts""" + with patch.object(self.extractor, 'call_api') as mock_call, \ + patch.object(self.extractor, 'download_all_posts') as mock_posts, \ + patch.object(self.extractor, 'download_all_highlights') as mock_highlights, \ + patch.object(self.extractor, 'download_all_tagged') as mock_tagged, \ + patch.object(self.extractor, '_download_stories_reusable') as mock_stories: + + self.extractor.full_profile = True + mock_call.side_effect = [ + mock_user_response, + mock_story_response + ] + mock_highlights.return_value = None + mock_stories.return_value = mock_story_response + mock_posts.return_value = None + mock_tagged.return_value = None + + result = self.extractor.download_profile(metadata, "test_user") + assert result.get("#stories") == len(mock_story_response) + mock_posts.assert_called_once_with(result, "123") + assert "errors" not in result.metadata + + def test_download_profile_not_found(self, metadata): + """Test profile not found error""" + with patch.object(self.extractor, 'call_api') as mock_call: + mock_call.return_value = {"user": None} + with pytest.raises(AssertionError) as exc_info: + self.extractor.download_profile(metadata, "invalid_user") + assert "User invalid_user not found" in str(exc_info.value) + + def test_download_profile_error_handling(self, metadata, mock_user_response): + """Test error handling in full profile mode""" + with (patch.object(self.extractor, 'call_api') as mock_call, \ + patch.object(self.extractor, 'download_all_highlights') as mock_highlights, \ + patch.object(self.extractor, 'download_all_tagged') as mock_tagged, \ + patch.object(self.extractor, '_download_stories_reusable') as stories_tagged, \ + patch.object(self.extractor, 'download_all_posts') as mock_posts + ): + self.extractor.full_profile = True + mock_call.side_effect = [ + mock_user_response, + Exception("Stories API failed"), + Exception("Posts API failed") + ] + mock_highlights.return_value = None + mock_tagged.return_value = None + stories_tagged.return_value = None + mock_posts.return_value = None + result = self.extractor.download_profile(metadata, "test_user") + + assert result.is_success() + assert "Error downloading stories for test_user" in result.metadata["errors"] + # assert "Error downloading posts for test_user" in result.metadata["errors"] \ No newline at end of file diff --git a/tests/feeders/test_gsheet_feeder.py b/tests/feeders/test_gsheet_feeder.py index dbd2416..62380f5 100644 --- a/tests/feeders/test_gsheet_feeder.py +++ b/tests/feeders/test_gsheet_feeder.py @@ -4,7 +4,7 @@ import gspread import pytest from unittest.mock import patch, MagicMock from auto_archiver.modules.gsheet_feeder import GsheetsFeeder -from auto_archiver.core import Metadata, Feeder, ArchivingContext +from auto_archiver.core import Metadata, Feeder def test_initialise_without_sheet_and_sheet_id(setup_module): @@ -100,21 +100,21 @@ def test__process_rows(gsheet_feeder: GsheetsFeeder): def test__set_metadata(gsheet_feeder: GsheetsFeeder, worksheet): gsheet_feeder._set_context(worksheet, 1) - assert ArchivingContext.get("gsheet") == {"row": 1, "worksheet": worksheet} + assert Metadata.get_context("gsheet") == {"row": 1, "worksheet": worksheet} @pytest.mark.skip(reason="Not recognising folder column") def test__set_metadata_with_folder_pickled(gsheet_feeder: GsheetsFeeder, worksheet): gsheet_feeder._set_context(worksheet, 7) - assert ArchivingContext.get("gsheet") == {"row": 1, "worksheet": worksheet} + assert Metadata.get_context("gsheet") == {"row": 1, "worksheet": worksheet} def test__set_metadata_with_folder(gsheet_feeder: GsheetsFeeder): testworksheet = TestWorksheet() testworksheet.wks.title = "TestSheet" gsheet_feeder._set_context(testworksheet, 6) - assert ArchivingContext.get("gsheet") == {"row": 6, "worksheet": testworksheet} - assert ArchivingContext.get("folder") == "some-folder/test-auto-archiver/testsheet" + assert Metadata.get_context("gsheet") == {"row": 6, "worksheet": testworksheet} + assert Metadata.get_context("folder") == "some-folder/test-auto-archiver/testsheet" @pytest.mark.usefixtures("setup_module") diff --git a/tests/test_metadata.py b/tests/test_metadata.py new file mode 100644 index 0000000..7270c80 --- /dev/null +++ b/tests/test_metadata.py @@ -0,0 +1,161 @@ +import pytest +from datetime import datetime, timezone +from dataclasses import dataclass +from typing import Any +from auto_archiver.core.metadata import Metadata + + +@pytest.fixture +def basic_metadata(): + m = Metadata() + m.set_url("https://example.com") + m.set("title", "Test Page") + return m + + +@dataclass +class MockMedia: + filename: str = "" + mimetype: str = "" + data: dict = None + + def get(self, key: str, default: Any = None) -> Any: + return self.data.get(key, default) if self.data else default + + def set(self, key: str, value: Any) -> None: + if not self.data: + self.data = {} + self.data[key] = value + + +@pytest.fixture +def media_file(): + def _create(filename="test.txt", mimetype="text/plain", hash_value=None): + m = MockMedia(filename=filename, mimetype=mimetype) + if hash_value: + m.set("hash", hash_value) + return m + + return _create + + +def test_initial_state(): + m = Metadata() + assert m.status == "no archiver" + assert m.metadata == {"_processed_at": m.get("_processed_at")} + assert m.media == [] + assert isinstance(m.get("_processed_at"), datetime) + + +def test_url_properties(basic_metadata): + assert basic_metadata.get_url() == "https://example.com" + assert basic_metadata.netloc == "example.com" + + +def test_simple_merge(basic_metadata): + right = Metadata(status="success") + right.set("title", "Test Title") + + basic_metadata.merge(right) + assert basic_metadata.status == "success" + assert basic_metadata.get("title") == "Test Title" + + +def test_left_merge(): + left = ( + Metadata() + .set("tags", ["a"]) + .set("stats", {"views": 10}) + .set("status", "success") + ) + right = ( + Metadata() + .set("tags", ["b"]) + .set("stats", {"likes": 5}) + .set("status", "no archiver") + ) + + left.merge(right, overwrite_left=True) + assert left.get("status") == "no archiver" + assert left.get("tags") == ["a", "b"] + assert left.get("stats") == {"views": 10, "likes": 5} + + +def test_media_management(basic_metadata, media_file): + media1 = media_file(hash_value="abc") + media2 = media_file(hash_value="abc") # Duplicate + media3 = media_file(hash_value="def") + + basic_metadata.add_media(media1, "m1") + basic_metadata.add_media(media2, "m2") + basic_metadata.add_media(media3) + + assert len(basic_metadata.media) == 3 + basic_metadata.remove_duplicate_media_by_hash() + assert len(basic_metadata.media) == 2 + assert basic_metadata.get_media_by_id("m1") == media1 + + +def test_success(): + m = Metadata() + assert not m.is_success() + m.success("context") + assert m.is_success() + assert m.status == "context: success" + + +def test_is_empty(): + m = Metadata() + assert m.is_empty() + # meaningless ids + ( + m.set("url", "example.com") + .set("total_bytes", 100) + .set("archive_duration_seconds", 10) + .set("_processed_at", datetime.now(timezone.utc)) + ) + assert m.is_empty() + + +def test_store(): + pass + +# Test Media operations + + +# Test custom getter/setters + + +def test_get_set_url(): + m = Metadata() + m.set_url("http://example.com") + assert m.get_url() == "http://example.com" + with pytest.raises(AssertionError): + m.set_url("") + assert m.get("url") == "http://example.com" + + +def test_set_content(): + m = Metadata() + m.set_content("Some content") + assert m.get("content") == "Some content" + # Test appending + m.set_content("New content") + # Do we want to add a line break to the method? + assert m.get("content") == "Some contentNew content" + + +def test_choose_most_complex(): + pass + + +def test_get_context(): + m = Metadata() + m.set_context("somekey", "somevalue") + assert m.get_context("somekey") == "somevalue" + assert m.get_context("nonexistent") is None + m.set_context("anotherkey", "anothervalue") + # check the previous is retained + assert m.get_context("somekey") == "somevalue" + assert m.get_context("anotherkey") == "anothervalue" + assert len(m._context) == 2 From 266c7a14e6606cfd1c478cb4ed0ece602646035d Mon Sep 17 00:00:00 2001 From: erinhmclark Date: Thu, 6 Feb 2025 16:53:00 +0000 Subject: [PATCH 29/62] Context related fixes, some more tests. --- .../modules/gsheet_feeder/gsheet_feeder.py | 4 +- .../modules/s3_storage/__manifest__.py | 3 +- .../modules/s3_storage/s3_storage.py | 6 +- src/auto_archiver/utils/gsheet.py | 53 ----- tests/enrichers/test_meta_enricher.py | 103 +++++++++ .../test_instagram_tbot_extractor.py | 88 +++---- tests/feeders/test_gsheet_feeder.py | 216 +++++++++--------- tests/storages/test_S3_storage.py | 123 ++++++++-- tests/storages/test_storage_base.py | 1 - 9 files changed, 370 insertions(+), 227 deletions(-) delete mode 100644 src/auto_archiver/utils/gsheet.py create mode 100644 tests/enrichers/test_meta_enricher.py diff --git a/src/auto_archiver/modules/gsheet_feeder/gsheet_feeder.py b/src/auto_archiver/modules/gsheet_feeder/gsheet_feeder.py index a51574e..50bf430 100644 --- a/src/auto_archiver/modules/gsheet_feeder/gsheet_feeder.py +++ b/src/auto_archiver/modules/gsheet_feeder/gsheet_feeder.py @@ -44,14 +44,14 @@ class GsheetsFeeder(Feeder): logger.info(f'Opening worksheet {ii=}: {worksheet.title=} header={self.header}') gw = GWorksheet(worksheet, header_row=self.header, columns=self.columns) if len(missing_cols := self.missing_required_columns(gw)): - logger.warning(f"SKIPPED worksheet '{wks.title}' due to missing required column(s) for {missing_cols}") + logger.warning(f"SKIPPED worksheet '{worksheet.title}' due to missing required column(s) for {missing_cols}") continue # process and yield metadata here: yield from self._process_rows(gw) logger.success(f'Finished worksheet {worksheet.title}') - def _process_rows(self, gw: GWorksheet) -> Metadata: + def _process_rows(self, gw: GWorksheet): for row in range(1 + self.header, gw.count_rows() + 1): url = gw.get_cell(row, 'url').strip() if not len(url): continue diff --git a/src/auto_archiver/modules/s3_storage/__manifest__.py b/src/auto_archiver/modules/s3_storage/__manifest__.py index df05055..bf032e7 100644 --- a/src/auto_archiver/modules/s3_storage/__manifest__.py +++ b/src/auto_archiver/modules/s3_storage/__manifest__.py @@ -3,7 +3,7 @@ "type": ["storage"], "requires_setup": True, "dependencies": { - "python": ["boto3", "loguru"], + "python": ["hash_enricher", "boto3", "loguru"], }, "configs": { "path_generator": { @@ -49,5 +49,6 @@ - Requires S3 credentials (API key and secret) and a bucket name to function. - The `random_no_duplicate` option ensures no duplicate uploads by leveraging hash-based folder structures. - Uses `boto3` for interaction with the S3 API. + - Depends on the `HashEnricher` module for hash calculation. """ } diff --git a/src/auto_archiver/modules/s3_storage/s3_storage.py b/src/auto_archiver/modules/s3_storage/s3_storage.py index f324d5c..0c0e275 100644 --- a/src/auto_archiver/modules/s3_storage/s3_storage.py +++ b/src/auto_archiver/modules/s3_storage/s3_storage.py @@ -9,10 +9,11 @@ from auto_archiver.core import Media from auto_archiver.core import Storage from auto_archiver.modules.hash_enricher import HashEnricher from auto_archiver.utils.misc import random_str +from auto_archiver.core.module import get_module NO_DUPLICATES_FOLDER = "no-dups/" -class S3Storage(Storage, HashEnricher): +class S3Storage(Storage): def setup(self, config: dict) -> None: super().setup(config) @@ -49,7 +50,8 @@ class S3Storage(Storage, HashEnricher): def is_upload_needed(self, media: Media) -> bool: if self.random_no_duplicate: # checks if a folder with the hash already exists, if so it skips the upload - hd = self.calculate_hash(media.filename) + he = get_module('hash_enricher', self.config) + hd = he.calculate_hash(media.filename) path = os.path.join(NO_DUPLICATES_FOLDER, hd[:24]) if existing_key:=self.file_in_folder(path): diff --git a/src/auto_archiver/utils/gsheet.py b/src/auto_archiver/utils/gsheet.py deleted file mode 100644 index 7a8862f..0000000 --- a/src/auto_archiver/utils/gsheet.py +++ /dev/null @@ -1,53 +0,0 @@ -import json, gspread - -from ..core import BaseModule - - -class Gsheets(BaseModule): - name = "gsheets" - - def __init__(self, config: dict) -> None: - # without this STEP.__init__ is not called - super().__init__(config) - self.gsheets_client = gspread.service_account(filename=self.service_account) - # TODO: config should be responsible for conversions - try: self.header = int(self.header) - except: pass - assert type(self.header) == int, f"header ({self.header}) value must be an integer not {type(self.header)}" - assert self.sheet is not None or self.sheet_id is not None, "You need to define either a 'sheet' name or a 'sheet_id' in your orchestration file when using gsheets." - - # TODO merge this into gsheets processors manifest - @staticmethod - def configs() -> dict: - return { - "sheet": {"default": None, "help": "name of the sheet to archive"}, - "sheet_id": {"default": None, "help": "(alternative to sheet name) the id of the sheet to archive"}, - "header": {"default": 1, "help": "index of the header row (starts at 1)"}, - "service_account": {"default": "secrets/service_account.json", "help": "service account JSON file path"}, - "columns": { - "default": { - 'url': 'link', - 'status': 'archive status', - 'folder': 'destination folder', - 'archive': 'archive location', - 'date': 'archive date', - 'thumbnail': 'thumbnail', - 'timestamp': 'upload timestamp', - 'title': 'upload title', - 'text': 'text content', - 'screenshot': 'screenshot', - 'hash': 'hash', - 'pdq_hash': 'perceptual hashes', - 'wacz': 'wacz', - 'replaywebpage': 'replaywebpage', - }, - "help": "names of columns in the google sheet (stringified JSON object)", - "cli_set": lambda cli_val, cur_val: dict(cur_val, **json.loads(cli_val)) - }, - } - - def open_sheet(self): - if self.sheet: - return self.gsheets_client.open(self.sheet) - else: # self.sheet_id - return self.gsheets_client.open_by_key(self.sheet_id) diff --git a/tests/enrichers/test_meta_enricher.py b/tests/enrichers/test_meta_enricher.py new file mode 100644 index 0000000..a09aaa9 --- /dev/null +++ b/tests/enrichers/test_meta_enricher.py @@ -0,0 +1,103 @@ +import datetime +from datetime import datetime, timedelta, timezone +from unittest.mock import MagicMock, patch + +import pytest + +from auto_archiver.core import Metadata, Media +from auto_archiver.modules.meta_enricher import MetaEnricher + + +@pytest.fixture +def mock_metadata(): + """Creates a mock Metadata object.""" + mock: Metadata = MagicMock(spec=Metadata) + mock.get_url.return_value = "https://example.com" + mock.is_empty.return_value = False # Default to not empty + mock.get_all_media.return_value = [] + return mock + +@pytest.fixture +def mock_media(): + """Creates a mock Media object.""" + mock: Media = MagicMock(spec=Media) + mock.filename = "mock_file.txt" + return mock + +@pytest.fixture +def metadata(): + m = Metadata() + m.set_url("https://example.com") + m.set_title("Test Title") + m.set_content("Test Content") + return m + + +@pytest.fixture(autouse=True) +def meta_enricher(setup_module): + return setup_module(MetaEnricher, {}) + + +def test_enrich_skips_empty_metadata(meta_enricher, mock_metadata): + """Test that enrich() does nothing when Metadata is empty.""" + mock_metadata.is_empty.return_value = True + meta_enricher.enrich(mock_metadata) + mock_metadata.get_url.assert_called_once() + + +def test_enrich_file_sizes(meta_enricher, metadata, tmp_path): + """Test that enrich_file_sizes() calculates and sets file sizes correctly.""" + file1 = tmp_path / "testfile_1.txt" + file2 = tmp_path / "testfile_2.txt" + file1.write_text("A" * 1000) + file2.write_text("B" * 2000) + metadata.add_media(Media(str(file1))) + metadata.add_media(Media(str(file2))) + + meta_enricher.enrich_file_sizes(metadata) + + # Verify individual media file sizes + media1 = metadata.get_all_media()[0] + media2 = metadata.get_all_media()[1] + + assert media1.get("bytes") == 1000 + assert media1.get("size") == "1000.0 bytes" + assert media2.get("bytes") == 2000 + assert media2.get("size") == "2.0 KB" + + assert metadata.get("total_bytes") == 3000 + assert metadata.get("total_size") == "2.9 KB" + +@pytest.mark.parametrize( + "size, expected", + [ + (500, "500.0 bytes"), + (1024, "1.0 KB"), + (2048, "2.0 KB"), + (1048576, "1.0 MB"), + (1073741824, "1.0 GB"), + ], +) +def test_human_readable_bytes(size, expected): + """Test that human_readable_bytes() converts sizes correctly.""" + enricher = MetaEnricher() + assert enricher.human_readable_bytes(size) == expected + +def test_enrich_file_sizes_no_media(meta_enricher, metadata): + """Test that enrich_file_sizes() handles empty media list gracefully.""" + meta_enricher.enrich_file_sizes(metadata) + assert metadata.get("total_bytes") == 0 + assert metadata.get("total_size") == "0.0 bytes" + + +def test_enrich_archive_duration(meta_enricher, metadata): + # Set fixed "processed at" time in the past + processed_at = datetime.now(timezone.utc) - timedelta(minutes=10, seconds=30) + metadata.set("_processed_at", processed_at) + # patch datetime + with patch("datetime.datetime") as mock_datetime: + mock_now = datetime.now(timezone.utc) + mock_datetime.now.return_value = mock_now + meta_enricher.enrich_archive_duration(metadata) + + assert metadata.get("archive_duration_seconds") == 630 \ No newline at end of file diff --git a/tests/extractors/test_instagram_tbot_extractor.py b/tests/extractors/test_instagram_tbot_extractor.py index 4fe80be..b82641d 100644 --- a/tests/extractors/test_instagram_tbot_extractor.py +++ b/tests/extractors/test_instagram_tbot_extractor.py @@ -5,15 +5,16 @@ from unittest.mock import patch, MagicMock import pytest +from auto_archiver.core import Metadata from auto_archiver.core.extractor import Extractor from auto_archiver.modules.instagram_tbot_extractor import InstagramTbotExtractor - +from tests.extractors.test_extractor_base import TestExtractorBase TESTFILES = os.path.join(os.path.dirname(__file__), "testfiles") @pytest.fixture -def test_session_file(tmpdir): +def session_file(tmpdir): """Fixture to create a test session file.""" session_file = os.path.join(tmpdir, "test_session.session") with open(session_file, "w") as f: @@ -21,27 +22,34 @@ def test_session_file(tmpdir): return session_file.replace(".session", "") -@pytest.mark.incremental -class TestInstagramTbotExtractor(object): - """ - Test suite for InstagramTbotExtractor. - """ +@pytest.fixture(autouse=True) +def patch_extractor_methods(request, setup_module): + with patch.object(InstagramTbotExtractor, '_prepare_session_file', return_value=None), \ + patch.object(InstagramTbotExtractor, '_initialize_telegram_client', return_value=None): + if hasattr(request, 'cls') and hasattr(request.cls, 'config'): + request.cls.extractor = setup_module("instagram_tbot_extractor", request.cls.config) + + yield + +@pytest.fixture +def metadata_sample(): + m = Metadata() + m.set_title("Test Title") + m.set_timestamp("2021-01-01T00:00:00Z") + m.set_url("https://www.instagram.com/p/1234567890") + return m + + +class TestInstagramTbotExtractor: extractor_module = "instagram_tbot_extractor" extractor: InstagramTbotExtractor config = { "api_id": 12345, "api_hash": "test_api_hash", - # "session_file" + "session_file": "test_session", } - @pytest.fixture(autouse=True) - def setup_extractor(self, setup_module): - assert self.extractor_module is not None, "self.extractor_module must be set on the subclass" - assert self.config is not None, "self.config must be a dict set on the subclass" - extractor: Type[Extractor] = setup_module(self.extractor_module, self.config) - return extractor - @pytest.fixture def mock_telegram_client(self): """Fixture to mock TelegramClient interactions.""" @@ -50,22 +58,11 @@ class TestInstagramTbotExtractor(object): mock_client.return_value = instance yield instance - - # @pytest.fixture - # def mock_session_file(self, temp_session_file): - # """Patch the extractor’s session file setup to use a temporary path.""" - # with patch.object(InstagramTbotExtractor, "session_file", temp_session_file): - # with patch.object(InstagramTbotExtractor, "_prepare_session_file", return_value=None): - # yield # Mocks are applied for the duration of the test - - @pytest.fixture - def metadata_sample(self): - """Loads a Metadata object from a pickle file.""" - with open(os.path.join(TESTFILES, "metadata_item.pkl"), "rb") as f: - return pickle.load(f) + def test_extractor_is_initialized(self): + assert self.extractor is not None - @pytest.mark.download + @patch("time.sleep") @pytest.mark.parametrize("url, expected_status, bot_responses", [ ("https://www.instagram.com/p/C4QgLbrIKXG", "insta-via-bot: success", [MagicMock(id=101, media=None, message="Are you new to Bellingcat? - The way we share our investigations is different. 💭\nWe want you to read our story but also learn ou")]), ("https://www.instagram.com/reel/DEVLK8qoIbg/", "insta-via-bot: success", [MagicMock(id=101, media=None, message="Our volunteer community is at the centre of many incredible Bellingcat investigations and tools. Stephanie Ladel is one such vol")]), @@ -74,32 +71,19 @@ class TestInstagramTbotExtractor(object): ("https://www.youtube.com/watch?v=ymCMy8OffHM", False, []), ("https://www.instagram.com/p/INVALID", False, [MagicMock(id=101, media=None, message="You must enter a URL to a post")]), ]) - def test_download(self, url, expected_status, bot_responses, metadata_sample): + def test_download(self, mock_sleep, url, expected_status, bot_responses, metadata_sample): """Test the `download()` method with various Instagram URLs.""" metadata_sample.set_url(url) - self.extractor.initialise() + self.extractor.client = MagicMock() result = self.extractor.download(metadata_sample) - if expected_status: - assert result.is_success() - assert result.status == expected_status - assert result.metadata.get("title") in [msg.message[:128] for msg in bot_responses if msg.message] - else: - assert result is False - # self.extractor.cleanup() - - # @patch.object(InstagramTbotExtractor, '_send_url_to_bot') - # @patch.object(InstagramTbotExtractor, '_process_messages') - # def test_download_invalid_link_returns_false( - # self, mock_process, mock_send, extractor, metadata_instagram - # ): - # # Setup Mocks - # # _send_url_to_bot -> simulate it returns (chat=MagicMock, since_id=100) - # mock_chat = MagicMock() - # mock_send.return_value = (mock_chat, 100) - # # _process_messages -> simulate it returns the text "You must enter a URL to a post" - # mock_process.return_value = "You must enter a URL to a post" - # result = extractor.download(metadata_instagram) - # assert result is False, "Should return False if message includes 'You must enter a URL to a post'" + pass + # TODO fully mock or use as authenticated test + # if expected_status: + # assert result.is_success() + # assert result.status == expected_status + # assert result.metadata.get("title") in [msg.message[:128] for msg in bot_responses if msg.message] + # else: + # assert result is False diff --git a/tests/feeders/test_gsheet_feeder.py b/tests/feeders/test_gsheet_feeder.py index 62380f5..103610e 100644 --- a/tests/feeders/test_gsheet_feeder.py +++ b/tests/feeders/test_gsheet_feeder.py @@ -9,57 +9,52 @@ from auto_archiver.core import Metadata, Feeder def test_initialise_without_sheet_and_sheet_id(setup_module): """Ensure initialise() raises AssertionError if neither sheet nor sheet_id is set. - (shouldn't really be asserting in there) + (shouldn't really be asserting in there) """ with patch("gspread.service_account"): - feeder = setup_module("gsheet_feeder", - {"service_account": "dummy.json", - "sheet": None, - "sheet_id": None}) with pytest.raises(AssertionError): - feeder.initialise() + setup_module( + "gsheet_feeder", + {"service_account": "dummy.json", "sheet": None, "sheet_id": None}, + ) @pytest.fixture def gsheet_feeder(setup_module) -> GsheetsFeeder: - feeder = setup_module("gsheet_feeder", - {"service_account": "dummy.json", - "sheet": "test-auto-archiver", - "sheet_id": None, - "header": 1, - "columns": { - "url": "link", - "status": "archive status", - "folder": "destination folder", - "archive": "archive location", - "date": "archive date", - "thumbnail": "thumbnail", - "timestamp": "upload timestamp", - "title": "upload title", - "text": "text content", - "screenshot": "screenshot", - "hash": "hash", - "pdq_hash": "perceptual hashes", - "wacz": "wacz", - "replaywebpage": "replaywebpage", - }, - "allow_worksheets": set(), - "block_worksheets": set(), - "use_sheet_names_in_stored_paths": True, - } - ) + with patch("gspread.service_account"): + feeder = setup_module( + "gsheet_feeder", + { + "service_account": "dummy.json", + "sheet": "test-auto-archiver", + "sheet_id": None, + "header": 1, + "columns": { + "url": "link", + "status": "archive status", + "folder": "destination folder", + "archive": "archive location", + "date": "archive date", + "thumbnail": "thumbnail", + "timestamp": "upload timestamp", + "title": "upload title", + "text": "text content", + "screenshot": "screenshot", + "hash": "hash", + "pdq_hash": "perceptual hashes", + "wacz": "wacz", + "replaywebpage": "replaywebpage", + }, + "allow_worksheets": set(), + "block_worksheets": set(), + "use_sheet_names_in_stored_paths": True, + }, + ) feeder.gsheets_client = MagicMock() return feeder -@pytest.fixture() -def worksheet(unpickle): - # Load the worksheet data from the pickle file - # only works for simple usage, cant reauthenticate but give structure - return unpickle("test_worksheet.pickle") - - -class TestWorksheet(): +class TestWorksheet: """ mimics the bits we need from gworksheet """ @@ -68,12 +63,17 @@ class TestWorksheet(): title = "TestSheet" rows = [ - { "row": 2, "url": "http://example.com", "status": "", "folder": "" }, - { "row": 3, "url": "http://example.com", "status": "", "folder": "" }, - { "row": 4, "url": "", "status": "", "folder": "" }, - { "row": 5, "url": "https://another.com", "status": None, "folder": "" }, - { "row": 6, "url": "https://another.com", "status": "success", "folder": "some_folder" }, - ] + {"row": 2, "url": "http://example.com", "status": "", "folder": ""}, + {"row": 3, "url": "http://example.com", "status": "", "folder": ""}, + {"row": 4, "url": "", "status": "", "folder": ""}, + {"row": 5, "url": "https://another.com", "status": None, "folder": ""}, + { + "row": 6, + "url": "https://another.com", + "status": "success", + "folder": "some_folder", + }, + ] def __init__(self): self.wks = self.SheetSheet() @@ -91,6 +91,7 @@ class TestWorksheet(): matching = next((r for r in self.rows if r["row"] == row), {}) return matching.get(col_name, default) + def test__process_rows(gsheet_feeder: GsheetsFeeder): testworksheet = TestWorksheet() metadata_items = list(gsheet_feeder._process_rows(testworksheet)) @@ -98,9 +99,12 @@ def test__process_rows(gsheet_feeder: GsheetsFeeder): assert isinstance(metadata_items[0], Metadata) assert metadata_items[0].get("url") == "http://example.com" -def test__set_metadata(gsheet_feeder: GsheetsFeeder, worksheet): - gsheet_feeder._set_context(worksheet, 1) - assert Metadata.get_context("gsheet") == {"row": 1, "worksheet": worksheet} + +def test__set_metadata(gsheet_feeder: GsheetsFeeder): + worksheet = TestWorksheet() + metadata = Metadata() + gsheet_feeder._set_context(metadata, worksheet, 1) + assert metadata.get_context("gsheet") == {"row": 1, "worksheet": worksheet} @pytest.mark.skip(reason="Not recognising folder column") @@ -111,18 +115,24 @@ def test__set_metadata_with_folder_pickled(gsheet_feeder: GsheetsFeeder, workshe def test__set_metadata_with_folder(gsheet_feeder: GsheetsFeeder): testworksheet = TestWorksheet() + metadata = Metadata() testworksheet.wks.title = "TestSheet" - gsheet_feeder._set_context(testworksheet, 6) - assert Metadata.get_context("gsheet") == {"row": 6, "worksheet": testworksheet} - assert Metadata.get_context("folder") == "some-folder/test-auto-archiver/testsheet" + gsheet_feeder._set_context(metadata, testworksheet, 6) + assert metadata.get_context("gsheet") == {"row": 6, "worksheet": testworksheet} + assert metadata.get_context("folder") == "some-folder/test-auto-archiver/testsheet" @pytest.mark.usefixtures("setup_module") -@pytest.mark.parametrize("sheet, sheet_id, expected_method, expected_arg, description", [ - ("TestSheet", None, "open", "TestSheet", "opening by sheet name"), - (None, "ABC123", "open_by_key", "ABC123", "opening by sheet ID") -]) -def test_open_sheet_with_name_or_id(setup_module, sheet, sheet_id, expected_method, expected_arg, description): +@pytest.mark.parametrize( + "sheet, sheet_id, expected_method, expected_arg, description", + [ + ("TestSheet", None, "open", "TestSheet", "opening by sheet name"), + (None, "ABC123", "open_by_key", "ABC123", "opening by sheet ID"), + ], +) +def test_open_sheet_with_name_or_id( + setup_module, sheet, sheet_id, expected_method, expected_arg, description +): """Ensure open_sheet() correctly opens by name or ID based on configuration.""" with patch("gspread.service_account") as mock_service_account: mock_client = MagicMock() @@ -131,15 +141,16 @@ def test_open_sheet_with_name_or_id(setup_module, sheet, sheet_id, expected_meth mock_client.open_by_key.return_value = "MockSheet" # Setup module with parameterized values - feeder = setup_module("gsheet_feeder", { - "service_account": "dummy.json", - "sheet": sheet, - "sheet_id": sheet_id - }) + feeder = setup_module( + "gsheet_feeder", + {"service_account": "dummy.json", "sheet": sheet, "sheet_id": sheet_id}, + ) feeder.initialise() sheet_result = feeder.open_sheet() # Validate the correct method was called - getattr(mock_client, expected_method).assert_called_once_with(expected_arg), f"Failed: {description}" + getattr(mock_client, expected_method).assert_called_once_with( + expected_arg + ), f"Failed: {description}" assert sheet_result == "MockSheet", f"Failed: {description}" @@ -150,10 +161,10 @@ def test_open_sheet_with_sheet_id(setup_module): mock_client = MagicMock() mock_service_account.return_value = mock_client mock_client.open_by_key.return_value = "MockSheet" - feeder = setup_module("gsheet_feeder", - {"service_account": "dummy.json", - "sheet": None, - "sheet_id": "ABC123"}) + feeder = setup_module( + "gsheet_feeder", + {"service_account": "dummy.json", "sheet": None, "sheet_id": "ABC123"}, + ) feeder.initialise() sheet = feeder.open_sheet() mock_client.open_by_key.assert_called_once_with("ABC123") @@ -161,47 +172,51 @@ def test_open_sheet_with_sheet_id(setup_module): def test_should_process_sheet(setup_module): - gdb = setup_module("gsheet_feeder", {"service_account": "dummy.json", - "sheet": "TestSheet", - "sheet_id": None, - "allow_worksheets": {"TestSheet", "Sheet2"}, - "block_worksheets": {"Sheet3"}} - ) + with patch("gspread.service_account"): + gdb = setup_module( + "gsheet_feeder", + { + "service_account": "dummy.json", + "sheet": "TestSheet", + "sheet_id": None, + "allow_worksheets": {"TestSheet", "Sheet2"}, + "block_worksheets": {"Sheet3"}, + }, + ) assert gdb.should_process_sheet("TestSheet") == True assert gdb.should_process_sheet("Sheet3") == False # False if allow_worksheets is set assert gdb.should_process_sheet("AnotherSheet") == False - -@pytest.mark.skip +# @pytest.mark.skip(reason="Requires a real connection") class TestGSheetsFeederReal: + """Testing GSheetsFeeder class""" - """ Testing GSheetsFeeder class """ - module_name: str = 'gsheet_feeder' + module_name: str = "gsheet_feeder" feeder: GsheetsFeeder + # You must follow the setup process explain in the docs for this to work config: dict = { - # TODO: Create test creds "service_account": "secrets/service_account.json", "sheet": "test-auto-archiver", "sheet_id": None, "header": 1, "columns": { - "url": "link", - "status": "archive status", - "folder": "destination folder", - "archive": "archive location", - "date": "archive date", - "thumbnail": "thumbnail", - "timestamp": "upload timestamp", - "title": "upload title", - "text": "text content", - "screenshot": "screenshot", - "hash": "hash", - "pdq_hash": "perceptual hashes", - "wacz": "wacz", - "replaywebpage": "replaywebpage", - }, + "url": "link", + "status": "archive status", + "folder": "destination folder", + "archive": "archive location", + "date": "archive date", + "thumbnail": "thumbnail", + "timestamp": "upload timestamp", + "title": "upload title", + "text": "text content", + "screenshot": "screenshot", + "hash": "hash", + "pdq_hash": "perceptual hashes", + "wacz": "wacz", + "replaywebpage": "replaywebpage", + }, "allow_worksheets": set(), "block_worksheets": set(), "use_sheet_names_in_stored_paths": True, @@ -213,9 +228,7 @@ class TestGSheetsFeederReal: self.module_name is not None ), "self.module_name must be set on the subclass" assert self.config is not None, "self.config must be a dict set on the subclass" - self.feeder: Type[Feeder] = setup_module( - self.module_name, self.config - ) + self.feeder: Type[Feeder] = setup_module(self.module_name, self.config) def reset_test_sheet(self): """Clears test sheet and re-adds headers to ensure consistent test results.""" @@ -225,19 +238,17 @@ class TestGSheetsFeederReal: worksheet.clear() worksheet.append_row(["Link", "Archive Status"]) - def test_initialise(self): - self.feeder.initialise() + def test_setup(self): assert hasattr(self.feeder, "gsheets_client") - @pytest.mark.download def test_open_sheet_real_connection(self): """Ensure open_sheet() connects to a real Google Sheets instance.""" - self.feeder.initialise() sheet = self.feeder.open_sheet() assert sheet is not None, "open_sheet() should return a valid sheet instance" - assert hasattr(sheet, "worksheets"), "Returned object should have worksheets method" + assert hasattr( + sheet, "worksheets" + ), "Returned object should have worksheets method" - @pytest.mark.download def test_iter_yields_metadata_real_data(self): """Ensure __iter__() yields Metadata objects for real test sheet data.""" self.reset_test_sheet() @@ -260,7 +271,6 @@ class TestGSheetsFeederReal: assert metadata_list[0].metadata.get("url") == "https://example.com" - # TODO # Test two sheets diff --git a/tests/storages/test_S3_storage.py b/tests/storages/test_S3_storage.py index df1c1f1..60b40e6 100644 --- a/tests/storages/test_S3_storage.py +++ b/tests/storages/test_S3_storage.py @@ -1,9 +1,101 @@ from typing import Type import pytest -from unittest.mock import MagicMock, patch, mock_open +from unittest.mock import MagicMock, patch, PropertyMock from auto_archiver.core import Media +from auto_archiver.modules.hash_enricher import HashEnricher from auto_archiver.modules.s3_storage import s3_storage -from tests.storages.test_storage_base import TestStorageBase + + +@patch('boto3.client') +@pytest.fixture +def s3_store(setup_module): + config: dict = { + "path_generator": "flat", + "filename_generator": "static", + "bucket": "test-bucket", + "region": "test-region", + "key": "test-key", + "secret": "test-secret", + "random_no_duplicate": False, + "endpoint_url": "https://{region}.example.com", + "cdn_url": "https://cdn.example.com/{key}", + "private": False, + } + s3_storage = setup_module("s3_storage", config) + return s3_storage + +def test_client_initialization(s3_store): + """Test that S3 client is initialized with correct parameters""" + assert s3_store.s3 is not None + assert s3_store.s3.meta.region_name == 'test-region' + + +def test_get_cdn_url_generation(s3_store): + """Test CDN URL formatting """ + media = Media("test.txt") + media.key = "path/to/file.txt" + url = s3_store.get_cdn_url(media) + assert url == "https://cdn.example.com/path/to/file.txt" + media.key = "another/path.jpg" + assert s3_store.get_cdn_url(media) == "https://cdn.example.com/another/path.jpg" + + +@patch.object(s3_storage.S3Storage, 'file_in_folder') +def test_skips_upload_when_duplicate_exists(mock_file_in_folder, s3_store): + """Test that upload skips when file_in_folder finds existing object""" + # Setup test-specific configuration + s3_store.random_no_duplicate = True + mock_file_in_folder.return_value = "existing_folder/existing_file.txt" + # Create test media with calculated hash + media = Media("test.txt") + media.key = "original_path.txt" + + # Mock hash calculation + with patch.object(s3_store, 'calculate_hash') as mock_calculate_hash: + mock_calculate_hash.return_value = "testhash123" + # Verify upload + assert s3_store.is_upload_needed(media) is False + assert media.key == "existing_folder/existing_file.txt" + assert media.get("previously archived") is True + + with patch.object(s3_store.s3, 'upload_fileobj') as mock_upload: + result = s3_store.uploadf(None, media) + mock_upload.assert_not_called() + assert result is True + +@patch.object(s3_storage.S3Storage, 'is_upload_needed') +def test_uploads_with_correct_parameters(mock_upload_needed, s3_store): + media = Media("test.txt") + mock_upload_needed.return_value = True + media.mimetype = 'image/png' + mock_file = MagicMock() + + with patch.object(s3_store.s3, 'upload_fileobj') as mock_upload: + s3_store.uploadf(mock_file, media) + + # Verify core upload parameters + mock_upload.assert_called_once_with( + mock_file, + Bucket='test-bucket', + # Key='original_key.txt', + Key=None, + ExtraArgs={ + 'ACL': 'public-read', + 'ContentType': 'image/png' + } + ) + + + + + + + + +# ============================================================ + + + class TestGDriveStorage: @@ -29,20 +121,13 @@ class TestGDriveStorage: @patch('boto3.client') @pytest.fixture(autouse=True) def setup_storage(self, setup_module): + he = HashEnricher() self.storage = setup_module(self.module_name, self.config) - self.storage.initialise() - @patch('boto3.client') - def test_client_initialization(self, mock_boto_client, setup_module): + def test_client_initialization(self, setup_storage): """Test that S3 client is initialized with correct parameters""" - self.storage.initialise() - mock_boto_client.assert_called_once_with( - 's3', - region_name='test-region', - endpoint_url='https://test-region.example.com', - aws_access_key_id='test-key', - aws_secret_access_key='test-secret' - ) + assert self.storage.s3 is not None + assert self.storage.s3.meta.region_name == 'test-region' def test_get_cdn_url_generation(self): """Test CDN URL formatting """ @@ -53,6 +138,18 @@ class TestGDriveStorage: media.key = "another/path.jpg" assert self.storage.get_cdn_url(media) == "https://cdn.example.com/another/path.jpg" + def test_upload_decision_logic(self): + """Test is_upload_needed under different conditions""" + media = Media("test.txt") + + # Test random_no_duplicate disabled + assert self.storage.is_upload_needed(media) is True + + # Test duplicate exists + self.storage.random_no_duplicate = True + with patch.object(self.storage, 'file_in_folder', return_value='existing.txt'): + assert self.storage.is_upload_needed(media) is False + assert media.key == 'existing.txt' @patch.object(s3_storage.S3Storage, 'file_in_folder') def test_skips_upload_when_duplicate_exists(self, mock_file_in_folder): diff --git a/tests/storages/test_storage_base.py b/tests/storages/test_storage_base.py index 50d8846..7578acd 100644 --- a/tests/storages/test_storage_base.py +++ b/tests/storages/test_storage_base.py @@ -2,7 +2,6 @@ from typing import Type import pytest -from auto_archiver.core.context import ArchivingContext from auto_archiver.core.metadata import Metadata from auto_archiver.core.storage import Storage From e9ad1e1b85dbea08354189e775ae4718b4ea52cb Mon Sep 17 00:00:00 2001 From: erinhmclark Date: Thu, 6 Feb 2025 22:01:55 +0000 Subject: [PATCH 30/62] Pass media to storage cdn_call --- src/auto_archiver/core/media.py | 2 +- .../modules/gdrive_storage/gdrive_storage.py | 11 +- tests/storages/test_S3_storage.py | 149 +++++------------- 3 files changed, 49 insertions(+), 113 deletions(-) diff --git a/src/auto_archiver/core/media.py b/src/auto_archiver/core/media.py index 2cb6fc9..952a025 100644 --- a/src/auto_archiver/core/media.py +++ b/src/auto_archiver/core/media.py @@ -65,7 +65,7 @@ class Media: def is_stored(self, in_storage) -> bool: # checks if the media is already stored in the given storage - return len(self.urls) > 0 and any([u for u in self.urls if in_storage.get_cdn_url() in u]) + return len(self.urls) > 0 and any([u for u in self.urls if in_storage.get_cdn_url(self) in u]) def set(self, key: str, value: Any) -> Media: self.properties[key] = value diff --git a/src/auto_archiver/modules/gdrive_storage/gdrive_storage.py b/src/auto_archiver/modules/gdrive_storage/gdrive_storage.py index b764f1d..cc9cf3d 100644 --- a/src/auto_archiver/modules/gdrive_storage/gdrive_storage.py +++ b/src/auto_archiver/modules/gdrive_storage/gdrive_storage.py @@ -74,7 +74,8 @@ class GDriveStorage(Storage): parent_id = folder_id # get id of file inside folder (or sub folder) - file_id = self._get_id_from_parent_and_name(folder_id, filename) + # TODO: supressing the error as being checked before first upload + file_id = self._get_id_from_parent_and_name(folder_id, filename, raise_on_missing=False) return f"https://drive.google.com/file/d/{file_id}/view?usp=sharing" def upload(self, media: Media, **kwargs) -> bool: @@ -106,7 +107,13 @@ class GDriveStorage(Storage): # must be implemented even if unused def uploadf(self, file: IO[bytes], key: str, **kwargs: dict) -> bool: pass - def _get_id_from_parent_and_name(self, parent_id: str, name: str, retries: int = 1, sleep_seconds: int = 10, use_mime_type: bool = False, raise_on_missing: bool = True, use_cache=False): + def _get_id_from_parent_and_name(self, parent_id: str, + name: str, + retries: int = 1, + sleep_seconds: int = 10, + use_mime_type: bool = False, + raise_on_missing: bool = True, + use_cache=False): """ Retrieves the id of a folder or file from its @name and the @parent_id folder Optionally does multiple @retries and sleeps @sleep_seconds between them diff --git a/tests/storages/test_S3_storage.py b/tests/storages/test_S3_storage.py index 60b40e6..2594e73 100644 --- a/tests/storages/test_S3_storage.py +++ b/tests/storages/test_S3_storage.py @@ -1,103 +1,11 @@ from typing import Type import pytest -from unittest.mock import MagicMock, patch, PropertyMock +from unittest.mock import MagicMock, patch from auto_archiver.core import Media from auto_archiver.modules.hash_enricher import HashEnricher from auto_archiver.modules.s3_storage import s3_storage -@patch('boto3.client') -@pytest.fixture -def s3_store(setup_module): - config: dict = { - "path_generator": "flat", - "filename_generator": "static", - "bucket": "test-bucket", - "region": "test-region", - "key": "test-key", - "secret": "test-secret", - "random_no_duplicate": False, - "endpoint_url": "https://{region}.example.com", - "cdn_url": "https://cdn.example.com/{key}", - "private": False, - } - s3_storage = setup_module("s3_storage", config) - return s3_storage - -def test_client_initialization(s3_store): - """Test that S3 client is initialized with correct parameters""" - assert s3_store.s3 is not None - assert s3_store.s3.meta.region_name == 'test-region' - - -def test_get_cdn_url_generation(s3_store): - """Test CDN URL formatting """ - media = Media("test.txt") - media.key = "path/to/file.txt" - url = s3_store.get_cdn_url(media) - assert url == "https://cdn.example.com/path/to/file.txt" - media.key = "another/path.jpg" - assert s3_store.get_cdn_url(media) == "https://cdn.example.com/another/path.jpg" - - -@patch.object(s3_storage.S3Storage, 'file_in_folder') -def test_skips_upload_when_duplicate_exists(mock_file_in_folder, s3_store): - """Test that upload skips when file_in_folder finds existing object""" - # Setup test-specific configuration - s3_store.random_no_duplicate = True - mock_file_in_folder.return_value = "existing_folder/existing_file.txt" - # Create test media with calculated hash - media = Media("test.txt") - media.key = "original_path.txt" - - # Mock hash calculation - with patch.object(s3_store, 'calculate_hash') as mock_calculate_hash: - mock_calculate_hash.return_value = "testhash123" - # Verify upload - assert s3_store.is_upload_needed(media) is False - assert media.key == "existing_folder/existing_file.txt" - assert media.get("previously archived") is True - - with patch.object(s3_store.s3, 'upload_fileobj') as mock_upload: - result = s3_store.uploadf(None, media) - mock_upload.assert_not_called() - assert result is True - -@patch.object(s3_storage.S3Storage, 'is_upload_needed') -def test_uploads_with_correct_parameters(mock_upload_needed, s3_store): - media = Media("test.txt") - mock_upload_needed.return_value = True - media.mimetype = 'image/png' - mock_file = MagicMock() - - with patch.object(s3_store.s3, 'upload_fileobj') as mock_upload: - s3_store.uploadf(mock_file, media) - - # Verify core upload parameters - mock_upload.assert_called_once_with( - mock_file, - Bucket='test-bucket', - # Key='original_key.txt', - Key=None, - ExtraArgs={ - 'ACL': 'public-read', - 'ContentType': 'image/png' - } - ) - - - - - - - - -# ============================================================ - - - - - class TestGDriveStorage: """ Test suite for GDriveStorage. @@ -121,10 +29,9 @@ class TestGDriveStorage: @patch('boto3.client') @pytest.fixture(autouse=True) def setup_storage(self, setup_module): - he = HashEnricher() self.storage = setup_module(self.module_name, self.config) - def test_client_initialization(self, setup_storage): + def test_client_initialization(self): """Test that S3 client is initialized with correct parameters""" assert self.storage.s3 is not None assert self.storage.s3.meta.region_name == 'test-region' @@ -138,37 +45,55 @@ class TestGDriveStorage: media.key = "another/path.jpg" assert self.storage.get_cdn_url(media) == "https://cdn.example.com/another/path.jpg" + def test_uploadf_sets_acl_public(self): + media = Media("test.txt") + mock_file = MagicMock() + with patch.object(self.storage.s3, 'upload_fileobj') as mock_s3_upload, \ + patch.object(self.storage, 'is_upload_needed', return_value=True): + self.storage.uploadf(mock_file, media) + mock_s3_upload.assert_called_once_with( + mock_file, + Bucket='test-bucket', + Key=media.key, + ExtraArgs={'ACL': 'public-read', 'ContentType': 'text/plain'} + ) + def test_upload_decision_logic(self): """Test is_upload_needed under different conditions""" media = Media("test.txt") - - # Test random_no_duplicate disabled + # Test default state (random_no_duplicate=False) assert self.storage.is_upload_needed(media) is True + # Set duplicate checking config to true: - # Test duplicate exists self.storage.random_no_duplicate = True - with patch.object(self.storage, 'file_in_folder', return_value='existing.txt'): + with patch('auto_archiver.modules.hash_enricher.HashEnricher.calculate_hash') as mock_calc_hash, \ + patch.object(self.storage, 'file_in_folder') as mock_file_in_folder: + mock_calc_hash.return_value = 'beepboop123beepboop123beepboop123' + mock_file_in_folder.return_value = 'existing_key.txt' + # Test duplicate result assert self.storage.is_upload_needed(media) is False - assert media.key == 'existing.txt' + assert media.key == 'existing_key.txt' + mock_file_in_folder.assert_called_with( + # (first 24 chars of hash) + 'no-dups/beepboop123beepboop123be' + ) + @patch.object(s3_storage.S3Storage, 'file_in_folder') def test_skips_upload_when_duplicate_exists(self, mock_file_in_folder): """Test that upload skips when file_in_folder finds existing object""" - # Setup test-specific configuration self.storage.random_no_duplicate = True mock_file_in_folder.return_value = "existing_folder/existing_file.txt" # Create test media with calculated hash media = Media("test.txt") media.key = "original_path.txt" - # Mock hash calculation - with patch.object(self.storage, 'calculate_hash') as mock_calculate_hash: - mock_calculate_hash.return_value = "testhash123" + with patch('auto_archiver.modules.hash_enricher.HashEnricher.calculate_hash') as mock_calculate_hash: + mock_calculate_hash.return_value = "beepboop123beepboop123beepboop123" # Verify upload assert self.storage.is_upload_needed(media) is False assert media.key == "existing_folder/existing_file.txt" assert media.get("previously archived") is True - with patch.object(self.storage.s3, 'upload_fileobj') as mock_upload: result = self.storage.uploadf(None, media) mock_upload.assert_not_called() @@ -177,21 +102,25 @@ class TestGDriveStorage: @patch.object(s3_storage.S3Storage, 'is_upload_needed') def test_uploads_with_correct_parameters(self, mock_upload_needed): media = Media("test.txt") + media.key = "original_key.txt" mock_upload_needed.return_value = True media.mimetype = 'image/png' mock_file = MagicMock() with patch.object(self.storage.s3, 'upload_fileobj') as mock_upload: self.storage.uploadf(mock_file, media) - - # Verify core upload parameters + # verify call occured with these params mock_upload.assert_called_once_with( mock_file, Bucket='test-bucket', - # Key='original_key.txt', - Key=None, + Key='original_key.txt', ExtraArgs={ 'ACL': 'public-read', 'ContentType': 'image/png' } - ) \ No newline at end of file + ) + + def test_file_in_folder_exists(self): + with patch.object(self.storage.s3, 'list_objects') as mock_list_objects: + mock_list_objects.return_value = {'Contents': [{'Key': 'path/to/file.txt'}]} + assert self.storage.file_in_folder('path/to/') == 'path/to/file.txt' \ No newline at end of file From 2920cf685f8c556cbdfa8d805f1eb20b8fe41d66 Mon Sep 17 00:00:00 2001 From: erinhmclark Date: Fri, 7 Feb 2025 12:35:40 +0000 Subject: [PATCH 31/62] Small fixes to whisper_enricher.py. --- src/auto_archiver/modules/whisper_enricher/__manifest__.py | 6 ++++-- .../modules/whisper_enricher/whisper_enricher.py | 2 +- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/src/auto_archiver/modules/whisper_enricher/__manifest__.py b/src/auto_archiver/modules/whisper_enricher/__manifest__.py index f7ad1b3..884de66 100644 --- a/src/auto_archiver/modules/whisper_enricher/__manifest__.py +++ b/src/auto_archiver/modules/whisper_enricher/__manifest__.py @@ -6,8 +6,10 @@ "python": ["s3_storage", "loguru", "requests"], }, "configs": { - "api_endpoint": {"default": None, "help": "WhisperApi api endpoint, eg: https://whisperbox-api.com/api/v1, a deployment of https://github.com/bellingcat/whisperbox-transcribe."}, - "api_key": {"default": None, "help": "WhisperApi api key for authentication"}, + "api_endpoint": {"required": True, + "help": "WhisperApi api endpoint, eg: https://whisperbox-api.com/api/v1, a deployment of https://github.com/bellingcat/whisperbox-transcribe."}, + "api_key": {"required": True, + "help": "WhisperApi api key for authentication"}, "include_srt": {"default": False, "help": "Whether to include a subtitle SRT (SubRip Subtitle file) for the video (can be used in video players)."}, "timeout": {"default": 90, "help": "How many seconds to wait at most for a successful job completion."}, "action": {"default": "translate", "help": "which Whisper operation to execute", "choices": ["transcribe", "translate", "language_detection"]}, diff --git a/src/auto_archiver/modules/whisper_enricher/whisper_enricher.py b/src/auto_archiver/modules/whisper_enricher/whisper_enricher.py index 8ca2131..a7298e4 100644 --- a/src/auto_archiver/modules/whisper_enricher/whisper_enricher.py +++ b/src/auto_archiver/modules/whisper_enricher/whisper_enricher.py @@ -110,7 +110,7 @@ class WhisperEnricher(Enricher): def _get_s3_storage(self) -> S3Storage: try: - return next(s for s in self.storages if s.__class__ == S3Storage) + return next(s for s in self.config['steps']['storages'] if s == 's3_storage') except: logger.warning("No S3Storage instance found in storages") return From 950624dd4bb0e917abbe58c98351bbabd26d0bb3 Mon Sep 17 00:00:00 2001 From: erinhmclark Date: Fri, 7 Feb 2025 20:26:00 +0000 Subject: [PATCH 32/62] Fix S3 storage to media in whisper_enricher.py. --- .../modules/whisper_enricher/__manifest__.py | 7 +++++-- .../whisper_enricher/whisper_enricher.py | 19 ++++++++----------- 2 files changed, 13 insertions(+), 13 deletions(-) diff --git a/src/auto_archiver/modules/whisper_enricher/__manifest__.py b/src/auto_archiver/modules/whisper_enricher/__manifest__.py index 884de66..1539df6 100644 --- a/src/auto_archiver/modules/whisper_enricher/__manifest__.py +++ b/src/auto_archiver/modules/whisper_enricher/__manifest__.py @@ -1,4 +1,4 @@ -{ +a={ "name": "Whisper Enricher", "type": ["enricher"], "requires_setup": True, @@ -12,7 +12,9 @@ "help": "WhisperApi api key for authentication"}, "include_srt": {"default": False, "help": "Whether to include a subtitle SRT (SubRip Subtitle file) for the video (can be used in video players)."}, "timeout": {"default": 90, "help": "How many seconds to wait at most for a successful job completion."}, - "action": {"default": "translate", "help": "which Whisper operation to execute", "choices": ["transcribe", "translate", "language_detection"]}, + "action": {"default": "translate", + "help": "which Whisper operation to execute", + "choices": ["transcribe", "translate", "language_detection"]}, }, "description": """ Integrates with a Whisper API service to transcribe, translate, or detect the language of audio and video files. @@ -27,6 +29,7 @@ ### Notes - Requires a Whisper API endpoint and API key for authentication. - Only compatible with S3-compatible storage systems for media file accessibility. + - ** This stores the media files in S3 prior to enriching them as Whisper requires public URLs to access the media files. - Handles multiple jobs and retries for failed or incomplete processing. """ } diff --git a/src/auto_archiver/modules/whisper_enricher/whisper_enricher.py b/src/auto_archiver/modules/whisper_enricher/whisper_enricher.py index a7298e4..004d91c 100644 --- a/src/auto_archiver/modules/whisper_enricher/whisper_enricher.py +++ b/src/auto_archiver/modules/whisper_enricher/whisper_enricher.py @@ -15,17 +15,21 @@ class WhisperEnricher(Enricher): """ def enrich(self, to_enrich: Metadata) -> None: - if not self._get_s3_storage(): + storages = self.config['steps']['storages'] + if not "s3_storage" in storages: logger.error("WhisperEnricher: To use the WhisperEnricher you need to use S3Storage so files are accessible publicly to the whisper service being called.") return + self.s3 = get_module("s3_storage", self.config) url = to_enrich.get_url() logger.debug(f"WHISPER[{self.action}]: iterating media items for {url=}.") job_results = {} for i, m in enumerate(to_enrich.media): if m.is_video() or m.is_audio(): - m.store(url=url, metadata=to_enrich, storages=self.storages) + # TODO: this used to pass all storage items to store now + # Now only passing S3, the rest will get added later in the usual order (?) + m.store(url=url, metadata=to_enrich, storages=[self.s3]) try: job_id = self.submit_job(m) job_results[job_id] = False @@ -53,8 +57,8 @@ class WhisperEnricher(Enricher): to_enrich.set_content(f"\n[automatic video transcript]: {v}") def submit_job(self, media: Media): - s3 = get_module("s3_storage", self.config) - s3_url = s3.get_cdn_url(media) + + s3_url = self.s3.get_cdn_url(media) assert s3_url in media.urls, f"Could not find S3 url ({s3_url}) in list of stored media urls " payload = { "url": s3_url, @@ -107,10 +111,3 @@ class WhisperEnricher(Enricher): logger.debug(f"DELETE whisper {job_id=} result: {r_del.status_code}") return result return False - - def _get_s3_storage(self) -> S3Storage: - try: - return next(s for s in self.config['steps']['storages'] if s == 's3_storage') - except: - logger.warning("No S3Storage instance found in storages") - return From 63aba6ad3994a27b7e95116dd9d6b8c4fd40e452 Mon Sep 17 00:00:00 2001 From: Patrick Robertson Date: Fri, 7 Feb 2025 21:54:49 +0100 Subject: [PATCH 33/62] Fix sphinx-autoapi imports --- src/auto_archiver/core/extractor.py | 2 +- src/auto_archiver/core/orchestrator.py | 2 +- .../modules/generic_extractor/generic_extractor.py | 2 +- src/auto_archiver/utils/gsheet.py | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/auto_archiver/core/extractor.py b/src/auto_archiver/core/extractor.py index 57320df..794c06c 100644 --- a/src/auto_archiver/core/extractor.py +++ b/src/auto_archiver/core/extractor.py @@ -17,7 +17,7 @@ from loguru import logger from retrying import retry import re -from ..core import Metadata, BaseModule +from auto_archiver.core import Metadata, BaseModule class Extractor(BaseModule): diff --git a/src/auto_archiver/core/orchestrator.py b/src/auto_archiver/core/orchestrator.py index 5ac091c..641f099 100644 --- a/src/auto_archiver/core/orchestrator.py +++ b/src/auto_archiver/core/orchestrator.py @@ -19,7 +19,7 @@ from rich_argparse import RichHelpFormatter from .metadata import Metadata, Media -from ..version import __version__ +from auto_archiver.version import __version__ from .config import yaml, read_yaml, store_yaml, to_dot_notation, merge_dicts, EMPTY_CONFIG, DefaultValidatingParser from .module import available_modules, LazyBaseModule, get_module, setup_paths from . import validators, Feeder, Extractor, Database, Storage, Formatter, Enricher diff --git a/src/auto_archiver/modules/generic_extractor/generic_extractor.py b/src/auto_archiver/modules/generic_extractor/generic_extractor.py index d1b1fb6..86e978f 100644 --- a/src/auto_archiver/modules/generic_extractor/generic_extractor.py +++ b/src/auto_archiver/modules/generic_extractor/generic_extractor.py @@ -6,7 +6,7 @@ from yt_dlp.extractor.common import InfoExtractor from loguru import logger from auto_archiver.core.extractor import Extractor -from ...core import Metadata, Media +from auto_archiver.core import Metadata, Media class GenericExtractor(Extractor): _dropins = {} diff --git a/src/auto_archiver/utils/gsheet.py b/src/auto_archiver/utils/gsheet.py index 7a8862f..c36a032 100644 --- a/src/auto_archiver/utils/gsheet.py +++ b/src/auto_archiver/utils/gsheet.py @@ -1,6 +1,6 @@ import json, gspread -from ..core import BaseModule +from auto_archiver.core import BaseModule class Gsheets(BaseModule): From 1fad37fd934ba26835c9eb20222d95210a1e513a Mon Sep 17 00:00:00 2001 From: Patrick Robertson Date: Fri, 7 Feb 2025 23:08:30 +0100 Subject: [PATCH 34/62] Remove blank file --- src/auto_archiver/core/authentication.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) delete mode 100644 src/auto_archiver/core/authentication.py diff --git a/src/auto_archiver/core/authentication.py b/src/auto_archiver/core/authentication.py deleted file mode 100644 index e69de29..0000000 From e9dd321dcd548cc02d7fa2a0d0171feed1226c51 Mon Sep 17 00:00:00 2001 From: Patrick Robertson Date: Mon, 10 Feb 2025 13:06:24 +0100 Subject: [PATCH 35/62] Fix setting cli_feeder as default feeder on clean install --- src/auto_archiver/core/config.py | 3 ++- src/auto_archiver/core/orchestrator.py | 20 +++++++++++++------- 2 files changed, 15 insertions(+), 8 deletions(-) diff --git a/src/auto_archiver/core/config.py b/src/auto_archiver/core/config.py index 2d462e4..8f36c54 100644 --- a/src/auto_archiver/core/config.py +++ b/src/auto_archiver/core/config.py @@ -36,6 +36,7 @@ steps:""" + "".join([f"\n {module}s: []" for module in BaseModule.MODULE_TYPES # a dictionary of authentication information that can be used by extractors to login to website. # you can use a comma separated list for multiple domains on the same line (common usecase: x.com,twitter.com) # Common login 'types' are username/password, cookie, api key/token. +# There are two special keys for using cookies, they are: cookies_file and cookies_from_browser. # Some Examples: # facebook.com: # username: "my_username" @@ -163,6 +164,6 @@ def read_yaml(yaml_filename: str) -> CommentedMap: def store_yaml(config: CommentedMap, yaml_filename: str) -> None: config_to_save = deepcopy(config) - config.pop('urls', None) + config_to_save.pop('urls', None) with open(yaml_filename, "w", encoding="utf-8") as outf: yaml.dump(config_to_save, outf) \ No newline at end of file diff --git a/src/auto_archiver/core/orchestrator.py b/src/auto_archiver/core/orchestrator.py index 641f099..20212ce 100644 --- a/src/auto_archiver/core/orchestrator.py +++ b/src/auto_archiver/core/orchestrator.py @@ -128,6 +128,10 @@ class ArchivingOrchestrator: elif basic_config.mode == 'simple': simple_modules = [module for module in available_modules(with_manifest=True) if not module.requires_setup] self.add_module_args(simple_modules, parser) + + # for simple mode, we use the cli_feeder and any modules that don't require setup + yaml_config['steps']['feeders'] = ['cli_feeder'] + # add them to the config for module in simple_modules: for module_type in module.type: @@ -237,18 +241,18 @@ class ArchivingOrchestrator: if log_file := logging_config['file']: logger.add(log_file) if not logging_config['rotation'] else logger.add(log_file, rotation=logging_config['rotation']) - - def install_modules(self): + def install_modules(self, modules_by_type): """ - Swaps out the previous 'strings' in the config with the actual modules and loads them + Traverses all modules in 'steps' and loads them into the orchestrator, storing them in the + orchestrator's attributes (self.feeders, self.extractors etc.). If no modules of a certain type + are loaded, the program will exit with an error message. """ invalid_modules = [] for module_type in BaseModule.MODULE_TYPES: step_items = [] - modules_to_load = self.config['steps'][f"{module_type}s"] - + modules_to_load = modules_by_type[f"{module_type}s"] assert modules_to_load, f"No {module_type}s were configured. Make sure to set at least one {module_type} in your configuration file or on the command line (using --{module_type}s)" def check_steps_ok(): @@ -264,9 +268,11 @@ class ArchivingOrchestrator: for module in modules_to_load: if module == 'cli_feeder': + # pseudo module, don't load it + breakpoint() urls = self.config['urls'] if not urls: - logger.error("No URLs provided. Please provide at least one URL to archive, or set up a feeder. Use --help for more information.") + logger.error("No URLs provided. Please provide at least one URL via the command line, or set up an alternative feeder. Use --help for more information.") exit() # cli_feeder is a pseudo module, it just takes the command line args def feed(self) -> Generator[Metadata]: @@ -330,7 +336,7 @@ class ArchivingOrchestrator: self.setup_complete_parser(basic_config, yaml_config, unused_args) logger.info(f"======== Welcome to the AUTO ARCHIVER ({__version__}) ==========") - self.install_modules() + self.install_modules(self.config['steps']) # log out the modules that were loaded for module_type in BaseModule.MODULE_TYPES: From 74207d7821e0306de8e3b6da00cf263edfe0293c Mon Sep 17 00:00:00 2001 From: Patrick Robertson Date: Mon, 10 Feb 2025 13:27:11 +0100 Subject: [PATCH 36/62] Implementation tests for auto-archiver --- src/auto_archiver/core/orchestrator.py | 2 -- tests/test_implementation.py | 35 ++++++++++++++++++++++++++ 2 files changed, 35 insertions(+), 2 deletions(-) create mode 100644 tests/test_implementation.py diff --git a/src/auto_archiver/core/orchestrator.py b/src/auto_archiver/core/orchestrator.py index 20212ce..a451443 100644 --- a/src/auto_archiver/core/orchestrator.py +++ b/src/auto_archiver/core/orchestrator.py @@ -111,7 +111,6 @@ class ArchivingOrchestrator: # if full, we'll load all modules # TODO: BUG** - basic_config won't have steps in it, since these args aren't added to 'basic_parser' # but should we add them? Or should we just add them to the 'complete' parser? - if yaml_config != EMPTY_CONFIG: # only load the modules enabled in config # TODO: if some steps are empty (e.g. 'feeders' is empty), should we default to the 'simple' ones? Or only if they are ALL empty? @@ -269,7 +268,6 @@ class ArchivingOrchestrator: for module in modules_to_load: if module == 'cli_feeder': # pseudo module, don't load it - breakpoint() urls = self.config['urls'] if not urls: logger.error("No URLs provided. Please provide at least one URL via the command line, or set up an alternative feeder. Use --help for more information.") diff --git a/tests/test_implementation.py b/tests/test_implementation.py new file mode 100644 index 0000000..82d5d0f --- /dev/null +++ b/tests/test_implementation.py @@ -0,0 +1,35 @@ +import sys +import pytest + +from auto_archiver.__main__ import main + +@pytest.fixture +def orchestration_file(tmp_path): + return (tmp_path / "example_orch.yaml").as_posix() + +@pytest.fixture +def autoarchiver(tmp_path, monkeypatch): + + def _autoarchiver(args=["--config", "example_orch.yaml"]): + # change dir to tmp_path + monkeypatch.chdir(tmp_path) + with monkeypatch.context() as m: + m.setattr(sys, "argv", ["auto-archiver"] + args) + return main() + + return _autoarchiver + + +def test_run_auto_archiver_no_args(caplog, autoarchiver): + with pytest.raises(SystemExit): + autoarchiver([]) + + assert "provide at least one URL via the command line, or set up an alternative feeder" in caplog.text + + +def test_run_auto_archiver_invalid_file(caplog, autoarchiver, monkeypatch): + # exec 'auto-archiver' on the command lin + with pytest.raises(SystemExit): + autoarchiver() + + assert "Make sure the file exists and try again, or run without th" in caplog.text \ No newline at end of file From f3f6b928172fe597c772e2c677a3f3d118f02bef Mon Sep 17 00:00:00 2001 From: Patrick Robertson Date: Mon, 10 Feb 2025 12:43:21 +0000 Subject: [PATCH 37/62] Implementation test cleanup --- tests/test_implementation.py | 45 ++++++++++++++++++++++++++++-------- 1 file changed, 36 insertions(+), 9 deletions(-) diff --git a/tests/test_implementation.py b/tests/test_implementation.py index 82d5d0f..7e33651 100644 --- a/tests/test_implementation.py +++ b/tests/test_implementation.py @@ -3,33 +3,60 @@ import pytest from auto_archiver.__main__ import main + @pytest.fixture -def orchestration_file(tmp_path): +def orchestration_file_path(tmp_path): return (tmp_path / "example_orch.yaml").as_posix() @pytest.fixture -def autoarchiver(tmp_path, monkeypatch): +def orchestration_file(orchestration_file_path): + def _orchestration_file(content=''): + with open(orchestration_file_path, "w") as f: + f.write(content) + return orchestration_file_path + + return _orchestration_file + +@pytest.fixture +def autoarchiver(tmp_path, monkeypatch, request): + def _autoarchiver(args=[]): + + def cleanup(): + from loguru import logger + if not logger._core.handlers.get(0): + logger._core.handlers_count = 0 + logger.add(sys.stderr) + + request.addfinalizer(cleanup) - def _autoarchiver(args=["--config", "example_orch.yaml"]): # change dir to tmp_path monkeypatch.chdir(tmp_path) with monkeypatch.context() as m: m.setattr(sys, "argv", ["auto-archiver"] + args) return main() - + return _autoarchiver def test_run_auto_archiver_no_args(caplog, autoarchiver): with pytest.raises(SystemExit): - autoarchiver([]) + autoarchiver() assert "provide at least one URL via the command line, or set up an alternative feeder" in caplog.text - -def test_run_auto_archiver_invalid_file(caplog, autoarchiver, monkeypatch): +def test_run_auto_archiver_invalid_file(caplog, autoarchiver): # exec 'auto-archiver' on the command lin with pytest.raises(SystemExit): - autoarchiver() + autoarchiver(["--config", "nonexistent_file.yaml"]) - assert "Make sure the file exists and try again, or run without th" in caplog.text \ No newline at end of file + assert "Make sure the file exists and try again, or run without th" in caplog.text + +def test_run_auto_archiver_empty_file(caplog, autoarchiver, orchestration_file): + # create a valid (empty) orchestration file + path = orchestration_file(content="") + # exec 'auto-archiver' on the command lin + with pytest.raises(SystemExit): + autoarchiver(["--config", path]) + + # should treat an empty file as if there is no file at all + assert " No URLs provided. Please provide at least one URL via the com" in caplog.text From 7c848046e88a12d6b9ea89c7b6b34ab76ef009e8 Mon Sep 17 00:00:00 2001 From: msramalho <19508417+msramalho@users.noreply.github.com> Date: Mon, 10 Feb 2025 14:59:32 +0000 Subject: [PATCH 38/62] adds better info about wrong/missing modules --- src/auto_archiver/core/module.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/src/auto_archiver/core/module.py b/src/auto_archiver/core/module.py index dec67e1..f3fbec5 100644 --- a/src/auto_archiver/core/module.py +++ b/src/auto_archiver/core/module.py @@ -13,7 +13,7 @@ import copy import sys from importlib.util import find_spec import os -from os.path import join, dirname +from os.path import join from loguru import logger import auto_archiver from .base_module import BaseModule @@ -64,8 +64,10 @@ def get_module_lazy(module_name: str, suppress_warnings: bool = False) -> LazyBa if module_name in _LAZY_LOADED_MODULES: return _LAZY_LOADED_MODULES[module_name] - module = available_modules(limit_to_modules=[module_name], suppress_warnings=suppress_warnings)[0] - return module + available = available_modules(limit_to_modules=[module_name], suppress_warnings=suppress_warnings) + if not available: + raise IndexError(f"Module '{module_name}' not found. Are you sure it's installed/exists?") + return available[0] def available_modules(with_manifest: bool=False, limit_to_modules: List[str]= [], suppress_warnings: bool = False) -> List[LazyBaseModule]: From 8fb3dc754b14b76833a12daa091ef608edf6a61c Mon Sep 17 00:00:00 2001 From: msramalho <19508417+msramalho@users.noreply.github.com> Date: Mon, 10 Feb 2025 14:59:51 +0000 Subject: [PATCH 39/62] fixing telethon extractor to use default entrypoint --- src/auto_archiver/modules/telethon_extractor/__init__.py | 2 +- .../modules/telethon_extractor/telethon_extractor.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/auto_archiver/modules/telethon_extractor/__init__.py b/src/auto_archiver/modules/telethon_extractor/__init__.py index a837fdf..2eaa57c 100644 --- a/src/auto_archiver/modules/telethon_extractor/__init__.py +++ b/src/auto_archiver/modules/telethon_extractor/__init__.py @@ -1 +1 @@ -from .telethon_extractor import TelethonArchiver \ No newline at end of file +from .telethon_extractor import TelethonExtractor \ No newline at end of file diff --git a/src/auto_archiver/modules/telethon_extractor/telethon_extractor.py b/src/auto_archiver/modules/telethon_extractor/telethon_extractor.py index 3e952e8..21fc4dc 100644 --- a/src/auto_archiver/modules/telethon_extractor/telethon_extractor.py +++ b/src/auto_archiver/modules/telethon_extractor/telethon_extractor.py @@ -6,14 +6,14 @@ from telethon.tl.functions.messages import ImportChatInviteRequest from telethon.errors.rpcerrorlist import UserAlreadyParticipantError, FloodWaitError, InviteRequestSentError, InviteHashExpiredError from loguru import logger from tqdm import tqdm -import re, time, json, os +import re, time, os from auto_archiver.core import Extractor from auto_archiver.core import Metadata, Media from auto_archiver.utils import random_str -class TelethonArchiver(Extractor): +class TelethonExtractor(Extractor): valid_url = re.compile(r"https:\/\/t\.me(\/c){0,1}\/(.+)\/(\d+)") invite_pattern = re.compile(r"t.me(\/joinchat){0,1}\/\+?(.+)") From 15abf686b1315b3a35a628df12f687b9aec431d5 Mon Sep 17 00:00:00 2001 From: msramalho <19508417+msramalho@users.noreply.github.com> Date: Mon, 10 Feb 2025 15:48:54 +0000 Subject: [PATCH 40/62] decouples s3_storage from hash_enricher --- src/auto_archiver/core/base_module.py | 2 +- .../modules/hash_enricher/hash_enricher.py | 8 ++------ src/auto_archiver/modules/s3_storage/s3_storage.py | 8 +++----- src/auto_archiver/utils/misc.py | 12 ++++++++++++ 4 files changed, 18 insertions(+), 12 deletions(-) diff --git a/src/auto_archiver/core/base_module.py b/src/auto_archiver/core/base_module.py index fcfe9ea..5c6ecbb 100644 --- a/src/auto_archiver/core/base_module.py +++ b/src/auto_archiver/core/base_module.py @@ -63,7 +63,7 @@ class BaseModule(ABC): def setup(self, config: dict): authentication = config.get('authentication', {}) - # extract out contatenated sites + # extract out concatenated sites for key, val in copy(authentication).items(): if "," in key: for site in key.split(","): diff --git a/src/auto_archiver/modules/hash_enricher/hash_enricher.py b/src/auto_archiver/modules/hash_enricher/hash_enricher.py index 58c6abe..b3ca8be 100644 --- a/src/auto_archiver/modules/hash_enricher/hash_enricher.py +++ b/src/auto_archiver/modules/hash_enricher/hash_enricher.py @@ -12,6 +12,7 @@ from loguru import logger from auto_archiver.core import Enricher from auto_archiver.core import Metadata +from auto_archiver.utils.misc import calculate_file_hash class HashEnricher(Enricher): @@ -35,9 +36,4 @@ class HashEnricher(Enricher): elif self.algorithm == "SHA3-512": hash = hashlib.sha3_512() else: return "" - with open(filename, "rb") as f: - while True: - buf = f.read(self.chunksize) - if not buf: break - hash.update(buf) - return hash.hexdigest() + return calculate_file_hash(filename, hash, self.chunksize) diff --git a/src/auto_archiver/modules/s3_storage/s3_storage.py b/src/auto_archiver/modules/s3_storage/s3_storage.py index f324d5c..2f85164 100644 --- a/src/auto_archiver/modules/s3_storage/s3_storage.py +++ b/src/auto_archiver/modules/s3_storage/s3_storage.py @@ -7,12 +7,11 @@ from loguru import logger from auto_archiver.core import Media from auto_archiver.core import Storage -from auto_archiver.modules.hash_enricher import HashEnricher -from auto_archiver.utils.misc import random_str +from auto_archiver.utils.misc import calculate_file_hash, random_str NO_DUPLICATES_FOLDER = "no-dups/" -class S3Storage(Storage, HashEnricher): +class S3Storage(Storage): def setup(self, config: dict) -> None: super().setup(config) @@ -42,14 +41,13 @@ class S3Storage(Storage, HashEnricher): extra_args['ContentType'] = media.mimetype except Exception as e: logger.warning(f"Unable to get mimetype for {media.key=}, error: {e}") - self.s3.upload_fileobj(file, Bucket=self.bucket, Key=media.key, ExtraArgs=extra_args) return True def is_upload_needed(self, media: Media) -> bool: if self.random_no_duplicate: # checks if a folder with the hash already exists, if so it skips the upload - hd = self.calculate_hash(media.filename) + hd = calculate_file_hash(media.filename) path = os.path.join(NO_DUPLICATES_FOLDER, hd[:24]) if existing_key:=self.file_in_folder(path): diff --git a/src/auto_archiver/utils/misc.py b/src/auto_archiver/utils/misc.py index 300a710..3af5a54 100644 --- a/src/auto_archiver/utils/misc.py +++ b/src/auto_archiver/utils/misc.py @@ -5,6 +5,7 @@ import json import uuid from datetime import datetime import requests +import hashlib from loguru import logger @@ -54,9 +55,20 @@ def update_nested_dict(dictionary, update_dict): else: dictionary[key] = value + def random_str(length: int = 32) -> str: assert length <= 32, "length must be less than 32 as UUID4 is used" return str(uuid.uuid4()).replace("-", "")[:length] + def json_loader(cli_val): return json.loads(cli_val) + + +def calculate_file_hash(filename: str, hash_algo = hashlib.sha256(), chunksize: int = 16000000) -> str: + with open(filename, "rb") as f: + while True: + buf = f.read(chunksize) + if not buf: break + hash_algo.update(buf) + return hash_algo.hexdigest() From f311621e58446983fb95d9e510249855a7687f61 Mon Sep 17 00:00:00 2001 From: erinhmclark Date: Mon, 10 Feb 2025 15:57:42 +0000 Subject: [PATCH 41/62] Small fixes. Add timestamp helper method. --- .../modules/gdrive_storage/gdrive_storage.py | 7 +- .../modules/gsheet_db/gsheet_db.py | 70 ++++++++++--------- .../telethon_extractor/telethon_extractor.py | 4 +- .../modules/whisper_enricher/__manifest__.py | 2 +- .../whisper_enricher/whisper_enricher.py | 13 ++-- src/auto_archiver/utils/misc.py | 36 +++++++++- tests/databases/test_gsheet_db.py | 8 ++- .../test_instagram_api_extractor.py | 3 +- .../test_instagram_tbot_extractor.py | 1 - tests/feeders/test_gsheet_feeder.py | 9 +-- tests/storages/test_gdrive_storage.py | 41 ++++++++--- tests/test_metadata.py | 4 ++ 12 files changed, 129 insertions(+), 69 deletions(-) diff --git a/src/auto_archiver/modules/gdrive_storage/gdrive_storage.py b/src/auto_archiver/modules/gdrive_storage/gdrive_storage.py index cc9cf3d..910f48b 100644 --- a/src/auto_archiver/modules/gdrive_storage/gdrive_storage.py +++ b/src/auto_archiver/modules/gdrive_storage/gdrive_storage.py @@ -70,12 +70,15 @@ class GDriveStorage(Storage): filename = path_parts[-1] logger.info(f"looking for folders for {path_parts[0:-1]} before getting url for {filename=}") for folder in path_parts[0:-1]: - folder_id = self._get_id_from_parent_and_name(parent_id, folder, use_mime_type=True, raise_on_missing=True) + folder_id = self._get_id_from_parent_and_name(parent_id, folder, use_mime_type=True, raise_on_missing=False) parent_id = folder_id - # get id of file inside folder (or sub folder) # TODO: supressing the error as being checked before first upload file_id = self._get_id_from_parent_and_name(folder_id, filename, raise_on_missing=False) + if not file_id: + # + logger.info(f"file {filename} not found in folder {folder_id}") + return None return f"https://drive.google.com/file/d/{file_id}/view?usp=sharing" def upload(self, media: Media, **kwargs) -> bool: diff --git a/src/auto_archiver/modules/gsheet_db/gsheet_db.py b/src/auto_archiver/modules/gsheet_db/gsheet_db.py index 3bb27b7..682eb94 100644 --- a/src/auto_archiver/modules/gsheet_db/gsheet_db.py +++ b/src/auto_archiver/modules/gsheet_db/gsheet_db.py @@ -1,6 +1,4 @@ from typing import Union, Tuple - -import datetime from urllib.parse import quote from loguru import logger @@ -8,33 +6,33 @@ from loguru import logger from auto_archiver.core import Database from auto_archiver.core import Metadata, Media from auto_archiver.modules.gsheet_feeder import GWorksheet +from auto_archiver.utils.misc import get_current_timestamp class GsheetsDb(Database): """ - NB: only works if GsheetFeeder is used. - could be updated in the future to support non-GsheetFeeder metadata + NB: only works if GsheetFeeder is used. + could be updated in the future to support non-GsheetFeeder metadata """ - def started(self, item: Metadata) -> None: logger.warning(f"STARTED {item}") gw, row = self._retrieve_gsheet(item) - gw.set_cell(row, 'status', 'Archive in progress') + gw.set_cell(row, "status", "Archive in progress") - def failed(self, item: Metadata, reason:str) -> None: + def failed(self, item: Metadata, reason: str) -> None: logger.error(f"FAILED {item}") - self._safe_status_update(item, f'Archive failed {reason}') + self._safe_status_update(item, f"Archive failed {reason}") def aborted(self, item: Metadata) -> None: logger.warning(f"ABORTED {item}") - self._safe_status_update(item, '') + self._safe_status_update(item, "") def fetch(self, item: Metadata) -> Union[Metadata, bool]: """check if the given item has been archived already""" return False - def done(self, item: Metadata, cached: bool=False) -> None: + def done(self, item: Metadata, cached: bool = False) -> None: """archival result ready - should be saved to DB""" logger.success(f"DONE {item.get_url()}") gw, row = self._retrieve_gsheet(item) @@ -46,23 +44,25 @@ class GsheetsDb(Database): def batch_if_valid(col, val, final_value=None): final_value = final_value or val try: - if val and gw.col_exists(col) and gw.get_cell(row_values, col) == '': + if val and gw.col_exists(col) and gw.get_cell(row_values, col) == "": cell_updates.append((row, col, final_value)) except Exception as e: logger.error(f"Unable to batch {col}={final_value} due to {e}") + status_message = item.status if cached: status_message = f"[cached] {status_message}" - cell_updates.append((row, 'status', status_message)) + cell_updates.append((row, "status", status_message)) media: Media = item.get_final_media() if hasattr(media, "urls"): - batch_if_valid('archive', "\n".join(media.urls)) - batch_if_valid('date', True, self._get_current_datetime_iso()) - batch_if_valid('title', item.get_title()) - batch_if_valid('text', item.get("content", "")) - batch_if_valid('timestamp', item.get_timestamp()) - if media: batch_if_valid('hash', media.get("hash", "not-calculated")) + batch_if_valid("archive", "\n".join(media.urls)) + batch_if_valid("date", True, get_current_timestamp()) + batch_if_valid("title", item.get_title()) + batch_if_valid("text", item.get("content", "")) + batch_if_valid("timestamp", item.get_timestamp()) + if media: + batch_if_valid("hash", media.get("hash", "not-calculated")) # merge all pdq hashes into a single string, if present pdq_hashes = [] @@ -71,31 +71,35 @@ class GsheetsDb(Database): if pdq := m.get("pdq_hash"): pdq_hashes.append(pdq) if len(pdq_hashes): - batch_if_valid('pdq_hash', ",".join(pdq_hashes)) + batch_if_valid("pdq_hash", ",".join(pdq_hashes)) - if (screenshot := item.get_media_by_id("screenshot")) and hasattr(screenshot, "urls"): - batch_if_valid('screenshot', "\n".join(screenshot.urls)) + if (screenshot := item.get_media_by_id("screenshot")) and hasattr( + screenshot, "urls" + ): + batch_if_valid("screenshot", "\n".join(screenshot.urls)) - if (thumbnail := item.get_first_image("thumbnail")): + if thumbnail := item.get_first_image("thumbnail"): if hasattr(thumbnail, "urls"): - batch_if_valid('thumbnail', f'=IMAGE("{thumbnail.urls[0]}")') + batch_if_valid("thumbnail", f'=IMAGE("{thumbnail.urls[0]}")') - if (browsertrix := item.get_media_by_id("browsertrix")): - batch_if_valid('wacz', "\n".join(browsertrix.urls)) - batch_if_valid('replaywebpage', "\n".join([f'https://replayweb.page/?source={quote(wacz)}#view=pages&url={quote(item.get_url())}' for wacz in browsertrix.urls])) + if browsertrix := item.get_media_by_id("browsertrix"): + batch_if_valid("wacz", "\n".join(browsertrix.urls)) + batch_if_valid( + "replaywebpage", + "\n".join( + [ + f"https://replayweb.page/?source={quote(wacz)}#view=pages&url={quote(item.get_url())}" + for wacz in browsertrix.urls + ] + ), + ) gw.batch_set_cell(cell_updates) - @staticmethod - def _get_current_datetime_iso() -> str: - """Helper method to generate the current datetime in ISO format.""" - return datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=datetime.timezone.utc).isoformat() - - def _safe_status_update(self, item: Metadata, new_status: str) -> None: try: gw, row = self._retrieve_gsheet(item) - gw.set_cell(row, 'status', new_status) + gw.set_cell(row, "status", new_status) except Exception as e: logger.debug(f"Unable to update sheet: {e}") diff --git a/src/auto_archiver/modules/telethon_extractor/telethon_extractor.py b/src/auto_archiver/modules/telethon_extractor/telethon_extractor.py index 0147ff2..947db9e 100644 --- a/src/auto_archiver/modules/telethon_extractor/telethon_extractor.py +++ b/src/auto_archiver/modules/telethon_extractor/telethon_extractor.py @@ -18,12 +18,14 @@ class TelethonExtractor(Extractor): invite_pattern = re.compile(r"t.me(\/joinchat){0,1}\/\+?(.+)") - def setup(self) -> None: + def setup(self, config: dict) -> None: + """ 1. makes a copy of session_file that is removed in cleanup 2. trigger login process for telegram or proceed if already saved in a session file 3. joins channel_invites where needed """ + super().setup(config) logger.info(f"SETUP {self.name} checking login...") # make a copy of the session that is used exclusively with this archiver instance diff --git a/src/auto_archiver/modules/whisper_enricher/__manifest__.py b/src/auto_archiver/modules/whisper_enricher/__manifest__.py index 1539df6..98e743e 100644 --- a/src/auto_archiver/modules/whisper_enricher/__manifest__.py +++ b/src/auto_archiver/modules/whisper_enricher/__manifest__.py @@ -1,4 +1,4 @@ -a={ +{ "name": "Whisper Enricher", "type": ["enricher"], "requires_setup": True, diff --git a/src/auto_archiver/modules/whisper_enricher/whisper_enricher.py b/src/auto_archiver/modules/whisper_enricher/whisper_enricher.py index 004d91c..a51ffc1 100644 --- a/src/auto_archiver/modules/whisper_enricher/whisper_enricher.py +++ b/src/auto_archiver/modules/whisper_enricher/whisper_enricher.py @@ -4,7 +4,6 @@ from loguru import logger from auto_archiver.core import Enricher from auto_archiver.core import Metadata, Media -from auto_archiver.modules.s3_storage import S3Storage from auto_archiver.core.module import get_module class WhisperEnricher(Enricher): @@ -14,13 +13,17 @@ class WhisperEnricher(Enricher): Only works if an S3 compatible storage is used """ - def enrich(self, to_enrich: Metadata) -> None: - storages = self.config['steps']['storages'] - if not "s3_storage" in storages: + def setup(self, config: dict) -> None: + super().setup(config) + self.stores = self.config['steps']['storages'] + self.s3 = get_module("s3_storage", self.config) + if not "s3_storage" in self.stores: logger.error("WhisperEnricher: To use the WhisperEnricher you need to use S3Storage so files are accessible publicly to the whisper service being called.") return - self.s3 = get_module("s3_storage", self.config) + + def enrich(self, to_enrich: Metadata) -> None: + url = to_enrich.get_url() logger.debug(f"WHISPER[{self.action}]: iterating media items for {url=}.") diff --git a/src/auto_archiver/utils/misc.py b/src/auto_archiver/utils/misc.py index 300a710..e4c214c 100644 --- a/src/auto_archiver/utils/misc.py +++ b/src/auto_archiver/utils/misc.py @@ -1,9 +1,7 @@ - - import os import json import uuid -from datetime import datetime +from datetime import datetime, timezone import requests from loguru import logger @@ -58,5 +56,37 @@ def random_str(length: int = 32) -> str: assert length <= 32, "length must be less than 32 as UUID4 is used" return str(uuid.uuid4()).replace("-", "")[:length] + def json_loader(cli_val): return json.loads(cli_val) + +def get_current_datetime_iso() -> str: + return datetime.now(timezone.utc).replace(tzinfo=timezone.utc).isoformat() + + +def get_datetime_from_str(dt_str: str, fmt: str | None = None) -> datetime | None: + # parse a datetime string with option of passing a specific format + try: + return datetime.strptime(dt_str, fmt) if fmt else datetime.fromisoformat(dt_str) + except ValueError as e: + logger.error(f"Unable to parse datestring {dt_str}: {e}") + return None + + +def get_timestamp(ts, utc=True, iso=True) -> str | datetime | None: + # Consistent parsing of timestamps + # If utc=True, the timezone is set to UTC, + # if iso=True, the output is an iso string + if not ts: return + try: + if isinstance(ts, str): ts = datetime.fromisoformat(ts) + if isinstance(ts, (int, float)): ts = datetime.fromtimestamp(ts) + if utc: ts = ts.replace(tzinfo=timezone.utc) + if iso: return ts.isoformat() + return ts + except Exception as e: + logger.error(f"Unable to parse timestamp {ts}: {e}") + return None + +def get_current_timestamp() -> str: + return get_timestamp(datetime.now()) \ No newline at end of file diff --git a/tests/databases/test_gsheet_db.py b/tests/databases/test_gsheet_db.py index bdc2811..0a655a8 100644 --- a/tests/databases/test_gsheet_db.py +++ b/tests/databases/test_gsheet_db.py @@ -103,19 +103,20 @@ def test_failed(gsheets_db, mock_metadata, mock_gworksheet): gsheets_db.failed(mock_metadata, reason) mock_gworksheet.set_cell.assert_called_once_with(1, 'status', f'Archive failed {reason}') + def test_aborted(gsheets_db, mock_metadata, mock_gworksheet): gsheets_db.aborted(mock_metadata) mock_gworksheet.set_cell.assert_called_once_with(1, 'status', '') def test_done(gsheets_db, metadata, mock_gworksheet, expected_calls): - with patch.object(gsheets_db, '_get_current_datetime_iso', return_value='2025-02-01T00:00:00+00:00'): + with patch("auto_archiver.modules.gsheet_db.gsheet_db.get_current_timestamp", return_value='2025-02-01T00:00:00+00:00'): gsheets_db.done(metadata) mock_gworksheet.batch_set_cell.assert_called_once_with(expected_calls) def test_done_cached(gsheets_db, metadata, mock_gworksheet): - with patch.object(gsheets_db, '_get_current_datetime_iso', return_value='2025-02-01T00:00:00+00:00'): + with patch("auto_archiver.modules.gsheet_db.gsheet_db.get_current_timestamp", return_value='2025-02-01T00:00:00+00:00'): gsheets_db.done(metadata, cached=True) # Verify the status message includes "[cached]" @@ -126,7 +127,8 @@ def test_done_cached(gsheets_db, metadata, mock_gworksheet): def test_done_missing_media(gsheets_db, metadata, mock_gworksheet): # clear media from metadata metadata.media = [] - with patch.object(gsheets_db, '_get_current_datetime_iso', return_value='2025-02-01T00:00:00+00:00'): + with patch("auto_archiver.modules.gsheet_db.gsheet_db.get_current_timestamp", + return_value='2025-02-01T00:00:00+00:00'): gsheets_db.done(metadata) # Verify nothing media-related gets updated call_args = mock_gworksheet.batch_set_cell.call_args[0][0] diff --git a/tests/extractors/test_instagram_api_extractor.py b/tests/extractors/test_instagram_api_extractor.py index d3f7bd6..c119e3f 100644 --- a/tests/extractors/test_instagram_api_extractor.py +++ b/tests/extractors/test_instagram_api_extractor.py @@ -185,5 +185,4 @@ class TestInstagramAPIExtractor(TestExtractorBase): result = self.extractor.download_profile(metadata, "test_user") assert result.is_success() - assert "Error downloading stories for test_user" in result.metadata["errors"] - # assert "Error downloading posts for test_user" in result.metadata["errors"] \ No newline at end of file + assert "Error downloading stories for test_user" in result.metadata["errors"] \ No newline at end of file diff --git a/tests/extractors/test_instagram_tbot_extractor.py b/tests/extractors/test_instagram_tbot_extractor.py index b82641d..d7a1e53 100644 --- a/tests/extractors/test_instagram_tbot_extractor.py +++ b/tests/extractors/test_instagram_tbot_extractor.py @@ -1,5 +1,4 @@ import os -import pickle from typing import Type from unittest.mock import patch, MagicMock diff --git a/tests/feeders/test_gsheet_feeder.py b/tests/feeders/test_gsheet_feeder.py index 103610e..ecf57f1 100644 --- a/tests/feeders/test_gsheet_feeder.py +++ b/tests/feeders/test_gsheet_feeder.py @@ -7,10 +7,8 @@ from auto_archiver.modules.gsheet_feeder import GsheetsFeeder from auto_archiver.core import Metadata, Feeder -def test_initialise_without_sheet_and_sheet_id(setup_module): - """Ensure initialise() raises AssertionError if neither sheet nor sheet_id is set. - (shouldn't really be asserting in there) - """ +def test_setup_without_sheet_and_sheet_id(setup_module): + # Ensure setup() raises AssertionError if neither sheet nor sheet_id is set. with patch("gspread.service_account"): with pytest.raises(AssertionError): setup_module( @@ -145,7 +143,6 @@ def test_open_sheet_with_name_or_id( "gsheet_feeder", {"service_account": "dummy.json", "sheet": sheet, "sheet_id": sheet_id}, ) - feeder.initialise() sheet_result = feeder.open_sheet() # Validate the correct method was called getattr(mock_client, expected_method).assert_called_once_with( @@ -165,7 +162,6 @@ def test_open_sheet_with_sheet_id(setup_module): "gsheet_feeder", {"service_account": "dummy.json", "sheet": None, "sheet_id": "ABC123"}, ) - feeder.initialise() sheet = feeder.open_sheet() mock_client.open_by_key.assert_called_once_with("ABC123") assert sheet == "MockSheet" @@ -263,7 +259,6 @@ class TestGSheetsFeederReal: ["https://example.com", "done"], ] worksheet.append_rows(test_rows) - self.feeder.initialise() metadata_list = list(self.feeder) # Validate that only the first row is processed diff --git a/tests/storages/test_gdrive_storage.py b/tests/storages/test_gdrive_storage.py index b7417ad..4259cb2 100644 --- a/tests/storages/test_gdrive_storage.py +++ b/tests/storages/test_gdrive_storage.py @@ -21,16 +21,6 @@ class TestGDriveStorage(TestStorageBase): 'service_account': 'fake_service_account.json' } - @pytest.mark.skip(reason="Requires real credentials") - @pytest.mark.download - def test_initialize_with_real_credentials(self): - """ - Test that the Google Drive service can be initialized with real credentials. - """ - self.storage.service_account = 'secrets/service_account.json' # Path to real credentials - self.storage.initialise() - assert self.storage.service is not None - def test_initialize_fails_with_non_existent_creds(self): """ @@ -38,6 +28,35 @@ class TestGDriveStorage(TestStorageBase): """ # Act and Assert with pytest.raises(FileNotFoundError) as exc_info: - self.storage.initialise() + self.storage.setup(self.config) assert "No such file or directory" in str(exc_info.value) + def test_path_parts(self): + media = Media(filename="test.jpg") + media.key = "folder1/folder2/test.jpg" + +# @pytest.mark.skip(reason="Requires real credentials") +@pytest.mark.download +class TestGDriveStorageConnected(TestStorageBase): + """ + 'Real' tests for GDriveStorage. + """ + + module_name: str = "gdrive_storage" + storage: Type[GDriveStorage] + config: dict = {'path_generator': 'url', + 'filename_generator': 'static', + # TODO: replace with real root folder id + 'root_folder_id': "1TVY_oJt95_dmRSEdP9m5zFy7l50TeCSk", + 'oauth_token': None, + 'service_account': 'secrets/service_account.json' + } + + + def test_initialize_with_real_credentials(self): + """ + Test that the Google Drive service can be initialized with real credentials. + """ + assert self.storage.service is not None + + diff --git a/tests/test_metadata.py b/tests/test_metadata.py index 7270c80..b07e107 100644 --- a/tests/test_metadata.py +++ b/tests/test_metadata.py @@ -159,3 +159,7 @@ def test_get_context(): assert m.get_context("somekey") == "somevalue" assert m.get_context("anotherkey") == "anothervalue" assert len(m._context) == 2 + + +def test_choose_most_complete(): + pass \ No newline at end of file From ab6cf52533c29b7c5815c94c8c27ca60b32f8ad7 Mon Sep 17 00:00:00 2001 From: msramalho <19508417+msramalho@users.noreply.github.com> Date: Mon, 10 Feb 2025 16:45:28 +0000 Subject: [PATCH 42/62] fixes bad hash initialization --- src/auto_archiver/modules/hash_enricher/hash_enricher.py | 8 ++++---- src/auto_archiver/utils/misc.py | 7 ++++--- 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/src/auto_archiver/modules/hash_enricher/hash_enricher.py b/src/auto_archiver/modules/hash_enricher/hash_enricher.py index b3ca8be..7a0587c 100644 --- a/src/auto_archiver/modules/hash_enricher/hash_enricher.py +++ b/src/auto_archiver/modules/hash_enricher/hash_enricher.py @@ -30,10 +30,10 @@ class HashEnricher(Enricher): to_enrich.media[i].set("hash", f"{self.algorithm}:{hd}") def calculate_hash(self, filename) -> str: - hash = None + hash_algo = None if self.algorithm == "SHA-256": - hash = hashlib.sha256() + hash_algo = hashlib.sha256 elif self.algorithm == "SHA3-512": - hash = hashlib.sha3_512() + hash_algo = hashlib.sha3_512 else: return "" - return calculate_file_hash(filename, hash, self.chunksize) + return calculate_file_hash(filename, hash_algo, self.chunksize) diff --git a/src/auto_archiver/utils/misc.py b/src/auto_archiver/utils/misc.py index 3af5a54..4d48372 100644 --- a/src/auto_archiver/utils/misc.py +++ b/src/auto_archiver/utils/misc.py @@ -65,10 +65,11 @@ def json_loader(cli_val): return json.loads(cli_val) -def calculate_file_hash(filename: str, hash_algo = hashlib.sha256(), chunksize: int = 16000000) -> str: +def calculate_file_hash(filename: str, hash_algo = hashlib.sha256, chunksize: int = 16000000) -> str: + hash = hash_algo() with open(filename, "rb") as f: while True: buf = f.read(chunksize) if not buf: break - hash_algo.update(buf) - return hash_algo.hexdigest() + hash.update(buf) + return hash.hexdigest() From 12f14cccc933d44b02906b2b9e239d1dd98af036 Mon Sep 17 00:00:00 2001 From: msramalho <19508417+msramalho@users.noreply.github.com> Date: Mon, 10 Feb 2025 16:58:35 +0000 Subject: [PATCH 43/62] fixes gsheet feeder<->db connection via context. --- src/auto_archiver/core/storage.py | 2 +- src/auto_archiver/modules/gsheet_feeder/gsheet_feeder.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/auto_archiver/core/storage.py b/src/auto_archiver/core/storage.py index 9f355f6..5dfa39d 100644 --- a/src/auto_archiver/core/storage.py +++ b/src/auto_archiver/core/storage.py @@ -1,6 +1,6 @@ from __future__ import annotations from abc import abstractmethod -from typing import IO, Optional +from typing import IO import os from loguru import logger diff --git a/src/auto_archiver/modules/gsheet_feeder/gsheet_feeder.py b/src/auto_archiver/modules/gsheet_feeder/gsheet_feeder.py index d129182..1b81385 100644 --- a/src/auto_archiver/modules/gsheet_feeder/gsheet_feeder.py +++ b/src/auto_archiver/modules/gsheet_feeder/gsheet_feeder.py @@ -68,7 +68,7 @@ class GsheetsFeeder(Feeder): folder = os.path.join(folder, slugify(self.sheet), slugify(wks.title)) m.set_context('folder', folder) - m.set_context('worksheet', {"row": row, "worksheet": gw}) + m.set_context('gsheet', {"row": row, "worksheet": gw}) yield m logger.success(f'Finished worksheet {wks.title}') From 2c3d1f591f4a721597e2cd9906c1cdc05db8a78e Mon Sep 17 00:00:00 2001 From: erinhmclark Date: Mon, 10 Feb 2025 17:25:15 +0000 Subject: [PATCH 44/62] Separate setup() and module_setup(). --- src/auto_archiver/core/base_module.py | 4 ++++ src/auto_archiver/core/module.py | 1 + src/auto_archiver/modules/gdrive_storage/gdrive_storage.py | 4 +--- src/auto_archiver/modules/gsheet_feeder/gsheet_feeder.py | 3 +-- src/auto_archiver/modules/html_formatter/html_formatter.py | 3 +-- .../instagram_api_extractor/instagram_api_extractor.py | 3 +-- .../modules/instagram_extractor/instagram_extractor.py | 3 +-- .../instagram_tbot_extractor/instagram_tbot_extractor.py | 3 +-- src/auto_archiver/modules/s3_storage/s3_storage.py | 3 +-- .../modules/telethon_extractor/telethon_extractor.py | 3 +-- .../modules/twitter_api_extractor/twitter_api_extractor.py | 4 +--- src/auto_archiver/modules/vk_extractor/vk_extractor.py | 3 +-- src/auto_archiver/modules/wacz_enricher/wacz_enricher.py | 3 +-- .../modules/whisper_enricher/whisper_enricher.py | 3 +-- 14 files changed, 17 insertions(+), 26 deletions(-) diff --git a/src/auto_archiver/core/base_module.py b/src/auto_archiver/core/base_module.py index 5c6ecbb..95575e3 100644 --- a/src/auto_archiver/core/base_module.py +++ b/src/auto_archiver/core/base_module.py @@ -80,6 +80,10 @@ class BaseModule(ABC): for key, val in config.get(self.name, {}).items(): setattr(self, key, val) + def module_setup(self): + # For any additional setup required by modules, e.g. autehntication + pass + def auth_for_site(self, site: str, extract_cookies=True) -> Mapping[str, Any]: """ Returns the authentication information for a given site. This is used to authenticate diff --git a/src/auto_archiver/core/module.py b/src/auto_archiver/core/module.py index f3fbec5..69f9fcc 100644 --- a/src/auto_archiver/core/module.py +++ b/src/auto_archiver/core/module.py @@ -242,6 +242,7 @@ class LazyBaseModule: default_config = dict((k, v['default']) for k, v in self.configs.items() if v.get('default')) config[self.name] = default_config | config.get(self.name, {}) instance.setup(config) + instance.module_setup() return instance def __repr__(self): diff --git a/src/auto_archiver/modules/gdrive_storage/gdrive_storage.py b/src/auto_archiver/modules/gdrive_storage/gdrive_storage.py index 910f48b..51c13c2 100644 --- a/src/auto_archiver/modules/gdrive_storage/gdrive_storage.py +++ b/src/auto_archiver/modules/gdrive_storage/gdrive_storage.py @@ -19,9 +19,7 @@ from auto_archiver.core import Storage class GDriveStorage(Storage): - def setup(self, config: dict) -> None: - # Step 1: Call the BaseModule setup to dynamically assign configs - super().setup(config) + def module_setup(self) -> None: self.scopes = ['https://www.googleapis.com/auth/drive'] # Initialize Google Drive service self._setup_google_drive_service() diff --git a/src/auto_archiver/modules/gsheet_feeder/gsheet_feeder.py b/src/auto_archiver/modules/gsheet_feeder/gsheet_feeder.py index 50bf430..dd98032 100644 --- a/src/auto_archiver/modules/gsheet_feeder/gsheet_feeder.py +++ b/src/auto_archiver/modules/gsheet_feeder/gsheet_feeder.py @@ -21,8 +21,7 @@ from . import GWorksheet class GsheetsFeeder(Feeder): - def setup(self, config: dict): - super().setup(config) + def module_setup(self) -> None: self.gsheets_client = gspread.service_account(filename=self.service_account) # TODO mv to validators assert self.sheet or self.sheet_id, ( diff --git a/src/auto_archiver/modules/html_formatter/html_formatter.py b/src/auto_archiver/modules/html_formatter/html_formatter.py index 4da82c8..bbba097 100644 --- a/src/auto_archiver/modules/html_formatter/html_formatter.py +++ b/src/auto_archiver/modules/html_formatter/html_formatter.py @@ -17,9 +17,8 @@ class HtmlFormatter(Formatter): environment: Environment = None template: any = None - def setup(self, config: dict) -> None: + def module_setup(self) -> None: """Sets up the Jinja2 environment and loads the template.""" - super().setup(config) # Ensure the base class logic is executed template_dir = os.path.join(pathlib.Path(__file__).parent.resolve(), "templates/") self.environment = Environment(loader=FileSystemLoader(template_dir), autoescape=True) diff --git a/src/auto_archiver/modules/instagram_api_extractor/instagram_api_extractor.py b/src/auto_archiver/modules/instagram_api_extractor/instagram_api_extractor.py index 5dad0ba..367cc75 100644 --- a/src/auto_archiver/modules/instagram_api_extractor/instagram_api_extractor.py +++ b/src/auto_archiver/modules/instagram_api_extractor/instagram_api_extractor.py @@ -32,8 +32,7 @@ class InstagramAPIExtractor(Extractor): r"(?:(?:http|https):\/\/)?(?:www.)?(?:instagram.com)\/(stories(?:\/highlights)?|p|reel)?\/?([^\/\?]*)\/?(\d+)?" ) - def setup(self, config: dict) -> None: - super().setup(config) + def module_setup(self) -> None: if self.api_endpoint[-1] == "/": self.api_endpoint = self.api_endpoint[:-1] diff --git a/src/auto_archiver/modules/instagram_extractor/instagram_extractor.py b/src/auto_archiver/modules/instagram_extractor/instagram_extractor.py index 3cf0362..e4e210f 100644 --- a/src/auto_archiver/modules/instagram_extractor/instagram_extractor.py +++ b/src/auto_archiver/modules/instagram_extractor/instagram_extractor.py @@ -25,8 +25,7 @@ class InstagramExtractor(Extractor): profile_pattern = re.compile(r"{valid_url}(\w+)".format(valid_url=valid_url)) # TODO: links to stories - def setup(self, config: dict) -> None: - super().setup(config) + def module_setup(self) -> None: self.insta = instaloader.Instaloader( download_geotags=True, download_comments=True, compress_json=False, dirname_pattern=self.download_folder, filename_pattern="{date_utc}_UTC_{target}__{typename}" diff --git a/src/auto_archiver/modules/instagram_tbot_extractor/instagram_tbot_extractor.py b/src/auto_archiver/modules/instagram_tbot_extractor/instagram_tbot_extractor.py index 5660cd2..707dcc3 100644 --- a/src/auto_archiver/modules/instagram_tbot_extractor/instagram_tbot_extractor.py +++ b/src/auto_archiver/modules/instagram_tbot_extractor/instagram_tbot_extractor.py @@ -27,12 +27,11 @@ class InstagramTbotExtractor(Extractor): https://t.me/instagram_load_bot """ - def setup(self, configs) -> None: + def module_setup(self) -> None: """ 1. makes a copy of session_file that is removed in cleanup 2. checks if the session file is valid """ - super().setup(configs) logger.info(f"SETUP {self.name} checking login...") self._prepare_session_file() self._initialize_telegram_client() diff --git a/src/auto_archiver/modules/s3_storage/s3_storage.py b/src/auto_archiver/modules/s3_storage/s3_storage.py index 2f85164..c77bbc3 100644 --- a/src/auto_archiver/modules/s3_storage/s3_storage.py +++ b/src/auto_archiver/modules/s3_storage/s3_storage.py @@ -13,8 +13,7 @@ NO_DUPLICATES_FOLDER = "no-dups/" class S3Storage(Storage): - def setup(self, config: dict) -> None: - super().setup(config) + def module_setup(self) -> None: self.s3 = boto3.client( 's3', region_name=self.region, diff --git a/src/auto_archiver/modules/telethon_extractor/telethon_extractor.py b/src/auto_archiver/modules/telethon_extractor/telethon_extractor.py index 97d3e94..3762f01 100644 --- a/src/auto_archiver/modules/telethon_extractor/telethon_extractor.py +++ b/src/auto_archiver/modules/telethon_extractor/telethon_extractor.py @@ -18,14 +18,13 @@ class TelethonExtractor(Extractor): invite_pattern = re.compile(r"t.me(\/joinchat){0,1}\/\+?(.+)") - def setup(self, config: dict) -> None: + def module_setup(self) -> None: """ 1. makes a copy of session_file that is removed in cleanup 2. trigger login process for telegram or proceed if already saved in a session file 3. joins channel_invites where needed """ - super().setup(config) logger.info(f"SETUP {self.name} checking login...") # make a copy of the session that is used exclusively with this archiver instance diff --git a/src/auto_archiver/modules/twitter_api_extractor/twitter_api_extractor.py b/src/auto_archiver/modules/twitter_api_extractor/twitter_api_extractor.py index 6573475..0b27e22 100644 --- a/src/auto_archiver/modules/twitter_api_extractor/twitter_api_extractor.py +++ b/src/auto_archiver/modules/twitter_api_extractor/twitter_api_extractor.py @@ -15,9 +15,7 @@ class TwitterApiExtractor(Extractor): valid_url: re.Pattern = re.compile(r"(?:twitter|x).com\/(?:\#!\/)?(\w+)\/status(?:es)?\/(\d+)") - def setup(self, config: dict) -> None: - super().setup(config) - + def module_setup(self) -> None: self.api_index = 0 self.apis = [] if len(self.bearer_tokens): diff --git a/src/auto_archiver/modules/vk_extractor/vk_extractor.py b/src/auto_archiver/modules/vk_extractor/vk_extractor.py index 2d09138..0d1fc04 100644 --- a/src/auto_archiver/modules/vk_extractor/vk_extractor.py +++ b/src/auto_archiver/modules/vk_extractor/vk_extractor.py @@ -12,8 +12,7 @@ class VkExtractor(Extractor): Currently only works for /wall posts """ - def setup(self, config: dict) -> None: - super().setup(config) + def module_setup(self) -> None: self.vks = VkScraper(self.username, self.password, session_file=self.session_file) def download(self, item: Metadata) -> Metadata: diff --git a/src/auto_archiver/modules/wacz_enricher/wacz_enricher.py b/src/auto_archiver/modules/wacz_enricher/wacz_enricher.py index 1586b75..7d91f43 100644 --- a/src/auto_archiver/modules/wacz_enricher/wacz_enricher.py +++ b/src/auto_archiver/modules/wacz_enricher/wacz_enricher.py @@ -18,8 +18,7 @@ class WaczExtractorEnricher(Enricher, Extractor): When used as an archiver it will extract the media from the .WACZ archive so it can be enriched. """ - def setup(self, configs) -> None: - super().setup(configs) + def module_setup(self) -> None: self.use_docker = os.environ.get('WACZ_ENABLE_DOCKER') or not os.environ.get('RUNNING_IN_DOCKER') self.docker_in_docker = os.environ.get('WACZ_ENABLE_DOCKER') and os.environ.get('RUNNING_IN_DOCKER') diff --git a/src/auto_archiver/modules/whisper_enricher/whisper_enricher.py b/src/auto_archiver/modules/whisper_enricher/whisper_enricher.py index a51ffc1..d83319e 100644 --- a/src/auto_archiver/modules/whisper_enricher/whisper_enricher.py +++ b/src/auto_archiver/modules/whisper_enricher/whisper_enricher.py @@ -13,8 +13,7 @@ class WhisperEnricher(Enricher): Only works if an S3 compatible storage is used """ - def setup(self, config: dict) -> None: - super().setup(config) + def module_setup(self) -> None: self.stores = self.config['steps']['storages'] self.s3 = get_module("s3_storage", self.config) if not "s3_storage" in self.stores: From e97ccf8a736fc6bd01a0efdf9a54c8cca16d5d97 Mon Sep 17 00:00:00 2001 From: erinhmclark Date: Mon, 10 Feb 2025 18:07:47 +0000 Subject: [PATCH 45/62] Separate setup() and module_setup(). --- src/auto_archiver/core/base_module.py | 6 +++--- src/auto_archiver/core/module.py | 6 +++--- src/auto_archiver/modules/gdrive_storage/gdrive_storage.py | 2 +- src/auto_archiver/modules/gsheet_feeder/gsheet_feeder.py | 2 +- src/auto_archiver/modules/html_formatter/html_formatter.py | 2 +- .../instagram_api_extractor/instagram_api_extractor.py | 2 +- .../modules/instagram_extractor/instagram_extractor.py | 2 +- .../instagram_tbot_extractor/instagram_tbot_extractor.py | 2 +- src/auto_archiver/modules/s3_storage/s3_storage.py | 2 +- .../modules/telethon_extractor/telethon_extractor.py | 2 +- .../modules/twitter_api_extractor/twitter_api_extractor.py | 2 +- src/auto_archiver/modules/vk_extractor/vk_extractor.py | 2 +- src/auto_archiver/modules/wacz_enricher/wacz_enricher.py | 2 +- .../modules/whisper_enricher/whisper_enricher.py | 2 +- 14 files changed, 18 insertions(+), 18 deletions(-) diff --git a/src/auto_archiver/core/base_module.py b/src/auto_archiver/core/base_module.py index 95575e3..ece4719 100644 --- a/src/auto_archiver/core/base_module.py +++ b/src/auto_archiver/core/base_module.py @@ -14,7 +14,7 @@ class BaseModule(ABC): Base module class. All modules should inherit from this class. The exact methods a class implements will depend on the type of module it is, - however all modules have a .setup(config: dict) method to run any setup code + however modules can have a .setup() method to run any setup code (e.g. logging in to a site, spinning up a browser etc.) See BaseModule.MODULE_TYPES for the types of modules you can create, noting that @@ -60,7 +60,7 @@ class BaseModule(ABC): def storages(self) -> list: return self.config.get('storages', []) - def setup(self, config: dict): + def config_setup(self, config: dict): authentication = config.get('authentication', {}) # extract out concatenated sites @@ -80,7 +80,7 @@ class BaseModule(ABC): for key, val in config.get(self.name, {}).items(): setattr(self, key, val) - def module_setup(self): + def setup(self): # For any additional setup required by modules, e.g. autehntication pass diff --git a/src/auto_archiver/core/module.py b/src/auto_archiver/core/module.py index 69f9fcc..c81e26a 100644 --- a/src/auto_archiver/core/module.py +++ b/src/auto_archiver/core/module.py @@ -58,7 +58,7 @@ def get_module_lazy(module_name: str, suppress_warnings: bool = False) -> LazyBa This has all the information about the module, but does not load the module itself or its dependencies - To load an actual module, call .setup() on a laz module + To load an actual module, call .setup() on a lazy module """ if module_name in _LAZY_LOADED_MODULES: @@ -241,8 +241,8 @@ class LazyBaseModule: # merge the default config with the user config default_config = dict((k, v['default']) for k, v in self.configs.items() if v.get('default')) config[self.name] = default_config | config.get(self.name, {}) - instance.setup(config) - instance.module_setup() + instance.config_setup(config) + instance.setup() return instance def __repr__(self): diff --git a/src/auto_archiver/modules/gdrive_storage/gdrive_storage.py b/src/auto_archiver/modules/gdrive_storage/gdrive_storage.py index 51c13c2..f38feb6 100644 --- a/src/auto_archiver/modules/gdrive_storage/gdrive_storage.py +++ b/src/auto_archiver/modules/gdrive_storage/gdrive_storage.py @@ -19,7 +19,7 @@ from auto_archiver.core import Storage class GDriveStorage(Storage): - def module_setup(self) -> None: + def setup(self) -> None: self.scopes = ['https://www.googleapis.com/auth/drive'] # Initialize Google Drive service self._setup_google_drive_service() diff --git a/src/auto_archiver/modules/gsheet_feeder/gsheet_feeder.py b/src/auto_archiver/modules/gsheet_feeder/gsheet_feeder.py index dd98032..8612d02 100644 --- a/src/auto_archiver/modules/gsheet_feeder/gsheet_feeder.py +++ b/src/auto_archiver/modules/gsheet_feeder/gsheet_feeder.py @@ -21,7 +21,7 @@ from . import GWorksheet class GsheetsFeeder(Feeder): - def module_setup(self) -> None: + def setup(self) -> None: self.gsheets_client = gspread.service_account(filename=self.service_account) # TODO mv to validators assert self.sheet or self.sheet_id, ( diff --git a/src/auto_archiver/modules/html_formatter/html_formatter.py b/src/auto_archiver/modules/html_formatter/html_formatter.py index bbba097..3691735 100644 --- a/src/auto_archiver/modules/html_formatter/html_formatter.py +++ b/src/auto_archiver/modules/html_formatter/html_formatter.py @@ -17,7 +17,7 @@ class HtmlFormatter(Formatter): environment: Environment = None template: any = None - def module_setup(self) -> None: + def setup(self) -> None: """Sets up the Jinja2 environment and loads the template.""" template_dir = os.path.join(pathlib.Path(__file__).parent.resolve(), "templates/") self.environment = Environment(loader=FileSystemLoader(template_dir), autoescape=True) diff --git a/src/auto_archiver/modules/instagram_api_extractor/instagram_api_extractor.py b/src/auto_archiver/modules/instagram_api_extractor/instagram_api_extractor.py index 367cc75..a75e065 100644 --- a/src/auto_archiver/modules/instagram_api_extractor/instagram_api_extractor.py +++ b/src/auto_archiver/modules/instagram_api_extractor/instagram_api_extractor.py @@ -32,7 +32,7 @@ class InstagramAPIExtractor(Extractor): r"(?:(?:http|https):\/\/)?(?:www.)?(?:instagram.com)\/(stories(?:\/highlights)?|p|reel)?\/?([^\/\?]*)\/?(\d+)?" ) - def module_setup(self) -> None: + def setup(self) -> None: if self.api_endpoint[-1] == "/": self.api_endpoint = self.api_endpoint[:-1] diff --git a/src/auto_archiver/modules/instagram_extractor/instagram_extractor.py b/src/auto_archiver/modules/instagram_extractor/instagram_extractor.py index e4e210f..0af2c32 100644 --- a/src/auto_archiver/modules/instagram_extractor/instagram_extractor.py +++ b/src/auto_archiver/modules/instagram_extractor/instagram_extractor.py @@ -25,7 +25,7 @@ class InstagramExtractor(Extractor): profile_pattern = re.compile(r"{valid_url}(\w+)".format(valid_url=valid_url)) # TODO: links to stories - def module_setup(self) -> None: + def setup(self) -> None: self.insta = instaloader.Instaloader( download_geotags=True, download_comments=True, compress_json=False, dirname_pattern=self.download_folder, filename_pattern="{date_utc}_UTC_{target}__{typename}" diff --git a/src/auto_archiver/modules/instagram_tbot_extractor/instagram_tbot_extractor.py b/src/auto_archiver/modules/instagram_tbot_extractor/instagram_tbot_extractor.py index 707dcc3..d4b7a8e 100644 --- a/src/auto_archiver/modules/instagram_tbot_extractor/instagram_tbot_extractor.py +++ b/src/auto_archiver/modules/instagram_tbot_extractor/instagram_tbot_extractor.py @@ -27,7 +27,7 @@ class InstagramTbotExtractor(Extractor): https://t.me/instagram_load_bot """ - def module_setup(self) -> None: + def setup(self) -> None: """ 1. makes a copy of session_file that is removed in cleanup 2. checks if the session file is valid diff --git a/src/auto_archiver/modules/s3_storage/s3_storage.py b/src/auto_archiver/modules/s3_storage/s3_storage.py index c77bbc3..6590ac9 100644 --- a/src/auto_archiver/modules/s3_storage/s3_storage.py +++ b/src/auto_archiver/modules/s3_storage/s3_storage.py @@ -13,7 +13,7 @@ NO_DUPLICATES_FOLDER = "no-dups/" class S3Storage(Storage): - def module_setup(self) -> None: + def setup(self) -> None: self.s3 = boto3.client( 's3', region_name=self.region, diff --git a/src/auto_archiver/modules/telethon_extractor/telethon_extractor.py b/src/auto_archiver/modules/telethon_extractor/telethon_extractor.py index 3762f01..65ea8cd 100644 --- a/src/auto_archiver/modules/telethon_extractor/telethon_extractor.py +++ b/src/auto_archiver/modules/telethon_extractor/telethon_extractor.py @@ -18,7 +18,7 @@ class TelethonExtractor(Extractor): invite_pattern = re.compile(r"t.me(\/joinchat){0,1}\/\+?(.+)") - def module_setup(self) -> None: + def setup(self) -> None: """ 1. makes a copy of session_file that is removed in cleanup diff --git a/src/auto_archiver/modules/twitter_api_extractor/twitter_api_extractor.py b/src/auto_archiver/modules/twitter_api_extractor/twitter_api_extractor.py index 0b27e22..72fd2f2 100644 --- a/src/auto_archiver/modules/twitter_api_extractor/twitter_api_extractor.py +++ b/src/auto_archiver/modules/twitter_api_extractor/twitter_api_extractor.py @@ -15,7 +15,7 @@ class TwitterApiExtractor(Extractor): valid_url: re.Pattern = re.compile(r"(?:twitter|x).com\/(?:\#!\/)?(\w+)\/status(?:es)?\/(\d+)") - def module_setup(self) -> None: + def setup(self) -> None: self.api_index = 0 self.apis = [] if len(self.bearer_tokens): diff --git a/src/auto_archiver/modules/vk_extractor/vk_extractor.py b/src/auto_archiver/modules/vk_extractor/vk_extractor.py index 0d1fc04..99527c4 100644 --- a/src/auto_archiver/modules/vk_extractor/vk_extractor.py +++ b/src/auto_archiver/modules/vk_extractor/vk_extractor.py @@ -12,7 +12,7 @@ class VkExtractor(Extractor): Currently only works for /wall posts """ - def module_setup(self) -> None: + def setup(self) -> None: self.vks = VkScraper(self.username, self.password, session_file=self.session_file) def download(self, item: Metadata) -> Metadata: diff --git a/src/auto_archiver/modules/wacz_enricher/wacz_enricher.py b/src/auto_archiver/modules/wacz_enricher/wacz_enricher.py index 7d91f43..c324c62 100644 --- a/src/auto_archiver/modules/wacz_enricher/wacz_enricher.py +++ b/src/auto_archiver/modules/wacz_enricher/wacz_enricher.py @@ -18,7 +18,7 @@ class WaczExtractorEnricher(Enricher, Extractor): When used as an archiver it will extract the media from the .WACZ archive so it can be enriched. """ - def module_setup(self) -> None: + def setup(self) -> None: self.use_docker = os.environ.get('WACZ_ENABLE_DOCKER') or not os.environ.get('RUNNING_IN_DOCKER') self.docker_in_docker = os.environ.get('WACZ_ENABLE_DOCKER') and os.environ.get('RUNNING_IN_DOCKER') diff --git a/src/auto_archiver/modules/whisper_enricher/whisper_enricher.py b/src/auto_archiver/modules/whisper_enricher/whisper_enricher.py index d83319e..89579f9 100644 --- a/src/auto_archiver/modules/whisper_enricher/whisper_enricher.py +++ b/src/auto_archiver/modules/whisper_enricher/whisper_enricher.py @@ -13,7 +13,7 @@ class WhisperEnricher(Enricher): Only works if an S3 compatible storage is used """ - def module_setup(self) -> None: + def setup(self) -> None: self.stores = self.config['steps']['storages'] self.s3 = get_module("s3_storage", self.config) if not "s3_storage" in self.stores: From 3dae2337a1e3a97b913780b58e45adbc1d0aff5a Mon Sep 17 00:00:00 2001 From: erinhmclark Date: Mon, 10 Feb 2025 18:56:46 +0000 Subject: [PATCH 46/62] remove cdn_url check before storage. --- src/auto_archiver/core/media.py | 2 +- src/auto_archiver/modules/gdrive_storage/gdrive_storage.py | 5 ++--- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/src/auto_archiver/core/media.py b/src/auto_archiver/core/media.py index 952a025..b6820ab 100644 --- a/src/auto_archiver/core/media.py +++ b/src/auto_archiver/core/media.py @@ -65,7 +65,7 @@ class Media: def is_stored(self, in_storage) -> bool: # checks if the media is already stored in the given storage - return len(self.urls) > 0 and any([u for u in self.urls if in_storage.get_cdn_url(self) in u]) + return len(self.urls) > 0 and len(self.urls) == len(in_storage.config["steps"]["storages"]) def set(self, key: str, value: Any) -> Media: self.properties[key] = value diff --git a/src/auto_archiver/modules/gdrive_storage/gdrive_storage.py b/src/auto_archiver/modules/gdrive_storage/gdrive_storage.py index f38feb6..4971030 100644 --- a/src/auto_archiver/modules/gdrive_storage/gdrive_storage.py +++ b/src/auto_archiver/modules/gdrive_storage/gdrive_storage.py @@ -68,11 +68,10 @@ class GDriveStorage(Storage): filename = path_parts[-1] logger.info(f"looking for folders for {path_parts[0:-1]} before getting url for {filename=}") for folder in path_parts[0:-1]: - folder_id = self._get_id_from_parent_and_name(parent_id, folder, use_mime_type=True, raise_on_missing=False) + folder_id = self._get_id_from_parent_and_name(parent_id, folder, use_mime_type=True, raise_on_missing=True) parent_id = folder_id # get id of file inside folder (or sub folder) - # TODO: supressing the error as being checked before first upload - file_id = self._get_id_from_parent_and_name(folder_id, filename, raise_on_missing=False) + file_id = self._get_id_from_parent_and_name(folder_id, filename, raise_on_missing=True) if not file_id: # logger.info(f"file {filename} not found in folder {folder_id}") From ed81dcdaf081613b44035fae9d2b9de9d6fbc5b1 Mon Sep 17 00:00:00 2001 From: Patrick Robertson Date: Mon, 10 Feb 2025 23:07:03 +0000 Subject: [PATCH 47/62] Remove dangling 'b = ' from config.py --- src/auto_archiver/core/config.py | 14 ++++---------- src/auto_archiver/core/orchestrator.py | 4 ++-- 2 files changed, 6 insertions(+), 12 deletions(-) diff --git a/src/auto_archiver/core/config.py b/src/auto_archiver/core/config.py index 8f36c54..9bb080f 100644 --- a/src/auto_archiver/core/config.py +++ b/src/auto_archiver/core/config.py @@ -15,15 +15,9 @@ from .module import BaseModule from typing import Any, List, Type, Tuple -yaml: YAML = YAML() +_yaml: YAML = YAML() -b = yaml.load(""" - # This is a comment - site.com,site2.com: - key: value - key2: value2 - """) -EMPTY_CONFIG = yaml.load(""" +EMPTY_CONFIG = _yaml.load(""" # Auto Archiver Configuration # Steps are the modules that will be run in the order they are defined @@ -149,7 +143,7 @@ def read_yaml(yaml_filename: str) -> CommentedMap: config = None try: with open(yaml_filename, "r", encoding="utf-8") as inf: - config = yaml.load(inf) + config = _yaml.load(inf) except FileNotFoundError: pass @@ -166,4 +160,4 @@ def store_yaml(config: CommentedMap, yaml_filename: str) -> None: config_to_save.pop('urls', None) with open(yaml_filename, "w", encoding="utf-8") as outf: - yaml.dump(config_to_save, outf) \ No newline at end of file + _yaml.dump(config_to_save, outf) \ No newline at end of file diff --git a/src/auto_archiver/core/orchestrator.py b/src/auto_archiver/core/orchestrator.py index a451443..473f882 100644 --- a/src/auto_archiver/core/orchestrator.py +++ b/src/auto_archiver/core/orchestrator.py @@ -20,7 +20,7 @@ from rich_argparse import RichHelpFormatter from .metadata import Metadata, Media from auto_archiver.version import __version__ -from .config import yaml, read_yaml, store_yaml, to_dot_notation, merge_dicts, EMPTY_CONFIG, DefaultValidatingParser +from .config import _yaml, read_yaml, store_yaml, to_dot_notation, merge_dicts, EMPTY_CONFIG, DefaultValidatingParser from .module import available_modules, LazyBaseModule, get_module, setup_paths from . import validators, Feeder, Extractor, Database, Storage, Formatter, Enricher from .module import BaseModule @@ -50,7 +50,7 @@ class AuthenticationJsonParseAction(JsonParseAction): auth_dict = json.load(f) except json.JSONDecodeError: # maybe it's yaml, try that - auth_dict = yaml.load(f) + auth_dict = _yaml.load(f) except: pass From a69ac3e509eed60f1801aca605531b6bc8f3e506 Mon Sep 17 00:00:00 2001 From: erinhmclark Date: Tue, 11 Feb 2025 09:46:22 +0000 Subject: [PATCH 48/62] Fix file hash reference in S3 tests --- tests/storages/test_S3_storage.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/tests/storages/test_S3_storage.py b/tests/storages/test_S3_storage.py index 2594e73..e532a18 100644 --- a/tests/storages/test_S3_storage.py +++ b/tests/storages/test_S3_storage.py @@ -2,13 +2,12 @@ from typing import Type import pytest from unittest.mock import MagicMock, patch from auto_archiver.core import Media -from auto_archiver.modules.hash_enricher import HashEnricher from auto_archiver.modules.s3_storage import s3_storage -class TestGDriveStorage: +class TestS3Storage: """ - Test suite for GDriveStorage. + Test suite for S3Storage. """ module_name: str = "s3_storage" storage: Type[s3_storage] @@ -66,7 +65,7 @@ class TestGDriveStorage: # Set duplicate checking config to true: self.storage.random_no_duplicate = True - with patch('auto_archiver.modules.hash_enricher.HashEnricher.calculate_hash') as mock_calc_hash, \ + with patch('auto_archiver.modules.s3_storage.s3_storage.calculate_file_hash') as mock_calc_hash, \ patch.object(self.storage, 'file_in_folder') as mock_file_in_folder: mock_calc_hash.return_value = 'beepboop123beepboop123beepboop123' mock_file_in_folder.return_value = 'existing_key.txt' @@ -87,8 +86,7 @@ class TestGDriveStorage: # Create test media with calculated hash media = Media("test.txt") media.key = "original_path.txt" - - with patch('auto_archiver.modules.hash_enricher.HashEnricher.calculate_hash') as mock_calculate_hash: + with patch('auto_archiver.modules.s3_storage.s3_storage.calculate_file_hash') as mock_calculate_hash: mock_calculate_hash.return_value = "beepboop123beepboop123beepboop123" # Verify upload assert self.storage.is_upload_needed(media) is False From 18666ff027526b99114d2b4ffb6304f9b3a83461 Mon Sep 17 00:00:00 2001 From: erinhmclark Date: Tue, 11 Feb 2025 11:28:24 +0000 Subject: [PATCH 49/62] skip authenticated tests in test_gsheet_feeder.py --- tests/feeders/test_gsheet_feeder.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/feeders/test_gsheet_feeder.py b/tests/feeders/test_gsheet_feeder.py index ecf57f1..bdf3e70 100644 --- a/tests/feeders/test_gsheet_feeder.py +++ b/tests/feeders/test_gsheet_feeder.py @@ -185,7 +185,7 @@ def test_should_process_sheet(setup_module): assert gdb.should_process_sheet("AnotherSheet") == False -# @pytest.mark.skip(reason="Requires a real connection") +@pytest.mark.skip(reason="Requires a real connection") class TestGSheetsFeederReal: """Testing GSheetsFeeder class""" From 1792e02d1d32c99ca1a59aeb0cab33a74d3a783e Mon Sep 17 00:00:00 2001 From: erinhmclark Date: Tue, 11 Feb 2025 11:34:36 +0000 Subject: [PATCH 50/62] skip authenticated tests in test_gdrive_storage.py --- tests/storages/test_gdrive_storage.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/storages/test_gdrive_storage.py b/tests/storages/test_gdrive_storage.py index 4259cb2..57480d0 100644 --- a/tests/storages/test_gdrive_storage.py +++ b/tests/storages/test_gdrive_storage.py @@ -35,7 +35,7 @@ class TestGDriveStorage(TestStorageBase): media = Media(filename="test.jpg") media.key = "folder1/folder2/test.jpg" -# @pytest.mark.skip(reason="Requires real credentials") +@pytest.mark.skip(reason="Requires real credentials") @pytest.mark.download class TestGDriveStorageConnected(TestStorageBase): """ From 89d9140d15eb9e4261abf27f9c71df47ef8efb07 Mon Sep 17 00:00:00 2001 From: erinhmclark Date: Tue, 11 Feb 2025 11:47:11 +0000 Subject: [PATCH 51/62] Fixed setup/ config_setup reference --- tests/storages/test_gdrive_storage.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/tests/storages/test_gdrive_storage.py b/tests/storages/test_gdrive_storage.py index 57480d0..aba0a25 100644 --- a/tests/storages/test_gdrive_storage.py +++ b/tests/storages/test_gdrive_storage.py @@ -7,7 +7,7 @@ from auto_archiver.core.metadata import Metadata from tests.storages.test_storage_base import TestStorageBase -class TestGDriveStorage(TestStorageBase): +class TestGDriveStorage: """ Test suite for GDriveStorage. """ @@ -21,6 +21,10 @@ class TestGDriveStorage(TestStorageBase): 'service_account': 'fake_service_account.json' } + @pytest.fixture(autouse=True) + def gdrive(self, setup_module): + with patch('google.oauth2.service_account.Credentials.from_service_account_file') as mock_creds: + self.storage = setup_module(self.module_name, self.config) def test_initialize_fails_with_non_existent_creds(self): """ @@ -28,13 +32,15 @@ class TestGDriveStorage(TestStorageBase): """ # Act and Assert with pytest.raises(FileNotFoundError) as exc_info: - self.storage.setup(self.config) + self.storage.setup() assert "No such file or directory" in str(exc_info.value) + def test_path_parts(self): media = Media(filename="test.jpg") media.key = "folder1/folder2/test.jpg" + @pytest.mark.skip(reason="Requires real credentials") @pytest.mark.download class TestGDriveStorageConnected(TestStorageBase): From f97ec6a9e0ac20268f045b661f2e080ff1eb8574 Mon Sep 17 00:00:00 2001 From: erinhmclark Date: Tue, 11 Feb 2025 11:58:28 +0000 Subject: [PATCH 52/62] Fixed S3 module import --- tests/storages/test_S3_storage.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/storages/test_S3_storage.py b/tests/storages/test_S3_storage.py index e532a18..2a5d026 100644 --- a/tests/storages/test_S3_storage.py +++ b/tests/storages/test_S3_storage.py @@ -2,7 +2,7 @@ from typing import Type import pytest from unittest.mock import MagicMock, patch from auto_archiver.core import Media -from auto_archiver.modules.s3_storage import s3_storage +from auto_archiver.modules.s3_storage import S3Storage class TestS3Storage: @@ -10,7 +10,7 @@ class TestS3Storage: Test suite for S3Storage. """ module_name: str = "s3_storage" - storage: Type[s3_storage] + storage: Type[S3Storage] s3: MagicMock config: dict = { "path_generator": "flat", @@ -78,7 +78,7 @@ class TestS3Storage: ) - @patch.object(s3_storage.S3Storage, 'file_in_folder') + @patch.object(S3Storage, 'file_in_folder') def test_skips_upload_when_duplicate_exists(self, mock_file_in_folder): """Test that upload skips when file_in_folder finds existing object""" self.storage.random_no_duplicate = True @@ -97,7 +97,7 @@ class TestS3Storage: mock_upload.assert_not_called() assert result is True - @patch.object(s3_storage.S3Storage, 'is_upload_needed') + @patch.object(S3Storage, 'is_upload_needed') def test_uploads_with_correct_parameters(self, mock_upload_needed): media = Media("test.txt") media.key = "original_key.txt" From 5e2e93382ffc47893183aae83ff138055b0edeb8 Mon Sep 17 00:00:00 2001 From: erinhmclark Date: Tue, 11 Feb 2025 12:17:42 +0000 Subject: [PATCH 53/62] Test fixes for 3.10 compliance. --- tests/databases/test_gsheet_db.py | 2 +- tests/feeders/test_gsheet_feeder.py | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/databases/test_gsheet_db.py b/tests/databases/test_gsheet_db.py index 0a655a8..32e8403 100644 --- a/tests/databases/test_gsheet_db.py +++ b/tests/databases/test_gsheet_db.py @@ -24,7 +24,7 @@ def mock_metadata(): metadata.status = "done" metadata.get_title.return_value = "Example Title" metadata.get.return_value = "Example Content" - metadata.get_timestamp.return_value = "2025-01-01T00:00:00Z" + metadata.get_timestamp.return_value = "2025-01-01T00:00:00" metadata.get_final_media.return_value = MagicMock(spec=Media) metadata.get_all_media.return_value = [] metadata.get_media_by_id.return_value = None diff --git a/tests/feeders/test_gsheet_feeder.py b/tests/feeders/test_gsheet_feeder.py index bdf3e70..b86e329 100644 --- a/tests/feeders/test_gsheet_feeder.py +++ b/tests/feeders/test_gsheet_feeder.py @@ -52,7 +52,7 @@ def gsheet_feeder(setup_module) -> GsheetsFeeder: return feeder -class TestWorksheet: +class MockWorksheet: """ mimics the bits we need from gworksheet """ @@ -91,7 +91,7 @@ class TestWorksheet: def test__process_rows(gsheet_feeder: GsheetsFeeder): - testworksheet = TestWorksheet() + testworksheet = MockWorksheet() metadata_items = list(gsheet_feeder._process_rows(testworksheet)) assert len(metadata_items) == 3 assert isinstance(metadata_items[0], Metadata) @@ -99,7 +99,7 @@ def test__process_rows(gsheet_feeder: GsheetsFeeder): def test__set_metadata(gsheet_feeder: GsheetsFeeder): - worksheet = TestWorksheet() + worksheet = MockWorksheet() metadata = Metadata() gsheet_feeder._set_context(metadata, worksheet, 1) assert metadata.get_context("gsheet") == {"row": 1, "worksheet": worksheet} @@ -112,7 +112,7 @@ def test__set_metadata_with_folder_pickled(gsheet_feeder: GsheetsFeeder, workshe def test__set_metadata_with_folder(gsheet_feeder: GsheetsFeeder): - testworksheet = TestWorksheet() + testworksheet = MockWorksheet() metadata = Metadata() testworksheet.wks.title = "TestSheet" gsheet_feeder._set_context(metadata, testworksheet, 6) From d1d6cde008861f508b8689ff6fd30cdde2fccd3a Mon Sep 17 00:00:00 2001 From: erinhmclark Date: Tue, 11 Feb 2025 12:27:48 +0000 Subject: [PATCH 54/62] Set mock timestamp without z format --- tests/databases/test_gsheet_db.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/databases/test_gsheet_db.py b/tests/databases/test_gsheet_db.py index 32e8403..18a22f1 100644 --- a/tests/databases/test_gsheet_db.py +++ b/tests/databases/test_gsheet_db.py @@ -41,7 +41,7 @@ def metadata(): metadata.set_title("Example Title") metadata.set_content("Example Content") metadata.success("my-archiver") - metadata.set("timestamp", "2025-01-01T00:00:00Z") + metadata.set("timestamp", "2025-01-01T00:00:00") metadata.set("date", "2025-02-04T18:22:24.909112+00:00") return metadata From 7309cd32e7df6ebf21b32ed0cba288ba8ecea297 Mon Sep 17 00:00:00 2001 From: msramalho <19508417+msramalho@users.noreply.github.com> Date: Tue, 11 Feb 2025 12:51:17 +0000 Subject: [PATCH 55/62] fix: context to be updated on Metadata.merge --- src/auto_archiver/core/metadata.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/auto_archiver/core/metadata.py b/src/auto_archiver/core/metadata.py index d20ea5e..a8d2ad4 100644 --- a/src/auto_archiver/core/metadata.py +++ b/src/auto_archiver/core/metadata.py @@ -44,6 +44,7 @@ class Metadata: if overwrite_left: if right.status and len(right.status): self.status = right.status + self._context.update(right._context) for k, v in right.metadata.items(): assert k not in self.metadata or type(v) == type(self.get(k)) if type(v) not in [dict, list, set] or k not in self.metadata: From e6594ad3dcb1f8e95919b1ef8a632ea321f7be7a Mon Sep 17 00:00:00 2001 From: msramalho <19508417+msramalho@users.noreply.github.com> Date: Tue, 11 Feb 2025 12:52:42 +0000 Subject: [PATCH 56/62] merge result into cached results for context preservation --- src/auto_archiver/core/orchestrator.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/auto_archiver/core/orchestrator.py b/src/auto_archiver/core/orchestrator.py index 473f882..bb5f9e3 100644 --- a/src/auto_archiver/core/orchestrator.py +++ b/src/auto_archiver/core/orchestrator.py @@ -424,8 +424,8 @@ class ArchivingOrchestrator: cached_result = None for d in self.databases: d.started(result) - if (local_result := d.fetch(result)): - cached_result = (cached_result or Metadata()).merge(local_result) + if local_result := d.fetch(result): + cached_result = (cached_result or Metadata()).merge(local_result).merge(result) if cached_result: logger.debug("Found previously archived entry") for d in self.databases: From 6fdd5f0e662293731ffe435d41b1d5e93d094cec Mon Sep 17 00:00:00 2001 From: msramalho <19508417+msramalho@users.noreply.github.com> Date: Tue, 11 Feb 2025 12:53:12 +0000 Subject: [PATCH 57/62] fix cases of single : vs :: in entrypoint --- src/auto_archiver/modules/api_db/__manifest__.py | 2 +- src/auto_archiver/modules/atlos_db/__manifest__.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/auto_archiver/modules/api_db/__manifest__.py b/src/auto_archiver/modules/api_db/__manifest__.py index 698c2e4..19129a4 100644 --- a/src/auto_archiver/modules/api_db/__manifest__.py +++ b/src/auto_archiver/modules/api_db/__manifest__.py @@ -1,7 +1,7 @@ { "name": "Auto-Archiver API Database", "type": ["database"], - "entry_point": "api_db:AAApiDb", + "entry_point": "api_db::AAApiDb", "requires_setup": True, "dependencies": { "python": ["requests", "loguru"], diff --git a/src/auto_archiver/modules/atlos_db/__manifest__.py b/src/auto_archiver/modules/atlos_db/__manifest__.py index 8f9473f..b9cabf2 100644 --- a/src/auto_archiver/modules/atlos_db/__manifest__.py +++ b/src/auto_archiver/modules/atlos_db/__manifest__.py @@ -1,7 +1,7 @@ { "name": "Atlos Database", "type": ["database"], - "entry_point": "atlos_db:AtlosDb", + "entry_point": "atlos_db::AtlosDb", "requires_setup": True, "dependencies": {"python": ["loguru", From 4eeb39477c4b3cf81be680fddbaa3ce91bfad8a1 Mon Sep 17 00:00:00 2001 From: msramalho <19508417+msramalho@users.noreply.github.com> Date: Tue, 11 Feb 2025 12:53:46 +0000 Subject: [PATCH 58/62] improves gsheetdb feedback on retrieve sheet failure --- src/auto_archiver/modules/gsheet_db/gsheet_db.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/auto_archiver/modules/gsheet_db/gsheet_db.py b/src/auto_archiver/modules/gsheet_db/gsheet_db.py index 5e1ed1e..5b270bf 100644 --- a/src/auto_archiver/modules/gsheet_db/gsheet_db.py +++ b/src/auto_archiver/modules/gsheet_db/gsheet_db.py @@ -97,6 +97,6 @@ class GsheetsDb(Database): gw: GWorksheet = gsheet.get("worksheet") row: int = gsheet.get("row") elif self.sheet_id: - print(self.sheet_id) + logger.error(f"Unable to retrieve Gsheet for {item.get_url()}, GsheetDB must be used alongside GsheetFeeder.") return gw, row From 5c590292212ffb3aeec56b11b0d854ad993be8e7 Mon Sep 17 00:00:00 2001 From: msramalho <19508417+msramalho@users.noreply.github.com> Date: Tue, 11 Feb 2025 12:53:58 +0000 Subject: [PATCH 59/62] updates api_db for new API endpoint --- src/auto_archiver/modules/api_db/api_db.py | 23 +++++++++++++--------- 1 file changed, 14 insertions(+), 9 deletions(-) diff --git a/src/auto_archiver/modules/api_db/api_db.py b/src/auto_archiver/modules/api_db/api_db.py index e1f67ce..374e755 100644 --- a/src/auto_archiver/modules/api_db/api_db.py +++ b/src/auto_archiver/modules/api_db/api_db.py @@ -16,10 +16,10 @@ class AAApiDb(Database): Helps avoid re-archiving the same URL multiple times. """ if not self.allow_rearchive: return - + params = {"url": item.get_url(), "limit": 15} headers = {"Authorization": f"Bearer {self.api_token}", "accept": "application/json"} - response = requests.get(os.path.join(self.api_endpoint, "tasks/search-url"), params=params, headers=headers) + response = requests.get(os.path.join(self.api_endpoint, "url/search"), params=params, headers=headers) if response.status_code == 200: if len(response.json()): @@ -30,21 +30,26 @@ class AAApiDb(Database): logger.error(f"AA API FAIL ({response.status_code}): {response.json()}") return False - - def done(self, item: Metadata, cached: bool=False) -> None: + def done(self, item: Metadata, cached: bool = False) -> None: """archival result ready - should be saved to DB""" if not self.store_results: return - if cached: + if cached: logger.debug(f"skipping saving archive of {item.get_url()} to the AA API because it was cached") return logger.debug(f"saving archive of {item.get_url()} to the AA API.") - payload = {'result': item.to_json(), 'public': self.public, 'author_id': self.author_id, 'group_id': self.group_id, 'tags': list(self.tags)} + payload = { + 'author_id': self.author_id, + 'url': item.get_url(), + 'public': self.public, + 'group_id': self.group_id, + 'tags': list(self.tags), + 'result': item.to_json(), + } headers = {"Authorization": f"Bearer {self.api_token}"} - response = requests.post(os.path.join(self.api_endpoint, "submit-archive"), json=payload, headers=headers) + response = requests.post(os.path.join(self.api_endpoint, "interop/submit-archive"), json=payload, headers=headers) - if response.status_code == 200: + if response.status_code == 201: logger.success(f"AA API: {response.json()}") else: logger.error(f"AA API FAIL ({response.status_code}): {response.json()}") - From 977f06c37a159a9557170409f726530b0903f0e9 Mon Sep 17 00:00:00 2001 From: msramalho <19508417+msramalho@users.noreply.github.com> Date: Tue, 11 Feb 2025 12:56:33 +0000 Subject: [PATCH 60/62] renames api_db property for clarity --- src/auto_archiver/modules/api_db/__manifest__.py | 4 ++-- src/auto_archiver/modules/api_db/api_db.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/auto_archiver/modules/api_db/__manifest__.py b/src/auto_archiver/modules/api_db/__manifest__.py index 19129a4..8359174 100644 --- a/src/auto_archiver/modules/api_db/__manifest__.py +++ b/src/auto_archiver/modules/api_db/__manifest__.py @@ -23,7 +23,7 @@ "default": None, "help": "which group of users have access to the archive in case public=false as author", }, - "allow_rearchive": { + "use_api_cache": { "default": True, "type": "bool", "help": "if False then the API database will be queried prior to any archiving operations and stop if the link has already been archived", @@ -43,7 +43,7 @@ ### Features - **API Integration**: Supports querying for existing archives and submitting results. -- **Duplicate Prevention**: Avoids redundant archiving when `allow_rearchive` is disabled. +- **Duplicate Prevention**: Avoids redundant archiving when `use_api_cache` is disabled. - **Configurable**: Supports settings like API endpoint, authentication token, tags, and permissions. - **Tagging and Metadata**: Adds tags and manages metadata for archives. - **Optional Storage**: Archives results conditionally based on configuration. diff --git a/src/auto_archiver/modules/api_db/api_db.py b/src/auto_archiver/modules/api_db/api_db.py index 374e755..753ff3f 100644 --- a/src/auto_archiver/modules/api_db/api_db.py +++ b/src/auto_archiver/modules/api_db/api_db.py @@ -15,7 +15,7 @@ class AAApiDb(Database): """ query the database for the existence of this item. Helps avoid re-archiving the same URL multiple times. """ - if not self.allow_rearchive: return + if not self.use_api_cache: return params = {"url": item.get_url(), "limit": 15} headers = {"Authorization": f"Bearer {self.api_token}", "accept": "application/json"} From d90d3cec28d2424a7370d232f6445965507b5d92 Mon Sep 17 00:00:00 2001 From: msramalho <19508417+msramalho@users.noreply.github.com> Date: Tue, 11 Feb 2025 13:03:18 +0000 Subject: [PATCH 61/62] fix telethon_extractor setup --- .../modules/telethon_extractor/telethon_extractor.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/auto_archiver/modules/telethon_extractor/telethon_extractor.py b/src/auto_archiver/modules/telethon_extractor/telethon_extractor.py index 21fc4dc..8088364 100644 --- a/src/auto_archiver/modules/telethon_extractor/telethon_extractor.py +++ b/src/auto_archiver/modules/telethon_extractor/telethon_extractor.py @@ -18,12 +18,13 @@ class TelethonExtractor(Extractor): invite_pattern = re.compile(r"t.me(\/joinchat){0,1}\/\+?(.+)") - def setup(self) -> None: + def setup(self, config: dict) -> None: """ 1. makes a copy of session_file that is removed in cleanup 2. trigger login process for telegram or proceed if already saved in a session file 3. joins channel_invites where needed """ + super().setup(config) logger.info(f"SETUP {self.name} checking login...") # make a copy of the session that is used exclusively with this archiver instance From 977618b4ceeb8c02d5a561905cf37f3391c3db3e Mon Sep 17 00:00:00 2001 From: msramalho <19508417+msramalho@users.noreply.github.com> Date: Tue, 11 Feb 2025 13:04:59 +0000 Subject: [PATCH 62/62] doc: adds note about telethon vs telegram extractors --- src/auto_archiver/modules/telegram_extractor/__manifest__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/auto_archiver/modules/telegram_extractor/__manifest__.py b/src/auto_archiver/modules/telegram_extractor/__manifest__.py index e1c49c2..cb0ee1e 100644 --- a/src/auto_archiver/modules/telegram_extractor/__manifest__.py +++ b/src/auto_archiver/modules/telegram_extractor/__manifest__.py @@ -13,7 +13,7 @@ The `TelegramExtractor` retrieves publicly available media content from Telegram message links without requiring login credentials. It processes URLs to fetch images and videos embedded in Telegram messages, ensuring a structured output using `Metadata` and `Media` objects. Recommended for scenarios where login-based archiving is not viable, although `telethon_archiver` - is advised for more comprehensive functionality. + is advised for more comprehensive functionality, and higher quality media extraction. ### Features - Extracts images and videos from public Telegram message links (`t.me`).