mirror of
https://github.com/lrnselfreliance/wrolpi.git
synced 2026-10-26 02:48:38 -04:00
Refresh walked the tree twice (find, then stat every unchanged group), and a two-directory refresh scanned the whole media directory. Compare now diffs in SQL on a TEMP table that already has mtime, accepts multiple roots, and groups by a stored stem that is versioned so a suffix-rule change cannot delete-and-reinsert tagged files. Deletes refresh_discover_paths and /api/files/refresh_progress, which nothing in production called. Video/doc/zim modelers use the same (no domain row OR indexed != True) gate.
3118 lines
112 KiB
Python
3118 lines
112 KiB
Python
import asyncio
|
||
import atexit
|
||
import contextlib
|
||
import functools
|
||
import inspect
|
||
import json
|
||
import logging
|
||
import logging.config
|
||
import multiprocessing
|
||
import os
|
||
import pathlib
|
||
import re
|
||
import shutil
|
||
import socket
|
||
import string
|
||
import sys
|
||
import tempfile
|
||
from asyncio import Task
|
||
from copy import deepcopy
|
||
from dataclasses import dataclass, field, fields
|
||
from datetime import datetime, date
|
||
from decimal import Decimal
|
||
from functools import wraps
|
||
from itertools import islice, filterfalse, tee
|
||
from multiprocessing.managers import DictProxy
|
||
from pathlib import Path
|
||
from types import GeneratorType
|
||
from typing import Optional, List, AsyncGenerator
|
||
from typing import Union, Callable, Tuple, Dict, Iterable, Generator, Any, Set
|
||
from urllib.parse import urlparse, urlunsplit
|
||
|
||
import aiohttp
|
||
import bs4
|
||
import yaml
|
||
from aiohttp import ClientResponse, ClientSession
|
||
from bs4 import BeautifulSoup
|
||
from selenium import webdriver
|
||
from sqlalchemy import types
|
||
from sqlalchemy.ext.declarative import declarative_base
|
||
from sqlalchemy.orm import Session
|
||
|
||
from wrolpi.dates import now, from_timestamp
|
||
from wrolpi.errors import WROLModeEnabled, NativeOnly, LogLevelError, InvalidConfig, \
|
||
ValidationError
|
||
from wrolpi.log_levels import TRACE_LEVEL
|
||
from wrolpi.vars import PYTEST, DOCKERIZED, CONFIG_DIR, MEDIA_DIRECTORY, DEFAULT_HTTP_HEADERS
|
||
|
||
|
||
def add_logging_level(level_name: str, level_int: int, methodName=None):
|
||
"""
|
||
Comprehensively adds a new logging level to the `logging` module and the
|
||
currently configured logging class.
|
||
|
||
`levelName` becomes an attribute of the `logging` module with the value
|
||
`levelNum`. `methodName` becomes a convenience method for both `logging`
|
||
itself and the class returned by `logging.getLoggerClass()` (usually just
|
||
`logging.Logger`). If `methodName` is not specified, `levelName.lower()` is
|
||
used.
|
||
|
||
To avoid accidental clobberings of existing attributes, this method will
|
||
raise an `AttributeError` if the level name is already an attribute of the
|
||
`logging` module or if the method name is already present
|
||
|
||
Example
|
||
-------
|
||
>>> add_logging_level('TRACE', logging.DEBUG - 5)
|
||
>>> logging.getLogger(__name__).setLevel("TRACE")
|
||
>>> logging.getLogger(__name__).trace('that worked')
|
||
>>> logging.trace('so did this')
|
||
>>> logging.TRACE
|
||
5
|
||
|
||
This function is taken directly from: https://stackoverflow.com/questions/2183233/how-to-add-a-custom-loglevel-to-pythons-logging-facility/35804945#35804945
|
||
"""
|
||
if not methodName:
|
||
methodName = level_name.lower()
|
||
|
||
if hasattr(logging, level_name):
|
||
raise AttributeError('{} already defined in logging module'.format(level_name))
|
||
if hasattr(logging, methodName):
|
||
raise AttributeError('{} already defined in logging module'.format(methodName))
|
||
if hasattr(logging.getLoggerClass(), methodName):
|
||
raise AttributeError('{} already defined in logger class'.format(methodName))
|
||
|
||
# This method was inspired by the answers to Stack Overflow post
|
||
# http://stackoverflow.com/q/2183233/2988730, especially
|
||
# http://stackoverflow.com/a/13638084/2988730
|
||
def logForLevel(self, message, *args, **kwargs):
|
||
if self.isEnabledFor(level_int):
|
||
kwargs.setdefault('stacklevel', 3)
|
||
self._log(level_int, message, args, **kwargs)
|
||
|
||
def logToRoot(message, *args, **kwargs):
|
||
kwargs.setdefault('stacklevel', 3)
|
||
logging.log(level_int, message, *args, **kwargs)
|
||
|
||
logging.addLevelName(level_int, level_name)
|
||
setattr(logging, level_name, level_int)
|
||
setattr(logging.getLoggerClass(), methodName, logForLevel)
|
||
setattr(logging, methodName, logToRoot)
|
||
|
||
|
||
def is_journal_stream() -> bool:
|
||
"""Check if stderr is connected to systemd journal.
|
||
|
||
Systemd sets JOURNAL_STREAM to device:inode of the journal socket when a service's
|
||
stdout/stderr is connected to the journal. We verify our stderr matches to avoid
|
||
false positives from subprocesses with redirected output.
|
||
"""
|
||
journal_stream = os.environ.get('JOURNAL_STREAM')
|
||
if not journal_stream:
|
||
return False
|
||
try:
|
||
device, inode = map(int, journal_stream.split(':'))
|
||
stat = os.fstat(sys.stderr.fileno())
|
||
return stat.st_dev == device and stat.st_ino == inode
|
||
except (ValueError, OSError):
|
||
return False
|
||
|
||
|
||
# Conditionally include timestamp/PID in log format.
|
||
# When running under systemd journal, these are provided by journald.
|
||
# When running in Docker or locally, we need to include them ourselves.
|
||
_LOG_PREFIX = '' if is_journal_stream() else '[%(asctime)s] [%(process)d] '
|
||
|
||
LOGGING_CONFIG = {
|
||
'version': 1,
|
||
'disable_existing_loggers': False,
|
||
'loggers': {
|
||
'sanic.root': {
|
||
'level': 'INFO',
|
||
'handlers': ['console'],
|
||
},
|
||
'sanic.error': {
|
||
'level': 'INFO',
|
||
'handlers': ['error_console'],
|
||
'qualname': 'sanic.error',
|
||
},
|
||
'sanic.access': {
|
||
'level': 'INFO',
|
||
'handlers': ['access_console'],
|
||
'qualname': 'sanic.access',
|
||
},
|
||
'sanic.server': {
|
||
'level': 'INFO',
|
||
'handlers': ['console'],
|
||
'qualname': 'sanic.server',
|
||
},
|
||
},
|
||
'handlers': {
|
||
'console': {
|
||
'class': 'logging.StreamHandler',
|
||
'formatter': 'generic',
|
||
'stream': sys.stdout,
|
||
'filters': ['empty_message_filter'],
|
||
},
|
||
'error_console': {
|
||
'class': 'logging.StreamHandler',
|
||
'formatter': 'generic',
|
||
'stream': sys.stderr,
|
||
},
|
||
'access_console': {
|
||
'class': 'logging.StreamHandler',
|
||
'formatter': 'access',
|
||
'stream': sys.stdout,
|
||
'filters': ['empty_message_filter'],
|
||
},
|
||
},
|
||
'formatters': {
|
||
'generic': {
|
||
'format': f'{_LOG_PREFIX}[%(name)s:%(lineno)d] [%(levelname)s] %(message)s',
|
||
'class': 'logging.Formatter',
|
||
},
|
||
'access': {
|
||
'format': f'{_LOG_PREFIX}[%(name)s:%(lineno)d] [%(levelname)s]: %(request)s %(message)s %(status)s %(byte)s',
|
||
'class': 'logging.Formatter',
|
||
},
|
||
},
|
||
'root': {
|
||
'handlers': ['console'],
|
||
'level': logging.DEBUG
|
||
},
|
||
'filters': {
|
||
'empty_message_filter': {
|
||
'()': 'wrolpi.logging_extra.EmptyMessageFilter',
|
||
}
|
||
}
|
||
}
|
||
|
||
|
||
# Used so IDE will suggest `logger.trace`
|
||
class TraceLogger(logging.Logger):
|
||
def trace(self, msg, *args, **kwargs):
|
||
return self.log(TRACE_LEVEL, msg, *args, **kwargs)
|
||
|
||
|
||
# Apply logging config.
|
||
logger: TraceLogger = logging.getLogger() # noqa Not really a TraceLogger, but I wanted IDE suggestions.
|
||
logging.config.dictConfig(LOGGING_CONFIG)
|
||
|
||
# Add extra verbose debug level (-vvv). TRACE_LEVEL imported from log_levels.
|
||
if not hasattr(logging, 'TRACE'):
|
||
add_logging_level('TRACE', TRACE_LEVEL)
|
||
|
||
logger_ = logger.getChild(__name__)
|
||
|
||
|
||
def set_log_level(level: int, warn_level: bool = True):
|
||
"""Set the level of the root logger so all children that have been created (or will be created) share the same
|
||
level.
|
||
|
||
@warning: Child processes will not share this level. See set_global_log_level
|
||
"""
|
||
logger.setLevel(level)
|
||
logger_.setLevel(level)
|
||
|
||
# Always warn about the log level, so we know what should have been logged.
|
||
effective_level = logger.getEffectiveLevel()
|
||
level_name = logging.getLevelName(effective_level)
|
||
if warn_level:
|
||
logger.warning(f'Logging level: {level_name}')
|
||
|
||
# Change log level for all handlers.
|
||
for handler in logger.handlers:
|
||
handler.setLevel(level)
|
||
|
||
# Enable debug logging in SQLAlchemy when logging is TRACE.
|
||
sa_logger = logging.getLogger('sqlalchemy.engine')
|
||
sa_level = logging.DEBUG if level == TRACE_LEVEL else logging.WARNING
|
||
sa_logger.setLevel(logging.NOTSET)
|
||
|
||
# Hide sanic access logs when not at least "INFO".
|
||
sanic_access = logging.getLogger('sanic.access')
|
||
sanic_level = logging.INFO if level <= 20 else logging.WARNING
|
||
sanic_access.setLevel(sanic_level)
|
||
|
||
|
||
def set_global_log_level(log_level: int):
|
||
"""Set the global (shared between processes) log level."""
|
||
if not isinstance(log_level, int) or not 0 <= log_level <= 40:
|
||
raise LogLevelError()
|
||
from wrolpi.api_utils import api_app
|
||
with api_app.shared_ctx.log_level.get_lock():
|
||
api_app.shared_ctx.log_level.value = log_level
|
||
|
||
|
||
@contextlib.contextmanager
|
||
def log_level_context(level):
|
||
starting_level = logger.getEffectiveLevel()
|
||
set_log_level(level, warn_level=False)
|
||
yield
|
||
set_log_level(starting_level, warn_level=False)
|
||
|
||
|
||
__all__ = [
|
||
'Base',
|
||
'ConfigFile',
|
||
'DownloadFileInfo',
|
||
'LOGGING_CONFIG',
|
||
'ModelHelper',
|
||
'TRACE_LEVEL',
|
||
'WROLPI_CONFIG',
|
||
'aiohttp_get',
|
||
'aiohttp_head',
|
||
'aiohttp_post',
|
||
'api_param_limiter',
|
||
'apply_modelers',
|
||
'apply_refresh_cleanup',
|
||
'background_task',
|
||
'can_connect_to_server',
|
||
'cancel_background_tasks',
|
||
'cancel_refresh_tasks',
|
||
'cancelable_wrapper',
|
||
'chain',
|
||
'chdir',
|
||
'check_media_directory',
|
||
'chunks',
|
||
'chunks_by_stem',
|
||
'create_empty_config_files',
|
||
'cum_timer',
|
||
'date_range',
|
||
'disable_wrol_mode',
|
||
'enable_wrol_mode',
|
||
'escape_file_name',
|
||
'extract_domain',
|
||
'extract_headlines',
|
||
'extract_html_text',
|
||
'format_html_string',
|
||
'format_json_file',
|
||
'get_absolute_media_path',
|
||
'get_all_configs',
|
||
'get_download_info',
|
||
'get_files_and_directories',
|
||
'get_global_statistics',
|
||
'get_html_soup',
|
||
'get_media_directory',
|
||
'get_relative_to_media_directory',
|
||
'get_title_from_html',
|
||
'get_warn_once',
|
||
'get_wrolpi_config',
|
||
'html_file_screenshot',
|
||
'html_screenshot',
|
||
'insert_parameter',
|
||
'iterify',
|
||
'limit_concurrent',
|
||
'logger',
|
||
'make_media_directory',
|
||
'minimize_dict',
|
||
'native_only',
|
||
'partition',
|
||
'recursive_map',
|
||
'register_modeler',
|
||
'register_refresh_cleanup',
|
||
'remove_whitespace',
|
||
'replace_file',
|
||
'resolve_generators',
|
||
'run_after',
|
||
'set_global_log_level',
|
||
'set_log_level',
|
||
'set_system_timezone',
|
||
'set_test_config',
|
||
'set_test_media_directory',
|
||
'slow_logger',
|
||
'split_lines_by_length',
|
||
'timer',
|
||
'trim_file_name',
|
||
'truncate_generator_bytes',
|
||
'truncate_object_bytes',
|
||
'unique_by_predicate',
|
||
'url_strip_host',
|
||
'walk',
|
||
'wrol_mode_check',
|
||
'wrol_mode_enabled',
|
||
'zig_zag',
|
||
]
|
||
|
||
# Base is used for all SQLAlchemy models.
|
||
Base = declarative_base()
|
||
|
||
|
||
class ModelHelper:
|
||
|
||
def dict(self, *_, **__) -> dict:
|
||
d = {i.name: getattr(self, i.name) for i in self.__table__.columns} # noqa
|
||
return d
|
||
|
||
@classmethod
|
||
def upsert(cls, session: Session, file) -> Base:
|
||
raise NotImplementedError('This model has not defined this method.')
|
||
|
||
@staticmethod
|
||
def get_by_path(session: Session, path) -> Base:
|
||
raise NotImplementedError('This model has not defined this method.')
|
||
|
||
@staticmethod
|
||
def get_by_id(session: Session, id_: int) -> Optional[Base]:
|
||
"""Attempts to get a model instance by its id. Returns None if no instance can be found."""
|
||
raise NotImplementedError('This model has not defined this method.')
|
||
|
||
@staticmethod
|
||
def find_by_id(session: Session, id_: int) -> Base:
|
||
"""Get a model instance by its id, raise an error if no instance is found."""
|
||
raise NotImplementedError('This model has not defined this method.')
|
||
|
||
@staticmethod
|
||
def can_model(file_group) -> bool:
|
||
raise NotImplementedError('This model has not defined this method.')
|
||
|
||
@staticmethod
|
||
def do_model(session: Session, file_group):
|
||
raise NotImplementedError('This model has not defined this method.')
|
||
|
||
def flush(self, session: Session = None):
|
||
"""A convenience method which flushes this record using its DB Session."""
|
||
session = session or Session.object_session(self)
|
||
if session:
|
||
session.flush([self])
|
||
else:
|
||
raise RuntimeError(f'{self} is not in a session!')
|
||
|
||
|
||
URL_CHARS = string.ascii_lowercase + string.digits
|
||
|
||
|
||
def find_file(directory: pathlib.Path, name: str, depth=1) -> Optional[pathlib.Path]:
|
||
"""Recursively searches a directory for a
|
||
file with the provided name."""
|
||
if depth == 0:
|
||
return
|
||
|
||
try:
|
||
entries = sorted(directory.iterdir())
|
||
except PermissionError:
|
||
return
|
||
|
||
for path in entries:
|
||
if path.is_file() and path.name == name:
|
||
return path
|
||
if path.is_dir() and (result := find_file(path, name, depth - 1)):
|
||
return result
|
||
|
||
|
||
TEST_MEDIA_DIRECTORY = None
|
||
|
||
|
||
def set_test_media_directory(path):
|
||
global TEST_MEDIA_DIRECTORY
|
||
if isinstance(path, pathlib.Path):
|
||
TEST_MEDIA_DIRECTORY = path
|
||
elif path:
|
||
TEST_MEDIA_DIRECTORY = pathlib.Path(path)
|
||
else:
|
||
TEST_MEDIA_DIRECTORY = None
|
||
|
||
|
||
def is_tempfile(path: pathlib.Path | str) -> bool:
|
||
path = str(path)
|
||
if path.startswith('/tmp'):
|
||
return True
|
||
elif path.startswith('/var/folders/'):
|
||
return True
|
||
elif path.startswith('/private/var/folders/'):
|
||
return True
|
||
return False
|
||
|
||
|
||
def get_media_directory() -> Path:
|
||
"""Get the media directory.
|
||
|
||
This will typically be /media/wrolpi, or something else from the .env.
|
||
|
||
During testing, this function returns TEST_MEDIA_DIRECTORY.
|
||
"""
|
||
global TEST_MEDIA_DIRECTORY
|
||
|
||
if PYTEST and not TEST_MEDIA_DIRECTORY:
|
||
raise RuntimeError('No test media directory set during testing!!')
|
||
|
||
if isinstance(TEST_MEDIA_DIRECTORY, pathlib.Path):
|
||
if not is_tempfile(TEST_MEDIA_DIRECTORY):
|
||
raise RuntimeError(f'Refusing to run test outside tmp directory: {TEST_MEDIA_DIRECTORY}')
|
||
return TEST_MEDIA_DIRECTORY.resolve()
|
||
|
||
return MEDIA_DIRECTORY
|
||
|
||
|
||
DB_CONFIG_FILE_NAMES = {'tags.yaml', 'channels.yaml', 'domains.yaml', 'download_manager.yaml',
|
||
'playlists.yaml'}
|
||
|
||
|
||
class ConfigFile:
|
||
"""
|
||
This is used to keep the DB and config files in sync. Importing the config reads the config from file, then
|
||
applies it to this config instance, AND the database. Dumping the config copies data from the database and saves it
|
||
to the config file. Updating this config is automatically saved to the config file.
|
||
|
||
Configs cannot be saved until they are imported without error.
|
||
"""
|
||
background_dump: callable = None
|
||
background_save: callable = None
|
||
default_config: dict = None
|
||
file_name: str = None
|
||
validator: dataclass = None
|
||
width: int = None
|
||
|
||
def __init__(self):
|
||
self.width = self.width or 90
|
||
|
||
# Do not load a global config on import while testing. A global instance should never be used for testing.
|
||
self._config = deepcopy({k: v for k, v in self.default_config.items()})
|
||
|
||
if self._config['version'] != 0:
|
||
raise RuntimeError('Configs should start at version 0')
|
||
|
||
from wrolpi.switches import register_switch_handler
|
||
|
||
@register_switch_handler(f'background_save_{self.file_name}')
|
||
def background_save(overwrite: bool = False):
|
||
self.save(overwrite=overwrite)
|
||
|
||
self.background_save = background_save
|
||
|
||
@register_switch_handler(f'background_dump_{self.file_name}')
|
||
def background_dump(file: pathlib.Path = None):
|
||
self.dump_config(file)
|
||
|
||
self.background_dump = background_dump
|
||
|
||
def __repr__(self):
|
||
return f'<{self.__class__.__name__} file={self.get_file()}>'
|
||
|
||
def config_status(self) -> dict:
|
||
has_backup_import = self.file_name in DB_CONFIG_FILE_NAMES
|
||
d = dict(
|
||
file_name=self.file_name,
|
||
rel_path=get_relative_to_media_directory(self.get_file()),
|
||
successful_import=self.successful_import,
|
||
valid=self.is_valid(),
|
||
has_backup_import=has_backup_import,
|
||
)
|
||
return d
|
||
|
||
def __json__(self) -> dict:
|
||
return dict(self._config)
|
||
|
||
def initialize(self, multiprocessing_dict: Optional[DictProxy] = None):
|
||
"""Initializes this config dict using the default config and the config file."""
|
||
# Use the provided multiprocessing.Manager().dict() when live, or, dict() for testing.
|
||
self._config = multiprocessing_dict if multiprocessing_dict is not None else dict()
|
||
# Use the default settings to initialize the config.
|
||
self._config.update(deepcopy(self.default_config))
|
||
|
||
file = self.get_file()
|
||
if file.is_file():
|
||
if not self.is_valid(file):
|
||
raise InvalidConfig(f'Config file is invalid: {str(self.get_relative_file())}')
|
||
# is_valid() already validated, just load the data
|
||
config_data = self.read_config_file(file)
|
||
config_data = {k: v for k, v in config_data.items() if k in self.default_config}
|
||
self._config.update(config_data)
|
||
return self
|
||
|
||
@property
|
||
def successful_import(self) -> bool:
|
||
from wrolpi.api_utils import api_app
|
||
return bool(api_app.shared_ctx.configs_imported.get(self.file_name))
|
||
|
||
@successful_import.setter
|
||
def successful_import(self, value: bool):
|
||
from wrolpi.api_utils import api_app
|
||
api_app.shared_ctx.configs_imported[self.file_name] = value
|
||
|
||
@property
|
||
def import_skipped(self) -> int:
|
||
"""Count of config entries that could not be imported (e.g. their files are not yet indexed).
|
||
|
||
While non-zero, `save()` refuses to overwrite the config file so the skipped entries are not
|
||
lost when a dump regenerates the config from the (incomplete) database."""
|
||
from wrolpi.api_utils import api_app
|
||
return int(api_app.shared_ctx.configs_import_skipped.get(self.file_name) or 0)
|
||
|
||
@import_skipped.setter
|
||
def import_skipped(self, value: int):
|
||
from wrolpi.api_utils import api_app
|
||
api_app.shared_ctx.configs_import_skipped[self.file_name] = int(value)
|
||
|
||
def read_config_file(self, file: pathlib.Path = None) -> dict:
|
||
file = file or self.get_file()
|
||
with file.open('rt') as fh:
|
||
try:
|
||
config_data = yaml.load(fh, Loader=yaml.CSafeLoader)
|
||
except yaml.constructor.ConstructorError:
|
||
logger_.warning(f'Config file contains Python objects, using slow loader: {file}')
|
||
fh.seek(0)
|
||
config_data = yaml.load(fh, Loader=yaml.Loader)
|
||
if not isinstance(config_data, dict):
|
||
raise InvalidConfig(f'Config file is invalid: {file}')
|
||
return config_data
|
||
|
||
def _get_backup_filename(self):
|
||
"""Returns the path for the backup file for today."""
|
||
# TODO what if the RPi clock is not working? Change this so the "version" of the config is used each day.
|
||
path = get_media_directory() / f'config/backup/{self.file_name}'
|
||
date_str = now().strftime('%Y%m%d')
|
||
name = f'{path.stem}-{date_str}{path.suffix}'
|
||
path = path.with_name(name)
|
||
return path
|
||
|
||
def get_backup_dates(self) -> list:
|
||
"""Scan config/backup/ for backup files matching this config, return sorted dates (newest first)."""
|
||
backup_dir = get_media_directory() / 'config/backup'
|
||
if not backup_dir.is_dir():
|
||
return []
|
||
stem = pathlib.Path(self.file_name).stem
|
||
suffix = pathlib.Path(self.file_name).suffix
|
||
dates = []
|
||
for f in backup_dir.iterdir():
|
||
if f.is_file() and f.name.startswith(f'{stem}-') and f.name.endswith(suffix):
|
||
# Extract date string between stem- and suffix
|
||
date_str = f.name[len(stem) + 1:-len(suffix)]
|
||
if len(date_str) == 8 and date_str.isdigit():
|
||
dates.append(date_str)
|
||
dates.sort(reverse=True)
|
||
return dates
|
||
|
||
def _get_backup_file(self, backup_date: str) -> pathlib.Path:
|
||
"""Return path to a specific backup file."""
|
||
stem = pathlib.Path(self.file_name).stem
|
||
suffix = pathlib.Path(self.file_name).suffix
|
||
return get_media_directory() / f'config/backup/{stem}-{backup_date}{suffix}'
|
||
|
||
def preview_backup_import(self, backup_date: str, mode: str) -> dict:
|
||
"""Preview what a backup import would do. Subclasses override this."""
|
||
raise NotImplementedError(f'{self.__class__.__name__} does not support backup import')
|
||
|
||
def _preserve_current_config(self):
|
||
"""Back up the current config file to today's backup slot before a restore overwrites it.
|
||
|
||
This ensures the user can undo a backup restore by restoring today's backup.
|
||
If today's backup already exists (from an earlier save), it is not overwritten.
|
||
"""
|
||
config_file = self.get_file()
|
||
if not config_file.is_file():
|
||
return
|
||
todays_backup = self._get_backup_filename()
|
||
if todays_backup.is_file():
|
||
return
|
||
todays_backup.parent.mkdir(parents=True, exist_ok=True)
|
||
shutil.copy2(config_file, todays_backup)
|
||
|
||
def import_backup(self, backup_date: str, mode: str, send_events: bool = False):
|
||
"""Import a backup file. Subclasses override this."""
|
||
raise NotImplementedError(f'{self.__class__.__name__} does not support backup import')
|
||
|
||
def save(self, file: pathlib.Path = None, send_events: bool = False, overwrite: bool = False):
|
||
"""
|
||
Write this config to its file.
|
||
|
||
@param file: The destination of the config file (defaults to `self.get_file()`).
|
||
@param send_events: Send failure Events to UI.
|
||
@param overwrite: Will overwrite the config file even if it was not imported successfully.
|
||
"""
|
||
from wrolpi.api_utils import api_app
|
||
from wrolpi.events import Events
|
||
|
||
file = file or self.get_file()
|
||
# Compare basenames so a nested file_name (e.g. 'inventory/catalog.yaml') is accepted.
|
||
if pathlib.Path(self.file_name).name != file.name:
|
||
raise RuntimeError(f'Refusing to save config to file which does not match {self.__class__.__name__}')
|
||
|
||
rel_path = get_relative_to_media_directory(file)
|
||
|
||
if file.exists() and overwrite is False:
|
||
if not self.successful_import:
|
||
Events.send_config_save_failed(f'Failed to save config: {rel_path}')
|
||
raise RuntimeError(f'Refusing to save config because it was never successfully imported! {rel_path}')
|
||
if self.import_skipped:
|
||
message = f'Refusing to save config because {self.import_skipped} entries were skipped' \
|
||
f' during import: {rel_path}'
|
||
Events.send_config_save_failed(message)
|
||
raise RuntimeError(message)
|
||
|
||
# Don't overwrite a real config while testing.
|
||
if PYTEST and not is_tempfile(file):
|
||
raise ValueError(f'Refusing to save config file while testing: {rel_path}')
|
||
|
||
logger_.debug(f'Save config called for {rel_path}')
|
||
|
||
# Only one process can write to a config.
|
||
with api_app.shared_ctx.config_save_lock:
|
||
# The version check reads the file, so it must hold the lock.
|
||
if file.exists() and overwrite is False:
|
||
version = self.read_config_file(file).get('version')
|
||
if version and version > self.version:
|
||
message = f'Refusing to overwrite newer config ({rel_path}): {version} > {self.version}'
|
||
Events.send_config_save_failed(message)
|
||
raise RuntimeError(message)
|
||
try:
|
||
# Config directory may not exist (parents=True for nested file_names like inventory/catalog.yaml).
|
||
if not file.parent.is_dir():
|
||
file.parent.mkdir(parents=True, exist_ok=True)
|
||
|
||
# Backup the existing config file.
|
||
if file.is_file():
|
||
backup_file = self._get_backup_filename()
|
||
if not backup_file.parent.exists():
|
||
backup_file.parent.mkdir(parents=True)
|
||
shutil.copy(file, backup_file)
|
||
logger_.debug(f'Copied backup config: {rel_path} -> {backup_file}')
|
||
|
||
# Write the config in-memory to the file. Track version changes to avoid overriding newer config.
|
||
self._config['version'] = (self._config['version'] or 0) + 1
|
||
config = deepcopy(self._config)
|
||
self.write_config_data(config, file)
|
||
|
||
# Set successful_import in case this was the first time the config was written.
|
||
self.successful_import = True
|
||
|
||
logger_.info(f'Saved config: {rel_path}')
|
||
except Exception as e:
|
||
# Configs are vital, raise a big error when this fails.
|
||
message = f'Failed to save config: {rel_path}'
|
||
logger_.critical(message, exc_info=e)
|
||
if send_events:
|
||
Events.send_config_save_failed(message)
|
||
raise e
|
||
|
||
def write_config_data(self, config: dict, config_file: pathlib.Path):
|
||
write_config_data(config, config_file, width=self.width)
|
||
|
||
def get_file(self) -> Path:
|
||
if not self.file_name:
|
||
raise NotImplementedError(f'You must define a file name for this {self.__class__.__name__} config.')
|
||
|
||
if PYTEST:
|
||
return get_media_directory() / f'config/{self.file_name}'
|
||
|
||
return CONFIG_DIR / self.file_name
|
||
|
||
def get_relative_file(self):
|
||
return get_relative_to_media_directory(self.get_file())
|
||
|
||
def import_config(self, file: pathlib.Path = None, send_events=False):
|
||
"""Read config file data, apply it to the in-memory config and database."""
|
||
file = file or self.get_file()
|
||
file_str = str(get_relative_to_media_directory(file))
|
||
# Caller will set to successful if it works.
|
||
self.successful_import = False
|
||
# Subclasses count entries they cannot import; a fresh import starts clean.
|
||
self.import_skipped = 0
|
||
if file.is_file():
|
||
if not self.is_valid(file):
|
||
raise InvalidConfig(f'Config is invalid: {file_str}')
|
||
|
||
data = self.read_config_file(file)
|
||
new_data, extra_data = partition(lambda i: i[0] in self.default_config, data.items())
|
||
new_data, extra_data = dict(new_data), dict(extra_data)
|
||
if extra_data:
|
||
logger_.warning(f'Ignoring extra config data ({file_str}): {extra_data}')
|
||
self._config.update(new_data)
|
||
# Import call above this will set successful_import.
|
||
# self.successful_import = True
|
||
else:
|
||
logger_.error(f'Failed to import {file_str} because it does not exist.')
|
||
|
||
def dump_config(self, file: pathlib.Path = None, send_events=False, overwrite=False):
|
||
"""Dump database data, copy it to this the in-memory config, then write it to the config file."""
|
||
# Copy the data as-is by default. Other classes will overwrite this and dump the database before saving.
|
||
self.save(file, send_events, overwrite)
|
||
|
||
def update(self, config: dict, overwrite: bool = False):
|
||
"""Update any values of this config. Save the config to its file."""
|
||
from wrolpi.api_utils import api_app
|
||
if not self.validate(config):
|
||
raise ValidationError(f'Invalid config: {config}')
|
||
|
||
with api_app.shared_ctx.config_update_lock:
|
||
config = {k: v for k, v in config.items() if k in self._config}
|
||
self._config.update(config)
|
||
self.background_save.activate_switch(context={'overwrite': overwrite})
|
||
|
||
def dict(self) -> dict:
|
||
"""Get a deepcopy of this config."""
|
||
if not hasattr(self, '_config'):
|
||
raise NotImplementedError('You cannot use a global config while testing!')
|
||
|
||
return deepcopy(self._config)
|
||
|
||
def validate(self, config: dict) -> bool:
|
||
allowed_fields = {i.name for i in fields(self.validator)}
|
||
default_fields = set(self.default_config.keys())
|
||
# Check for fields in default_config that are missing from validator
|
||
missing_from_validator = default_fields - allowed_fields
|
||
if missing_from_validator:
|
||
logger.error(
|
||
f'{self.file_name}: Fields in default_config but missing from validator: {missing_from_validator}')
|
||
try:
|
||
# Remove keys no longer in the config.
|
||
config_items = config.items()
|
||
extra_items = {k: v for k, v in config_items if k not in allowed_fields}
|
||
if extra_items:
|
||
logger.error(f'Invalid config items: {extra_items}')
|
||
config = {k: v for k, v in config_items if k in allowed_fields}
|
||
self.validator(**config)
|
||
return True
|
||
except Exception as e:
|
||
logger.debug(f'Failed to validate config', exc_info=e)
|
||
return False
|
||
|
||
def is_valid(self, file: pathlib.Path = None) -> bool | None:
|
||
if not self.validator:
|
||
raise NotImplementedError(f'Cannot validate {self.file_name} without validator!')
|
||
|
||
file = file or self.get_file()
|
||
if not file.is_file() or file.stat().st_size == 0:
|
||
return None
|
||
|
||
try:
|
||
config_data = self.read_config_file(file)
|
||
except InvalidConfig:
|
||
return False
|
||
|
||
return self.validate(config_data)
|
||
|
||
# `version` is used to prevent overwriting of newer configs. Cannot not be modified directly.
|
||
|
||
@property
|
||
def version(self) -> int:
|
||
return self._config['version']
|
||
|
||
|
||
def write_config_data(config: dict, config_file: pathlib.Path, width: int = 90):
|
||
"""Write a config dict to a YAML file atomically (readers never see a truncated config).
|
||
Decimals are converted to str; all other Python objects are rejected.
|
||
|
||
Shared by `ConfigFile` and `MultiFileConfig` so both serialize configs identically (Decimal-safe, fsync'd)."""
|
||
|
||
class _ConfigDumper(yaml.CSafeDumper):
|
||
"""CSafeDumper that converts Decimal to str. All other Python objects are rejected."""
|
||
pass
|
||
|
||
_ConfigDumper.add_representer(Decimal, lambda dumper, data: dumper.represent_str(str(data)))
|
||
|
||
# Same directory as the config so the rename is atomic.
|
||
fd, temporary_path = tempfile.mkstemp(dir=config_file.parent, prefix=f'{config_file.name}.', suffix='.tmp')
|
||
try:
|
||
# The rename keeps the temporary file's 0600, so give it the config's mode.
|
||
mode = config_file.stat().st_mode & 0o777 if config_file.exists() else 0o644
|
||
os.chmod(temporary_path, mode)
|
||
with os.fdopen(fd, 'wt') as fh:
|
||
fd = None # The file object owns the descriptor now.
|
||
try:
|
||
yaml.dump(config, fh, Dumper=_ConfigDumper, width=width, sort_keys=True)
|
||
except yaml.representer.RepresenterError as e:
|
||
raise ValueError(f'Config contains a Python object that cannot be serialized: {e}') from e
|
||
fh.flush()
|
||
os.fsync(fh.fileno())
|
||
os.replace(temporary_path, config_file)
|
||
except BaseException:
|
||
if fd is not None:
|
||
os.close(fd)
|
||
pathlib.Path(temporary_path).unlink(missing_ok=True)
|
||
raise
|
||
|
||
|
||
def read_config_data(file: pathlib.Path) -> dict:
|
||
"""Read a YAML config file into a dict. Mirrors `ConfigFile.read_config_file`, but is standalone so
|
||
`MultiFileConfig` can use it without a `ConfigFile` instance."""
|
||
with file.open('rt') as fh:
|
||
try:
|
||
config_data = yaml.load(fh, Loader=yaml.CSafeLoader)
|
||
except yaml.constructor.ConstructorError:
|
||
logger_.warning(f'Config file contains Python objects, using slow loader: {file}')
|
||
fh.seek(0)
|
||
config_data = yaml.load(fh, Loader=yaml.Loader)
|
||
if not isinstance(config_data, dict):
|
||
raise InvalidConfig(f'Config file is invalid: {file}')
|
||
return config_data
|
||
|
||
|
||
class MultiFileConfig:
|
||
"""Manages a *directory* of YAML config files, one file per entity (e.g. ``config/inventory/<slug>.yaml``).
|
||
|
||
Unlike `ConfigFile` (one file, mirrored to the database), a `MultiFileConfig` is the source-of-truth itself:
|
||
the YAML files are loaded into a single shared multiprocessing dict (keyed by slug) and served from memory.
|
||
There is no database.
|
||
|
||
Design notes / why this is NOT a `ConfigFile` subclass:
|
||
* `ConfigFile` assumes a single, statically-named file and registers per-`file_name` switch handlers at import.
|
||
A directory has a dynamic number of files whose names are not known at import time.
|
||
* `register_switch_handler` raises on duplicate names and `SWITCH_HANDLERS` is per-process, so dynamic
|
||
per-slug switch handlers cannot be created safely. Instead this class registers ONE directory-level
|
||
``background_save_<name>`` switch that dumps every entity to disk (cheap — entities are small, human-paced).
|
||
* In-memory state lives in one shared `manager.dict()` keyed by slug, so it is shared across Sanic workers
|
||
natively (top-level key reassignment propagates through the proxy).
|
||
|
||
Subclasses define: ``subdirectory`` (e.g. 'inventory'), ``name`` (switch/registry name), ``validator``
|
||
(a dataclass) and ``default_entity`` (default dict for a new entity, must include ``version``).
|
||
"""
|
||
subdirectory: str = None
|
||
name: str = None
|
||
validator: dataclass = None
|
||
default_entity: dict = None
|
||
glob_pattern: str = '*.yaml'
|
||
width: int = 90
|
||
|
||
def __init__(self):
|
||
# slug -> entity dict (the in-memory store; replaced by a shared dict in `initialize`).
|
||
self._configs: dict = dict()
|
||
|
||
from wrolpi.switches import register_switch_handler
|
||
|
||
@register_switch_handler(f'background_save_{self.name}')
|
||
def background_save(**context):
|
||
self.dump_all()
|
||
|
||
self.background_save = background_save
|
||
|
||
def __repr__(self):
|
||
return f'<{self.__class__.__name__} dir={self.get_directory()} count={len(self._configs)}>'
|
||
|
||
def get_directory(self) -> Path:
|
||
if not self.subdirectory:
|
||
raise NotImplementedError(f'You must define a subdirectory for this {self.__class__.__name__} config.')
|
||
if PYTEST:
|
||
return get_media_directory() / 'config' / self.subdirectory
|
||
return CONFIG_DIR / self.subdirectory
|
||
|
||
def get_file(self, slug: str) -> Path:
|
||
return self.get_directory() / f'{slug}.yaml'
|
||
|
||
def discover(self) -> List[str]:
|
||
"""Return the slugs of all entity files in the directory."""
|
||
directory = self.get_directory()
|
||
if not directory.is_dir():
|
||
return []
|
||
return sorted(f.stem for f in directory.glob(self.glob_pattern) if f.is_file())
|
||
|
||
def initialize(self, multiprocessing_dict: Optional[DictProxy] = None):
|
||
"""Assign the shared in-memory store, then import every entity file from disk."""
|
||
self._configs = multiprocessing_dict if multiprocessing_dict is not None else dict()
|
||
self._configs.clear()
|
||
self.import_all()
|
||
return self
|
||
|
||
def is_valid(self, data: dict) -> bool:
|
||
if not self.validator:
|
||
raise NotImplementedError(f'Cannot validate {self.name} without validator!')
|
||
allowed_fields = {i.name for i in fields(self.validator)}
|
||
try:
|
||
self.validator(**{k: v for k, v in data.items() if k in allowed_fields})
|
||
return True
|
||
except Exception as e:
|
||
logger_.debug(f'Failed to validate {self.name} config', exc_info=e)
|
||
return False
|
||
|
||
def import_all(self, send_events: bool = False):
|
||
"""Read every entity file in the directory into the in-memory store."""
|
||
for slug in self.discover():
|
||
file = self.get_file(slug)
|
||
try:
|
||
data = read_config_data(file)
|
||
if not self.is_valid(data):
|
||
logger_.error(f'Refusing to import invalid {self.name} config: {file}')
|
||
continue
|
||
self._configs[slug] = data
|
||
except Exception as e:
|
||
logger_.error(f'Failed to import {self.name} config: {file}', exc_info=e)
|
||
self.successful_import = True
|
||
|
||
@property
|
||
def successful_import(self) -> bool:
|
||
from wrolpi.api_utils import api_app
|
||
try:
|
||
return bool(api_app.shared_ctx.configs_imported.get(self.name))
|
||
except AttributeError:
|
||
# Shared context may not be attached (e.g. pure-unit tests); the flag is only informational here.
|
||
return False
|
||
|
||
@successful_import.setter
|
||
def successful_import(self, value: bool):
|
||
from wrolpi.api_utils import api_app
|
||
try:
|
||
api_app.shared_ctx.configs_imported[self.name] = value
|
||
except AttributeError:
|
||
pass
|
||
|
||
def get(self, slug: str) -> Optional[dict]:
|
||
"""Return a deepcopy of one entity's config, or None."""
|
||
data = self._configs.get(slug)
|
||
return deepcopy(data) if data is not None else None
|
||
|
||
def all(self) -> List[dict]:
|
||
return [deepcopy(v) for v in self._configs.values()]
|
||
|
||
@property
|
||
def slugs(self) -> List[str]:
|
||
return list(self._configs.keys())
|
||
|
||
def save_lock(self):
|
||
"""The cross-process save lock when the app context is attached; a no-op lock otherwise (pure-unit tests).
|
||
|
||
Hold this around a whole read-check-mutate-persist cycle (not just the file write) to avoid lost writes
|
||
under concurrent saves. It is NOT reentrant — call `_persist_slug` (not `_save_slug`) while holding it."""
|
||
from wrolpi.api_utils import api_app
|
||
return getattr(api_app.shared_ctx, 'config_save_lock', None) or contextlib.nullcontext()
|
||
|
||
def _persist_slug(self, slug: str):
|
||
"""Write one entity's in-memory config to its YAML file (version bump + daily backup).
|
||
|
||
The caller MUST already hold `save_lock()`."""
|
||
file = self.get_file(slug)
|
||
if PYTEST and not is_tempfile(file):
|
||
raise ValueError(f'Refusing to save config file while testing: {file}')
|
||
data = dict(self._configs[slug])
|
||
# Backup the existing file before overwriting.
|
||
if file.is_file():
|
||
backup_file = self._get_backup_filename(slug)
|
||
backup_file.parent.mkdir(parents=True, exist_ok=True)
|
||
shutil.copy(file, backup_file)
|
||
data['version'] = (data.get('version') or 0) + 1
|
||
file.parent.mkdir(parents=True, exist_ok=True)
|
||
write_config_data(data, file, width=self.width)
|
||
# Persist the bumped version back to the shared store.
|
||
self._configs[slug] = data
|
||
|
||
def _save_slug(self, slug: str):
|
||
"""Acquire the save lock and write one entity to its file."""
|
||
with self.save_lock():
|
||
self._persist_slug(slug)
|
||
|
||
def _get_backup_filename(self, slug: str) -> Path:
|
||
date_str = now().strftime('%Y%m%d')
|
||
return get_media_directory() / 'config' / 'backup' / self.subdirectory / f'{slug}-{date_str}.yaml'
|
||
|
||
def _backup_directory(self) -> Path:
|
||
return get_media_directory() / 'config' / 'backup' / self.subdirectory
|
||
|
||
def get_backup_dates(self, slug: str) -> list:
|
||
"""Scan the entity's backup directory for `<slug>-YYYYMMDD.yaml` files, return the dates newest-first."""
|
||
backup_dir = self._backup_directory()
|
||
if not backup_dir.is_dir():
|
||
return []
|
||
dates = []
|
||
for f in backup_dir.iterdir():
|
||
# The length+digit check rejects a different slug that merely shares this slug's prefix
|
||
# (e.g. slug 'food' must not match 'food-storage-20260101.yaml').
|
||
if f.is_file() and f.suffix == '.yaml' and f.name.startswith(f'{slug}-'):
|
||
date_str = f.name[len(slug) + 1:-len('.yaml')]
|
||
if len(date_str) == 8 and date_str.isdigit():
|
||
dates.append(date_str)
|
||
dates.sort(reverse=True)
|
||
return dates
|
||
|
||
def _get_backup_file(self, slug: str, backup_date: str) -> Path:
|
||
# `backup_date` is request-controlled; constrain it to YYYYMMDD so it cannot traverse out of the backup
|
||
# directory (e.g. '../../etc/something'). `slug` is already constrained to a known, in-memory entity.
|
||
if not re.fullmatch(r'\d{8}', backup_date or ''):
|
||
raise ValidationError('backup_date must be YYYYMMDD')
|
||
return self._backup_directory() / f'{slug}-{backup_date}.yaml'
|
||
|
||
def reimport(self) -> List[dict]:
|
||
"""Re-read every entity file from disk into memory so hand-edits and copied-in files are picked up.
|
||
|
||
The in-memory store is made to match the directory: entities whose file was deleted on disk are dropped,
|
||
edited files are reloaded, and newly-added files are discovered. Returns the resulting entities."""
|
||
discovered = set(self.discover())
|
||
for slug in list(self._configs.keys()):
|
||
if slug not in discovered:
|
||
self._configs.pop(slug, None)
|
||
self.import_all()
|
||
return self.all()
|
||
|
||
def dump_all(self):
|
||
"""Write every in-memory entity to disk."""
|
||
for slug in list(self._configs.keys()):
|
||
try:
|
||
self._save_slug(slug)
|
||
except Exception as e:
|
||
logger_.error(f'Failed to save {self.name} config: {slug}', exc_info=e)
|
||
|
||
def save_later(self, slug: str = None):
|
||
"""Schedule a background save of all entities (coalesces rapid edits)."""
|
||
self.background_save.activate_switch()
|
||
|
||
def create(self, slug: str, data: dict):
|
||
"""Add a new entity and persist it immediately."""
|
||
if slug in self._configs:
|
||
raise ValueError(f'{self.name} already exists: {slug}')
|
||
entity = deepcopy(self.default_entity)
|
||
entity.update(data)
|
||
entity['version'] = 0
|
||
if not self.is_valid(entity):
|
||
raise ValidationError(f'Invalid {self.name} config: {entity}')
|
||
self._configs[slug] = entity
|
||
self._save_slug(slug)
|
||
|
||
def replace(self, slug: str, data: dict):
|
||
"""Replace an entity's config wholesale and persist it (background)."""
|
||
if slug not in self._configs:
|
||
raise ValueError(f'No {self.name} exists: {slug}')
|
||
entity = dict(self._configs[slug])
|
||
entity.update(data)
|
||
if not self.is_valid(entity):
|
||
raise ValidationError(f'Invalid {self.name} config: {entity}')
|
||
self._configs[slug] = entity
|
||
self.save_later(slug)
|
||
|
||
def rename(self, old_slug: str, new_slug: str):
|
||
"""Rename an entity: write the new file, then remove the old one."""
|
||
if old_slug not in self._configs:
|
||
raise ValueError(f'No {self.name} exists: {old_slug}')
|
||
if new_slug in self._configs:
|
||
raise ValueError(f'{self.name} already exists: {new_slug}')
|
||
entity = dict(self._configs[old_slug])
|
||
entity['slug'] = new_slug
|
||
self._configs[new_slug] = entity
|
||
self._save_slug(new_slug)
|
||
self.delete(old_slug)
|
||
|
||
def delete(self, slug: str):
|
||
"""Remove an entity from memory and delete its file."""
|
||
self._configs.pop(slug, None)
|
||
file = self.get_file(slug)
|
||
if file.is_file():
|
||
if PYTEST and not is_tempfile(file):
|
||
raise ValueError(f'Refusing to delete config file while testing: {file}')
|
||
file.unlink()
|
||
|
||
|
||
@dataclass
|
||
class WROLPiConfigValidator:
|
||
archive_destination: str = None
|
||
check_for_upgrades: bool = None
|
||
download_daily_limit_global: int = None
|
||
download_daily_limit_per_domain: int = None
|
||
download_on_startup: bool = None
|
||
download_timeout: int = None
|
||
download_wait: int = None
|
||
download_window_start: str = None
|
||
download_window_end: str = None
|
||
hotspot_device: str = None
|
||
hotspot_on_startup: bool = None
|
||
hotspot_password: str = None
|
||
hotspot_ssid: str = None
|
||
ignore_outdated_zims: bool = None
|
||
log_level: str = None
|
||
map_default_location: dict = None
|
||
map_destination: str = None
|
||
nav_color: str = None
|
||
playlists_destination: str = None
|
||
require_cookies_unlocked: bool = None
|
||
require_media_mounted: bool = None
|
||
save_ffprobe_json: bool = None
|
||
tags_directory: bool = None
|
||
throttle_on_startup: bool = None
|
||
timezone: str = None
|
||
version: int = None
|
||
videos_destination: str = None
|
||
wrol_mode: bool = None
|
||
zims_destination: str = None
|
||
ignored_directories: list[str] = field(default_factory=list)
|
||
|
||
|
||
def get_all_configs() -> Dict[str, ConfigFile]:
|
||
all_configs = dict()
|
||
|
||
if wrolpi_config := get_wrolpi_config():
|
||
all_configs[wrolpi_config.file_name] = wrolpi_config
|
||
|
||
from wrolpi.tags import get_tags_config
|
||
if tags_config := get_tags_config():
|
||
all_configs[tags_config.file_name] = tags_config
|
||
|
||
from modules.videos.lib import get_channels_config
|
||
if channels_config := get_channels_config():
|
||
all_configs[channels_config.file_name] = channels_config
|
||
|
||
# NOTE: Inventory is no longer a single ConfigFile; it is a config-only MultiFileConfig (config/inventory/*.yaml)
|
||
# that is the source-of-truth itself (no DB sync), so it is not registered here. The shared food catalog,
|
||
# however, is a single ConfigFile and is registered.
|
||
from modules.inventory.catalog import get_catalog_config
|
||
if catalog_config := get_catalog_config():
|
||
all_configs[catalog_config.file_name] = catalog_config
|
||
|
||
from modules.videos.lib import get_videos_downloader_config
|
||
if videos_downloader_config := get_videos_downloader_config():
|
||
all_configs[videos_downloader_config.file_name] = videos_downloader_config
|
||
|
||
from wrolpi.downloader import get_download_manager_config
|
||
if download_manager_config := get_download_manager_config():
|
||
all_configs[download_manager_config.file_name] = download_manager_config
|
||
|
||
from modules.archive.lib import get_domains_config, get_archive_downloader_config
|
||
if domains_config := get_domains_config():
|
||
all_configs[domains_config.file_name] = domains_config
|
||
if archive_downloader_config := get_archive_downloader_config():
|
||
all_configs[archive_downloader_config.file_name] = archive_downloader_config
|
||
|
||
from wrolpi.downloader import get_download_cache_config
|
||
if download_cache_config := get_download_cache_config():
|
||
all_configs[download_cache_config.file_name] = download_cache_config
|
||
|
||
from modules.map.pins import get_map_pins_config
|
||
if map_pins_config := get_map_pins_config():
|
||
all_configs[map_pins_config.file_name] = map_pins_config
|
||
|
||
from modules.flasher.config import get_flasher_config
|
||
if flasher_config := get_flasher_config():
|
||
all_configs[flasher_config.file_name] = flasher_config
|
||
|
||
from wrolpi.collections.config import get_playlists_config
|
||
if playlists_config := get_playlists_config():
|
||
all_configs[playlists_config.file_name] = playlists_config
|
||
|
||
return all_configs
|
||
|
||
|
||
async def import_all_db_configs() -> dict[str, bool]:
|
||
"""
|
||
Import all configs that require the database.
|
||
|
||
Returns a dict mapping config name to success status.
|
||
Used by main.py startup and disaster recovery tests.
|
||
|
||
Import order matters:
|
||
1. Tags - no dependencies
|
||
2. Downloads - no dependencies
|
||
3. Channels - uses tags, linked to downloads post-import
|
||
4. Domains - uses tags, linked to downloads post-import
|
||
5. Inventories - no dependencies
|
||
"""
|
||
from wrolpi import tags
|
||
from wrolpi.downloader import import_downloads_config
|
||
from modules.videos.lib import import_channels_config
|
||
from modules.archive.lib import import_domains_config
|
||
from modules.inventory.common import import_inventories_config
|
||
|
||
results = {}
|
||
|
||
# Tags first (channels/domains depend on tags existing)
|
||
try:
|
||
tags.import_tags_config()
|
||
results['tags'] = True
|
||
logger.debug('tags config imported')
|
||
except Exception as e:
|
||
logger.warning(f'Failed to import tags config: {e}')
|
||
results['tags'] = False
|
||
|
||
# Downloads (channels link to downloads post-import)
|
||
try:
|
||
await import_downloads_config()
|
||
results['downloads'] = True
|
||
logger.debug('downloads config imported')
|
||
except Exception as e:
|
||
logger.warning(f'Failed to import downloads config: {e}')
|
||
results['downloads'] = False
|
||
|
||
# Channels (uses tags, links to downloads)
|
||
try:
|
||
import_channels_config()
|
||
results['channels'] = True
|
||
logger.debug('channels config imported')
|
||
except Exception as e:
|
||
logger.warning(f'Failed to import channels config: {e}')
|
||
results['channels'] = False
|
||
|
||
# Domains (uses tags)
|
||
try:
|
||
import_domains_config()
|
||
results['domains'] = True
|
||
logger.debug('domains config imported')
|
||
except Exception as e:
|
||
logger.warning(f'Failed to import domains config: {e}')
|
||
results['domains'] = False
|
||
|
||
# Inventories (config-only, no DB). Migrates any legacy inventories.yaml, then loads per-inventory files.
|
||
try:
|
||
import_inventories_config()
|
||
results['inventories'] = True
|
||
logger.debug('inventories config imported')
|
||
except Exception as e:
|
||
logger.warning(f'Failed to import inventories config: {e}')
|
||
results['inventories'] = False
|
||
|
||
# Playlists (kind='playlist' collections; uses tags for tag_name).
|
||
# NOTE: there is no generic collections config — channels/domains are owned by their own configs
|
||
# (channels.yaml/domains.yaml), and authors/subjects are re-derived on refresh. Playlists are
|
||
# user-curated with no other config home, so they own playlists.yaml.
|
||
try:
|
||
from wrolpi.collections.config import playlists_config
|
||
playlists_config.import_config()
|
||
results['playlists'] = True
|
||
logger.debug('playlists config imported')
|
||
except Exception as e:
|
||
logger.warning(f'Failed to import playlists config: {e}')
|
||
results['playlists'] = False
|
||
|
||
# Map pins (YAML-only, no DB)
|
||
try:
|
||
from modules.map.pins import get_map_pins_config
|
||
get_map_pins_config().import_config()
|
||
results['map_pins'] = True
|
||
logger.debug('map pins config imported')
|
||
except Exception as e:
|
||
logger.warning(f'Failed to import map pins config: {e}')
|
||
results['map_pins'] = False
|
||
|
||
# Flasher saved firmware configurations (YAML-only, no DB)
|
||
try:
|
||
from modules.flasher.config import get_flasher_config
|
||
get_flasher_config().import_config()
|
||
results['flasher'] = True
|
||
logger.debug('flasher config imported')
|
||
except Exception as e:
|
||
logger.warning(f'Failed to import flasher config: {e}')
|
||
results['flasher'] = False
|
||
|
||
return results
|
||
|
||
|
||
def get_config_by_file_name(file_name: str) -> ConfigFile:
|
||
configs = get_all_configs()
|
||
if file_name not in configs:
|
||
raise InvalidConfig(f'No config with {file_name} exists')
|
||
|
||
return configs[file_name]
|
||
|
||
|
||
class WROLPiConfig(ConfigFile):
|
||
file_name = 'wrolpi.yaml'
|
||
default_config = dict(
|
||
archive_destination='archive/%(domain_tag)s/%(domain)s',
|
||
download_daily_limit_global=None,
|
||
download_daily_limit_per_domain=None,
|
||
download_on_startup=True,
|
||
download_timeout=0,
|
||
download_wait=20,
|
||
download_window_start=None,
|
||
download_window_end=None,
|
||
hotspot_device='wlan0',
|
||
hotspot_on_startup=True,
|
||
hotspot_password='wrolpi hotspot',
|
||
hotspot_ssid='WROLPi',
|
||
check_for_upgrades=True,
|
||
ignore_outdated_zims=False,
|
||
ignored_directories=['config', 'tags'],
|
||
log_level='info',
|
||
map_default_location=None,
|
||
map_destination='map',
|
||
nav_color='violet',
|
||
playlists_destination='playlists',
|
||
require_cookies_unlocked=True,
|
||
require_media_mounted=True,
|
||
save_ffprobe_json=True,
|
||
tags_directory=True,
|
||
throttle_on_startup=False,
|
||
timezone=None,
|
||
version=0,
|
||
videos_destination='videos/%(channel_tag)s/%(channel_name)s',
|
||
wrol_mode=False,
|
||
zims_destination='zims',
|
||
)
|
||
validator = WROLPiConfigValidator
|
||
|
||
def get_file(self) -> Path:
|
||
"""WROLPiConfig must be discovered so that other config files can be found."""
|
||
if not self.file_name:
|
||
raise NotImplementedError(f'You must define a file name for this {self.__class__.__name__} config.')
|
||
|
||
if PYTEST:
|
||
return get_media_directory() / f'config/{self.file_name}'
|
||
|
||
# Use the usual "/media/wrolpi/config/wrolpi.yaml" if it exists.
|
||
default_config_path = CONFIG_DIR / self.file_name
|
||
if default_config_path.is_file():
|
||
return default_config_path
|
||
|
||
# Search the media directory for the special "wrolpi.yaml" file. Assume the first one found is the config.
|
||
# Use a depth of 3; multiple drives may exist and the config directory may be down in a second drive.
|
||
if config_path := find_file(get_media_directory(), self.file_name, 3):
|
||
return config_path
|
||
|
||
# Not testing, and can't find file deep in media directory. Use the default.
|
||
return default_config_path
|
||
|
||
def import_config(self, file: pathlib.Path = None, send_events=False):
|
||
try:
|
||
# WROLPiConfig does not sync to database.
|
||
super().import_config(file)
|
||
|
||
# Destinations cannot be empty.
|
||
self._config['archive_destination'] = self.archive_destination or self.default_config['archive_destination']
|
||
self._config['map_destination'] = self.map_destination or self.default_config['map_destination']
|
||
self._config['videos_destination'] = self.videos_destination or self.default_config['videos_destination']
|
||
self._config['zims_destination'] = self.zims_destination or self.default_config['zims_destination']
|
||
self._config['playlists_destination'] = \
|
||
self.playlists_destination or self.default_config['playlists_destination']
|
||
|
||
# The playlists directory holds WROLPi-managed hardlinks/stubs; never index it.
|
||
ignored = list(self._config.get('ignored_directories') or [])
|
||
if self._config['playlists_destination'] not in ignored:
|
||
ignored.append(self._config['playlists_destination'])
|
||
self._config['ignored_directories'] = ignored
|
||
|
||
self.successful_import = True
|
||
except Exception as e:
|
||
self.successful_import = False
|
||
message = f'Failed to import {self.file_name}'
|
||
logger.error(message, exc_info=e)
|
||
if send_events:
|
||
from wrolpi.events import Events
|
||
Events.send_config_import_failed(message)
|
||
raise
|
||
|
||
@property
|
||
def download_on_startup(self) -> bool:
|
||
return self._config['download_on_startup']
|
||
|
||
@download_on_startup.setter
|
||
def download_on_startup(self, value: bool):
|
||
self.update({'download_on_startup': value})
|
||
|
||
@property
|
||
def require_cookies_unlocked(self) -> bool:
|
||
return self._config.get('require_cookies_unlocked', True)
|
||
|
||
@require_cookies_unlocked.setter
|
||
def require_cookies_unlocked(self, value: bool):
|
||
self.update({'require_cookies_unlocked': value})
|
||
|
||
@property
|
||
def require_media_mounted(self) -> bool:
|
||
return self._config.get('require_media_mounted', True)
|
||
|
||
@require_media_mounted.setter
|
||
def require_media_mounted(self, value: bool):
|
||
self.update({'require_media_mounted': value})
|
||
|
||
@property
|
||
def download_timeout(self) -> int:
|
||
return self._config['download_timeout']
|
||
|
||
@download_timeout.setter
|
||
def download_timeout(self, value: int):
|
||
self.update({'download_timeout': value})
|
||
|
||
@property
|
||
def download_wait(self) -> int:
|
||
return self._config.get('download_wait', 20)
|
||
|
||
@download_wait.setter
|
||
def download_wait(self, value: int):
|
||
self.update({'download_wait': value})
|
||
|
||
@property
|
||
def download_daily_limit_per_domain(self) -> Optional[int]:
|
||
"""Maximum non-recurring downloads allowed per domain per day. ``None``/``0`` means unlimited."""
|
||
return self._config.get('download_daily_limit_per_domain')
|
||
|
||
@download_daily_limit_per_domain.setter
|
||
def download_daily_limit_per_domain(self, value: Optional[int]):
|
||
self.update({'download_daily_limit_per_domain': value})
|
||
|
||
@property
|
||
def download_daily_limit_global(self) -> Optional[int]:
|
||
"""Maximum non-recurring downloads allowed across all domains per day. ``None``/``0`` means unlimited."""
|
||
return self._config.get('download_daily_limit_global')
|
||
|
||
@download_daily_limit_global.setter
|
||
def download_daily_limit_global(self, value: Optional[int]):
|
||
self.update({'download_daily_limit_global': value})
|
||
|
||
@staticmethod
|
||
def _normalize_time_value(value) -> str:
|
||
"""YAML may parse HH:MM values as sexagesimal integers (e.g. 21:00 -> 1260). Convert back to HH:MM."""
|
||
if value is None:
|
||
return None
|
||
if isinstance(value, int):
|
||
hours, minutes = divmod(value, 60)
|
||
return f'{hours:02d}:{minutes:02d}'
|
||
return str(value)
|
||
|
||
@property
|
||
def download_window_start(self) -> str:
|
||
return self._normalize_time_value(self._config.get('download_window_start'))
|
||
|
||
@download_window_start.setter
|
||
def download_window_start(self, value: str):
|
||
self.update({'download_window_start': value})
|
||
|
||
@property
|
||
def download_window_end(self) -> str:
|
||
return self._normalize_time_value(self._config.get('download_window_end'))
|
||
|
||
@download_window_end.setter
|
||
def download_window_end(self, value: str):
|
||
self.update({'download_window_end': value})
|
||
|
||
@property
|
||
def hotspot_device(self) -> str:
|
||
return self._config['hotspot_device']
|
||
|
||
@hotspot_device.setter
|
||
def hotspot_device(self, value: str):
|
||
self.update({'hotspot_device': value})
|
||
|
||
@property
|
||
def hotspot_on_startup(self) -> bool:
|
||
return self._config['hotspot_on_startup']
|
||
|
||
@hotspot_on_startup.setter
|
||
def hotspot_on_startup(self, value: bool):
|
||
self.update({'hotspot_on_startup': value})
|
||
|
||
@property
|
||
def hotspot_password(self) -> str:
|
||
return self._config['hotspot_password']
|
||
|
||
@hotspot_password.setter
|
||
def hotspot_password(self, value: str):
|
||
self.update({'hotspot_password': value})
|
||
|
||
@property
|
||
def hotspot_ssid(self) -> str:
|
||
return self._config['hotspot_ssid']
|
||
|
||
@hotspot_ssid.setter
|
||
def hotspot_ssid(self, value: str):
|
||
self.update({'hotspot_ssid': value})
|
||
|
||
@property
|
||
def ignore_outdated_zims(self) -> bool:
|
||
return self._config['ignore_outdated_zims']
|
||
|
||
@ignore_outdated_zims.setter
|
||
def ignore_outdated_zims(self, value: bool):
|
||
self.update({'ignore_outdated_zims': value})
|
||
|
||
@property
|
||
def save_ffprobe_json(self) -> bool:
|
||
return self._config['save_ffprobe_json']
|
||
|
||
@save_ffprobe_json.setter
|
||
def save_ffprobe_json(self, value: bool):
|
||
self.update({'save_ffprobe_json': value})
|
||
|
||
@property
|
||
def log_level(self) -> str:
|
||
return self._config.get('log_level', 'info')
|
||
|
||
@log_level.setter
|
||
def log_level(self, value: str):
|
||
self.update({'log_level': value})
|
||
|
||
@property
|
||
def check_for_upgrades(self) -> bool:
|
||
return self._config.get('check_for_upgrades', True)
|
||
|
||
@check_for_upgrades.setter
|
||
def check_for_upgrades(self, value: bool):
|
||
self.update({'check_for_upgrades': value})
|
||
|
||
@property
|
||
def throttle_on_startup(self) -> bool:
|
||
return self._config['throttle_on_startup']
|
||
|
||
@throttle_on_startup.setter
|
||
def throttle_on_startup(self, value: bool):
|
||
self.update({'throttle_on_startup': value})
|
||
|
||
@property
|
||
def wrol_mode(self) -> bool:
|
||
return self._config['wrol_mode']
|
||
|
||
@wrol_mode.setter
|
||
def wrol_mode(self, value: bool):
|
||
self.update({'wrol_mode': value})
|
||
|
||
@property
|
||
def ignored_directories(self) -> List[str]:
|
||
return self._config['ignored_directories']
|
||
|
||
@ignored_directories.setter
|
||
def ignored_directories(self, value: List[str]):
|
||
self.update({'ignored_directories': value})
|
||
|
||
@property
|
||
def videos_destination(self) -> str:
|
||
return self._config['videos_destination']
|
||
|
||
@videos_destination.setter
|
||
def videos_destination(self, value: str):
|
||
self.update({'videos_destination': value})
|
||
|
||
@property
|
||
def archive_destination(self) -> str:
|
||
return self._config['archive_destination']
|
||
|
||
@archive_destination.setter
|
||
def archive_destination(self, value: str):
|
||
self.update({'archive_destination': value})
|
||
|
||
@property
|
||
def playlists_destination(self) -> str:
|
||
return self._config['playlists_destination']
|
||
|
||
@playlists_destination.setter
|
||
def playlists_destination(self, value: str):
|
||
self.update({'playlists_destination': value})
|
||
|
||
@property
|
||
def map_default_location(self) -> dict:
|
||
return self._config.get('map_default_location')
|
||
|
||
@map_default_location.setter
|
||
def map_default_location(self, value: dict):
|
||
self.update({'map_default_location': value})
|
||
|
||
@property
|
||
def map_destination(self) -> str:
|
||
return self._config['map_destination']
|
||
|
||
@map_destination.setter
|
||
def map_destination(self, value: str):
|
||
self.update({'map_destination': value})
|
||
|
||
@property
|
||
def zims_destination(self) -> str:
|
||
return self._config['zims_destination']
|
||
|
||
@zims_destination.setter
|
||
def zims_destination(self, value: str):
|
||
self.update({'zims_destination': value})
|
||
|
||
@property
|
||
def nav_color(self) -> str:
|
||
return self._config['nav_color']
|
||
|
||
@nav_color.setter
|
||
def nav_color(self, value: str):
|
||
self.update({'nav_color': value})
|
||
|
||
@property
|
||
def tags_directory(self) -> bool:
|
||
return self._config.get('tags_directory', True)
|
||
|
||
@tags_directory.setter
|
||
def tags_directory(self, value: bool):
|
||
self.update({'tags_directory': value})
|
||
|
||
@property
|
||
def timezone(self) -> str:
|
||
return self._config.get('timezone')
|
||
|
||
@timezone.setter
|
||
def timezone(self, value: str):
|
||
self.update({'timezone': value})
|
||
|
||
|
||
WROLPI_CONFIG: WROLPiConfig = WROLPiConfig()
|
||
TEST_WROLPI_CONFIG: WROLPiConfig = None
|
||
|
||
|
||
def get_wrolpi_config() -> WROLPiConfig:
|
||
"""Read the global WROLPi yaml config file."""
|
||
global TEST_WROLPI_CONFIG
|
||
if isinstance(TEST_WROLPI_CONFIG, WROLPiConfig):
|
||
return TEST_WROLPI_CONFIG
|
||
|
||
global WROLPI_CONFIG
|
||
return WROLPI_CONFIG
|
||
|
||
|
||
def set_test_config(enable: bool):
|
||
global TEST_WROLPI_CONFIG
|
||
TEST_WROLPI_CONFIG = WROLPiConfig() if enable else None
|
||
|
||
|
||
def wrol_mode_enabled() -> bool:
|
||
"""Return True if WROL Mode is enabled."""
|
||
return get_wrolpi_config().wrol_mode
|
||
|
||
|
||
def wrol_mode_check(func):
|
||
"""Wraps a function so that it cannot be called when WROL Mode is enabled."""
|
||
|
||
@wraps(func)
|
||
def check(*a, **kw):
|
||
if wrol_mode_enabled():
|
||
raise WROLModeEnabled()
|
||
|
||
# WROL Mode is not enabled, run the function as normal.
|
||
result = func(*a, **kw)
|
||
return result
|
||
|
||
return check
|
||
|
||
|
||
def enable_wrol_mode():
|
||
"""
|
||
Modify config to enable WROL Mode.
|
||
|
||
Stop downloads and Download Manager.
|
||
"""
|
||
logger_.warning('ENABLING WROL MODE')
|
||
get_wrolpi_config().wrol_mode = True
|
||
from wrolpi.downloader import download_manager
|
||
download_manager.stop()
|
||
|
||
|
||
async def disable_wrol_mode():
|
||
"""
|
||
Modify config to disable WROL Mode.
|
||
|
||
Start downloads and Download Manager.
|
||
"""
|
||
logger_.warning('DISABLING WROL MODE')
|
||
get_wrolpi_config().wrol_mode = False
|
||
from wrolpi.downloader import download_manager
|
||
await download_manager.enable()
|
||
|
||
|
||
def insert_parameter(func: Callable, parameter_name: str, item, args: Tuple, kwargs: Dict) -> Tuple[Tuple, Dict]:
|
||
"""Insert a parameter wherever it fits in the Callable's signature."""
|
||
sig = inspect.signature(func)
|
||
if parameter_name not in sig.parameters:
|
||
raise TypeError(f'Function {func} MUST have a {parameter_name} parameter!')
|
||
|
||
args = list(args)
|
||
|
||
index = list(sig.parameters).index(parameter_name)
|
||
args.insert(index, item)
|
||
args = tuple(args)
|
||
|
||
return args, kwargs
|
||
|
||
|
||
def iterify(kind: type = list):
|
||
"""Convenience function to convert the output of the wrapped function to the type provided."""
|
||
|
||
def wrapper(func):
|
||
@wraps(func)
|
||
def wrapped(*a, **kw):
|
||
result = func(*a, **kw)
|
||
return kind(result)
|
||
|
||
return wrapped
|
||
|
||
return wrapper
|
||
|
||
|
||
dt_or_d = Union[datetime, date]
|
||
|
||
|
||
def date_range(start: dt_or_d, end: dt_or_d, steps: int) -> List[dt_or_d]:
|
||
delta = (end - start) // steps
|
||
return [start + (delta * i) for i in range(steps)]
|
||
|
||
|
||
WHITESPACE = re.compile(r'\s')
|
||
|
||
|
||
def remove_whitespace(s: str) -> str:
|
||
return WHITESPACE.sub('', s)
|
||
|
||
|
||
RUN_AFTER = False
|
||
|
||
|
||
def run_after(after: callable, *args, **kwargs) -> callable:
|
||
"""Run the `after` function sometime in the future ofter the wrapped function returns."""
|
||
if not inspect.iscoroutinefunction(after):
|
||
synchronous_after = after
|
||
|
||
async def after(*a, **kw):
|
||
if logger.isEnabledFor(logging.DEBUG):
|
||
logger.debug(f'run_after synchronous_after called for {after}')
|
||
return synchronous_after(*a, **kw)
|
||
|
||
def wrapper(func: callable):
|
||
if PYTEST and RUN_AFTER is False:
|
||
return func
|
||
|
||
if inspect.iscoroutinefunction(func):
|
||
@wraps(func)
|
||
async def wrapped(*a, **kw):
|
||
results = await func(*a, **kw)
|
||
if logger.isEnabledFor(logging.DEBUG):
|
||
logger.debug(f'run_after async called for {func}')
|
||
coro = after(*args, **kwargs)
|
||
asyncio.ensure_future(coro)
|
||
return results
|
||
else:
|
||
@wraps(func)
|
||
def wrapped(*a, **kw):
|
||
results = func(*a, **kw)
|
||
if logger.isEnabledFor(logging.DEBUG):
|
||
logger.debug(f'run_after sync called for {func}')
|
||
coro = after(*args, **kwargs)
|
||
asyncio.ensure_future(coro)
|
||
return results
|
||
|
||
return wrapped
|
||
|
||
return wrapper
|
||
|
||
|
||
MEDIA_DIRECTORY_PERMISSIONS = 0o40755
|
||
|
||
|
||
def check_media_directory():
|
||
"""
|
||
Check that the media directory exists and has the correct permissions. Log an error if the directory
|
||
is unusable.
|
||
"""
|
||
result = True
|
||
media_directory = get_media_directory()
|
||
if not media_directory.is_dir():
|
||
logger_.error(f'Media directory does not exist: {media_directory}')
|
||
return False
|
||
|
||
permissions = media_directory.stat().st_mode
|
||
if permissions != MEDIA_DIRECTORY_PERMISSIONS:
|
||
logger_.error(f'Media directory has the wrong permissions: {oct(permissions)}')
|
||
result = False
|
||
|
||
if not media_directory.is_mount():
|
||
logger_.warning('Media directory is not a mount.')
|
||
|
||
try:
|
||
# Write a file into the media directory, log an error if this is not possible. This file should be cleaned up
|
||
# by tempfile.
|
||
with tempfile.NamedTemporaryFile(dir=media_directory) as tf:
|
||
tf.write(b'test')
|
||
except Exception as e:
|
||
logger_.error(f'Could not write to media directory: {media_directory}', exc_info=e)
|
||
result = False
|
||
|
||
return result
|
||
|
||
|
||
def get_absolute_media_path(path: str | Path) -> Path:
|
||
"""
|
||
Get the absolute path of file/directory within the config media directory.
|
||
|
||
>>> get_media_directory()
|
||
Path('/media')
|
||
>>> get_absolute_media_path('videos/blender')
|
||
Path('/media/videos/blender')
|
||
|
||
:raises UnknownDirectory: the directory/path doesn't exist
|
||
"""
|
||
if not path:
|
||
raise ValueError('Path cannot be empty!')
|
||
path = get_media_directory() / path
|
||
return path
|
||
|
||
|
||
def get_relative_to_media_directory(path: Union[str, Path]) -> Path:
|
||
"""
|
||
Get the path for a file/directory relative to the config media directory.
|
||
|
||
>>> get_media_directory()
|
||
Path('/media')
|
||
>>> get_relative_to_media_directory('/media/videos/blender')
|
||
Path('videos/blender')
|
||
|
||
:raises UnknownDirectory: the directory/path doesn't exist
|
||
"""
|
||
absolute = get_absolute_media_path(path)
|
||
return absolute.relative_to(get_media_directory())
|
||
|
||
|
||
def get_paths_in_parent_directory(paths: List[Union[str, Path]], parent_directory: pathlib.Path) -> List[Path]:
|
||
"""Return only those of the provided paths that are children of the provided parent directory."""
|
||
paths = [i if isinstance(i, Path) else Path(i) for i in paths]
|
||
parent_directory = parent_directory.resolve()
|
||
new_paths = []
|
||
for path in paths:
|
||
if parent_directory in path.resolve().parents and path != parent_directory:
|
||
new_paths.append(path)
|
||
return new_paths
|
||
|
||
|
||
def get_paths_in_media_directory(paths: List[Union[str, Path]]) -> List[Path]:
|
||
"""Return only those of the provided paths that are children of the media directory."""
|
||
return get_paths_in_parent_directory(paths, get_media_directory())
|
||
|
||
|
||
def minimize_dict(d: dict, keys: Iterable) -> Optional[dict]:
|
||
"""Return a new dictionary that contains only the keys provided."""
|
||
if d is None:
|
||
return
|
||
return {k: d[k] for k in set(keys) & d.keys()}
|
||
|
||
|
||
def make_media_directory(path: Union[str, Path]):
|
||
"""Make a directory relative within the media directory."""
|
||
media_dir = get_media_directory()
|
||
path = media_dir / str(path)
|
||
path.mkdir(parents=True)
|
||
|
||
|
||
def extract_domain(url: str) -> str:
|
||
"""
|
||
Extract the domain from a URL. Remove leading www and port.
|
||
|
||
>>> extract_domain('https://www.example.com/foo')
|
||
'example.com'
|
||
>>> extract_domain('https://example.com:443/foo')
|
||
'example.com'
|
||
"""
|
||
parsed = urlparse(url)
|
||
domain = parsed.netloc
|
||
if not domain:
|
||
raise ValueError(f'URL does not have a domain: {url=}')
|
||
domain = domain.decode() if hasattr(domain, 'decode') else domain
|
||
# Remove port if present
|
||
if ':' in domain:
|
||
domain = domain.split(':')[0]
|
||
if domain.startswith('www.'):
|
||
# Remove leading www.
|
||
domain = domain[4:]
|
||
return domain
|
||
|
||
|
||
# Common subdomains stripped when grouping downloads by domain (e.g. for daily download limits).
|
||
COMMON_DOWNLOAD_SUBDOMAINS = ('m.', 'mobile.')
|
||
|
||
|
||
def normalize_domain(url: str) -> str:
|
||
"""Return the canonical domain used to group downloads (e.g. for daily download limits).
|
||
|
||
Strips the port, a leading ``www.`` (via :func:`extract_domain`), and other common subdomains
|
||
(``m.``, ``mobile.``) so that, for example, ``m.example.com`` and ``www.example.com`` share one bucket.
|
||
|
||
Downloaders may apply type-specific URL normalization before calling this (see
|
||
``Downloader.normalize_domain``).
|
||
|
||
>>> normalize_domain('https://m.example.com/foo')
|
||
'example.com'
|
||
>>> normalize_domain('https://www.example.com:443/foo')
|
||
'example.com'
|
||
"""
|
||
domain = extract_domain(url)
|
||
for subdomain in COMMON_DOWNLOAD_SUBDOMAINS:
|
||
if domain.startswith(subdomain):
|
||
domain = domain[len(subdomain):]
|
||
break
|
||
return domain
|
||
|
||
|
||
def api_param_limiter(maximum: int, default: int = 20) -> callable:
|
||
"""Create a function which restricts the maximum number that can be returned.
|
||
Useful for restricting API limit params.
|
||
|
||
>>> limiter = api_param_limiter(100)
|
||
>>> limiter(0)
|
||
0
|
||
>>> limiter(100)
|
||
100
|
||
>>> limiter(150)
|
||
100
|
||
"""
|
||
|
||
def limiter_(i: int) -> int:
|
||
if not i:
|
||
# Limit is not valid, use the default.
|
||
return default
|
||
return min(int(i), maximum)
|
||
|
||
return limiter_
|
||
|
||
|
||
def partition(pred, iterable):
|
||
"""Use a predicate to partition entries into false entries and true entries
|
||
|
||
>>> is_even = lambda i: not i % 2
|
||
>>> partition(is_even, range(10))
|
||
([0, 2, 4, 6, 8], [1, 3, 5, 7, 9])
|
||
"""
|
||
t1, t2 = tee(iterable)
|
||
return list(filter(pred, t2)), list(filterfalse(pred, t1))
|
||
|
||
|
||
@contextlib.contextmanager
|
||
def chdir(directory: Union[pathlib.Path, str, None] = None, with_home: bool = False):
|
||
"""
|
||
Change the Python working directory in this context. If no directory is passed, a TemporaryDirectory will be used.
|
||
|
||
This also changes the $HOME environment variable to match.
|
||
"""
|
||
home_exists = 'HOME' in os.environ
|
||
original_home = os.environ.get('HOME')
|
||
cwd = os.getcwd()
|
||
if directory is None:
|
||
with tempfile.TemporaryDirectory() as d:
|
||
d = os.path.realpath(d)
|
||
try:
|
||
os.chdir(d)
|
||
if with_home:
|
||
os.environ['HOME'] = d
|
||
yield
|
||
finally:
|
||
os.chdir(cwd)
|
||
if home_exists:
|
||
os.environ['HOME'] = original_home
|
||
else:
|
||
del os.environ['HOME']
|
||
return
|
||
|
||
try:
|
||
os.chdir(directory)
|
||
if with_home:
|
||
os.environ['HOME'] = str(directory)
|
||
yield
|
||
finally:
|
||
os.chdir(cwd)
|
||
if home_exists:
|
||
os.environ['HOME'] = original_home
|
||
else:
|
||
del os.environ['HOME']
|
||
return
|
||
|
||
|
||
ZIG_TYPE = Union[int, float, complex, Decimal, datetime]
|
||
|
||
|
||
def zig_zag(low: ZIG_TYPE, high: ZIG_TYPE) -> Generator[ZIG_TYPE, None, None]:
|
||
"""
|
||
Generate numbers between `low` and `high` that are
|
||
spread out evenly. Produces infinite unique results.
|
||
|
||
>>> import itertools
|
||
>>> list(itertools.islice(zig_zag(0, 10), 10))
|
||
[0, 5, 2, 7, 1, 3, 6, 8, 4, 9]
|
||
>>> list(itertools.islice(zig_zag(0, 5), 5))
|
||
[0, 2, 1, 3, 4]
|
||
>>> list(itertools.islice(zig_zag(50.0, 100.0), 11))
|
||
[50.0, 75.0, 62.5, 87.5, 56.25, 68.75, 81.25, 93.75, 53.125, 59.375, 65.625]
|
||
"""
|
||
output_type = type(low)
|
||
if not isinstance(high, type(low)):
|
||
raise ValueError(f'high and low must be same type')
|
||
if isinstance(low, datetime) and isinstance(high, datetime):
|
||
low, high = low.timestamp(), high.timestamp()
|
||
output_type = from_timestamp
|
||
|
||
# Special thanks to my wife for helping me solve this! :*
|
||
results = set()
|
||
num = low
|
||
divisor = 2
|
||
diff = high - low
|
||
while True:
|
||
output_value = output_type(num)
|
||
if output_value not in results:
|
||
yield output_value
|
||
results.add(output_value)
|
||
num += diff / divisor
|
||
if num >= high:
|
||
divisor *= 2
|
||
num = low + (diff / divisor)
|
||
|
||
|
||
def walk(path: Path) -> Generator[Path, None, None]:
|
||
"""Recursively walk a directory structure yielding all files and directories."""
|
||
if not path.is_dir():
|
||
raise ValueError('Can only walk a directory.')
|
||
|
||
for path in path.iterdir():
|
||
yield path
|
||
if path.is_dir():
|
||
yield from walk(path)
|
||
|
||
|
||
def get_files_and_directories(directory: Path):
|
||
"""Walk a directory, return a tuple of all files and all directories within. (Not recursive)."""
|
||
if not directory.exists():
|
||
raise FileNotFoundError(f'Directory does not exist: {directory}')
|
||
if not directory.is_dir():
|
||
raise ValueError('Can only walk a directory.')
|
||
|
||
directories, files = partition(lambda i: i.is_dir(), directory.iterdir())
|
||
return files, directories
|
||
|
||
|
||
# These characters are invalid in Windows or Linux.
|
||
INVALID_FILE_CHARS = re.compile(r'[/<>:|"\\?*%!\n\r]')
|
||
|
||
SPACE_FILE_CHARS = re.compile(r'( +)|(\t+)')
|
||
|
||
|
||
def escape_file_name(name: str) -> str:
|
||
"""Remove any invalid characters in a file name."""
|
||
# Replace forward-slash (linux directories) with unicode Big Solidus (U+29F8)
|
||
name = name.replace('/', '⧸')
|
||
# Replace commonly used pipe with dash.
|
||
name = name.replace(' | ', ' - ')
|
||
# Replace multiple spaces or tabs with a single space.
|
||
name = SPACE_FILE_CHARS.sub(' ', name)
|
||
name = INVALID_FILE_CHARS.sub('', name)
|
||
return name.strip()
|
||
|
||
|
||
# Maximum bytes for a filename. Linux ext4/xfs limit is 255 bytes.
|
||
# We use 200 to leave room for FileGroup suffix variations (e.g., .readability.html is 17 bytes vs .html at 5 bytes).
|
||
MAXIMUM_FILE_BYTES = 200
|
||
|
||
|
||
def trim_file_name(path: str | pathlib.Path) -> str | pathlib.Path:
|
||
"""Shorten the file name only if it is longer than the file system supports. Trim from the end of the name until
|
||
the name is short enough (preserving any suffix).
|
||
|
||
Uses BYTE length (not character length) because UTF-8 multi-byte characters (e.g., Devanagari, Chinese)
|
||
can exceed the 255 byte filesystem limit even with fewer than 140 characters.
|
||
"""
|
||
name_type = pathlib.Path if isinstance(path, pathlib.Path) else str
|
||
parent = path.parent if isinstance(path, pathlib.Path) else '/'.join(i for i in path.split('/')[:-1])
|
||
filename = path.name if isinstance(path, pathlib.Path) else path.split('/')[-1] if '/' in path else path
|
||
|
||
filename_bytes = len(filename.encode('utf-8'))
|
||
if filename_bytes <= MAXIMUM_FILE_BYTES:
|
||
# No trimming needed, but preserve parent directory if present
|
||
if parent:
|
||
if name_type == pathlib.Path:
|
||
return parent / filename
|
||
return f'{parent}/{filename}'
|
||
return name_type(filename)
|
||
|
||
# Trim based on byte length, not character length.
|
||
# This is because a FileGroup will have varying filename lengths.
|
||
from wrolpi.files.lib import split_path_stem_and_suffix
|
||
stem, suffix = split_path_stem_and_suffix(filename)
|
||
suffix_bytes = len(suffix.encode('utf-8'))
|
||
target_stem_bytes = MAXIMUM_FILE_BYTES - suffix_bytes
|
||
|
||
# Trim characters from end of stem until byte length fits
|
||
trimmed_stem = stem
|
||
while len(trimmed_stem.encode('utf-8')) > target_stem_bytes and trimmed_stem:
|
||
trimmed_stem = trimmed_stem[:-1]
|
||
trimmed_stem = trimmed_stem.strip()
|
||
|
||
new_name = trimmed_stem + suffix
|
||
if parent:
|
||
if name_type == pathlib.Path:
|
||
return parent / new_name
|
||
return f'{parent}/{new_name}'
|
||
return name_type(new_name)
|
||
|
||
|
||
def native_only(func: callable):
|
||
"""Wraps a function. Raises NativeOnly if run while Dockerized."""
|
||
|
||
@wraps(func)
|
||
def wrapped(*a, **kw):
|
||
if DOCKERIZED:
|
||
raise NativeOnly('Only supported on a Raspberry Pi')
|
||
return func(*a, **kw)
|
||
|
||
return wrapped
|
||
|
||
|
||
def recursive_map(obj: Any, func: callable):
|
||
"""
|
||
Apply `func` to all values of any dictionaries; or items in any list/set/tuple.
|
||
|
||
>>> recursive_map({'foo ': ' bar '}, lambda i: i.strip() if hasattr(i, 'strip') else i)
|
||
{'foo ': 'bar'}
|
||
>>> recursive_map(['foo ', ' bar '], lambda i: i.strip() if hasattr(i, 'strip') else i)
|
||
['foo', 'bar']
|
||
"""
|
||
if isinstance(obj, dict):
|
||
obj = obj.copy()
|
||
for key, value in obj.items():
|
||
obj[key] = recursive_map(value, func)
|
||
return obj
|
||
elif isinstance(obj, (list, set, tuple)):
|
||
type_ = type(obj)
|
||
return type_(map(lambda i: recursive_map(i, func), obj))
|
||
return func(obj)
|
||
|
||
|
||
@contextlib.asynccontextmanager
|
||
async def aiohttp_session(timeout: int = None) -> AsyncGenerator[ClientSession, Any]:
|
||
"""Convenience function because aiohttp timeout cannot be None."""
|
||
if timeout:
|
||
timeout = aiohttp.ClientTimeout(total=timeout) if timeout else None
|
||
async with aiohttp.ClientSession(timeout=timeout) as session:
|
||
yield session
|
||
else:
|
||
async with aiohttp.ClientSession() as session:
|
||
yield session
|
||
|
||
|
||
@contextlib.asynccontextmanager
|
||
async def aiohttp_post(url: str, json_, timeout: int = None, headers: dict = None) -> AsyncGenerator[
|
||
ClientResponse, Any]:
|
||
"""Perform an async aiohttp POST request. Yield the response in a session."""
|
||
headers = headers or DEFAULT_HTTP_HEADERS
|
||
async with aiohttp_session(timeout) as session:
|
||
async with session.post(url, json=json_, headers=headers) as response:
|
||
yield response
|
||
|
||
|
||
@contextlib.asynccontextmanager
|
||
async def aiohttp_get(url: str, timeout: int = None, headers: dict = None) -> AsyncGenerator[ClientResponse, Any]:
|
||
"""Perform an async aiohttp GET request. Yield the response in a session."""
|
||
headers = headers or DEFAULT_HTTP_HEADERS
|
||
async with aiohttp_session(timeout) as session:
|
||
async with session.get(url, headers=headers) as response:
|
||
yield response
|
||
|
||
|
||
@contextlib.asynccontextmanager
|
||
async def aiohttp_head(url: str, timeout: int = None, headers: dict = None) -> AsyncGenerator[ClientResponse, Any]:
|
||
"""Perform an async aiohttp HEAD request. Yield the response in a session."""
|
||
headers = headers or DEFAULT_HTTP_HEADERS
|
||
async with aiohttp_session(timeout) as session:
|
||
async with session.head(url, headers=headers) as response:
|
||
yield response
|
||
|
||
|
||
CONTROLLER_URL = 'http://127.0.0.1'
|
||
|
||
|
||
async def set_system_timezone(timezone: str):
|
||
"""Ask the Controller to apply the timezone to the system.
|
||
|
||
This is best-effort: if the Controller is unavailable or rejects the change,
|
||
the application-level timezone in wrolpi.yaml still takes effect."""
|
||
try:
|
||
async with aiohttp_post(f'{CONTROLLER_URL}/api/timezone/set',
|
||
json_={'timezone': timezone}, timeout=10) as resp:
|
||
if resp.status != 200:
|
||
contents = await resp.json()
|
||
logger.warning(f'Controller failed to set system timezone: {contents}')
|
||
except Exception as e:
|
||
logger.warning(f'Could not reach Controller to set system timezone: {e}')
|
||
|
||
|
||
@dataclass
|
||
class DownloadFileInfo:
|
||
"""Information about a file that can be downloaded."""
|
||
name: str = None
|
||
size: int = None
|
||
type: str = None
|
||
accept_ranges: str = None
|
||
status: int = None
|
||
location: str = None
|
||
|
||
|
||
FILENAME_MATCHER = re.compile(r'.*filename="(.*)"')
|
||
|
||
|
||
async def get_download_info(url: str, timeout: int = 60) -> DownloadFileInfo:
|
||
"""Gets information (name, size, etc.) about a downloadable file at the provided URL."""
|
||
async with aiohttp_head(url, timeout) as response:
|
||
download_logger.debug(f'{response.headers=}')
|
||
|
||
info = DownloadFileInfo(
|
||
type=response.headers.get('Content-Type'),
|
||
size=int(response.headers['Content-Length']) if 'Content-Length' in response.headers else None,
|
||
accept_ranges=response.headers.get('Accept-Ranges'),
|
||
status=response.status,
|
||
location=response.headers.get('Location'),
|
||
)
|
||
|
||
disposition = response.headers.get('Content-Disposition')
|
||
|
||
if disposition and 'filename' in disposition:
|
||
if (match := FILENAME_MATCHER.match(disposition)) and (groups := match.groups()):
|
||
info.name = groups[0]
|
||
else:
|
||
# No Content-Disposition with filename, use the Location or URL name.
|
||
if info.location:
|
||
parsed = urlparse(info.location)
|
||
else:
|
||
parsed = urlparse(url)
|
||
info.name = parsed.path.split('/')[-1]
|
||
|
||
# If the resolved name has no extension but the original URL does, prefer the URL's filename.
|
||
# This prevents UUID-like names from CDN redirects.
|
||
url_name = urlparse(url).path.split('/')[-1]
|
||
if info.name and '.' not in info.name and '.' in url_name:
|
||
info.name = url_name
|
||
|
||
return info
|
||
|
||
|
||
download_logger = logger.getChild('download')
|
||
|
||
|
||
def human_bandwidth(bps: int) -> str:
|
||
"""Convert bits per second to a more readable format.
|
||
|
||
>>> human_bandwidth(2000)
|
||
# '2 Kbps'
|
||
"""
|
||
units = ['bps', 'Kbps', 'Mbps', 'Gbps', 'Tbps']
|
||
unit = 0
|
||
while bps > 1000 and unit <= len(units):
|
||
bps /= 1000
|
||
unit += 1
|
||
|
||
return f'{int(bps)} {units[unit]}'
|
||
|
||
|
||
modelers = []
|
||
|
||
|
||
def register_modeler(modeler: callable):
|
||
"""Register a FileGroup modeler. `modeler` must accept `progress_callback=None`."""
|
||
modelers.append(modeler)
|
||
return modeler
|
||
|
||
|
||
async def apply_modelers(progress_callback: callable = None):
|
||
"""Run all registered modelers.
|
||
|
||
Args:
|
||
progress_callback: Optional callback(processed, total) called after each modeler completes,
|
||
or after each batch if the modeler supports progress_callback.
|
||
"""
|
||
from wrolpi.db import get_db_session
|
||
from wrolpi.files.models import FileGroup
|
||
|
||
# Get initial count of unindexed files for progress tracking
|
||
initial_unindexed = 0
|
||
if progress_callback:
|
||
with get_db_session() as session:
|
||
initial_unindexed = session.query(FileGroup).filter(FileGroup.indexed != True).count()
|
||
progress_callback(0, initial_unindexed)
|
||
|
||
# Track cumulative progress across all modelers
|
||
cumulative_processed = 0
|
||
|
||
for modeler in modelers:
|
||
logger_.info(f'Applying modeler {modeler.__name__}')
|
||
try:
|
||
# Modelers must accept progress_callback= (see register_modeler). The leftover
|
||
# indexer (apply_indexers) runs after all of these.
|
||
if progress_callback and initial_unindexed > 0:
|
||
def make_modeler_callback(current_cumulative):
|
||
def modeler_progress_callback(batch_processed: int):
|
||
nonlocal cumulative_processed
|
||
cumulative_processed = current_cumulative + batch_processed
|
||
progress_callback(cumulative_processed, initial_unindexed)
|
||
|
||
return modeler_progress_callback
|
||
|
||
await modeler(progress_callback=make_modeler_callback(cumulative_processed))
|
||
else:
|
||
await modeler()
|
||
except Exception as e:
|
||
logger_.error(f'Modeler {modeler.__name__} raised error', exc_info=e)
|
||
if PYTEST:
|
||
raise
|
||
|
||
# Update progress after each modeler completes (fallback for modelers without progress_callback)
|
||
if progress_callback and initial_unindexed > 0:
|
||
with get_db_session() as session:
|
||
remaining = session.query(FileGroup).filter(FileGroup.indexed != True).count()
|
||
cumulative_processed = initial_unindexed - remaining
|
||
progress_callback(cumulative_processed, initial_unindexed)
|
||
|
||
# Sleep to catch cancel.
|
||
await asyncio.sleep(0)
|
||
|
||
|
||
REFRESH_CLEANUP = []
|
||
|
||
|
||
def register_refresh_cleanup(func: callable):
|
||
REFRESH_CLEANUP.append(func)
|
||
return func
|
||
|
||
|
||
async def apply_refresh_cleanup():
|
||
# TODO convert all functions to async so refresh functions can be used.
|
||
for func in REFRESH_CLEANUP:
|
||
slow_message = f'After refresh cleanup {func.__name__} took %(elapsed)s seconds'
|
||
with slow_logger(5, slow_message, logger__=logger_):
|
||
try:
|
||
logger_.info(f'Applying refresh cleanup {func.__name__}')
|
||
func()
|
||
except Exception as e:
|
||
logger_.error(f'Refresh cleanup {func.__name__} failed!', exc_info=e)
|
||
if PYTEST:
|
||
raise
|
||
# Sleep to catch cancel.
|
||
await asyncio.sleep(0)
|
||
|
||
|
||
def chunks(it: Iterable, size: int):
|
||
"""Split an iterable into iterables of the defined length."""
|
||
it = iter(it)
|
||
return iter(lambda: tuple(islice(it, size)), ())
|
||
|
||
|
||
def chunks_by_stem(it: List[Union[pathlib.Path, str, int]], size: int) -> Generator[List[pathlib.Path], None, None]:
|
||
"""
|
||
Attempt to split a list of paths near the defined size. Keep groups of files together when they share
|
||
matching names.
|
||
|
||
>>> files = ['1.mp4', '1.txt', '2.mp4', '2.txt', '2.png', '3.mp4']
|
||
>>> chunks_by_stem(files, 2)
|
||
[['1.mp4', '1.txt'], ['2.mp4', '2.txt', '2.png'], ['3.mp4']]
|
||
>>> chunks_by_stem(files, 3)
|
||
[['1.mp4', '1.txt', '2.mp4', '2.txt', '2.png'], ['3.mp4']]
|
||
"""
|
||
if not isinstance(size, int) or size < 1:
|
||
raise ValueError('size must be a positive integer')
|
||
|
||
if not it or len(it) < size:
|
||
yield sorted(it)
|
||
return
|
||
|
||
if not isinstance(it, list):
|
||
raise ValueError('iterable must a list')
|
||
|
||
from wrolpi.files.lib import split_path_stem_and_suffix
|
||
it = sorted(it.copy())
|
||
index = size
|
||
last_stem = None
|
||
while it:
|
||
if index >= len(it):
|
||
# Ran out of items, yield what is left.
|
||
yield it
|
||
return
|
||
path = it[index]
|
||
stem, _ = split_path_stem_and_suffix(path)
|
||
if last_stem and stem != last_stem:
|
||
# Found a break in the path names, yield and reset.
|
||
chunk, it = it[:index], it[index:]
|
||
index = size
|
||
last_stem = None
|
||
yield chunk
|
||
continue
|
||
# Didn't find a name change, try again.
|
||
last_stem = stem
|
||
index += 1
|
||
|
||
|
||
@contextlib.contextmanager
|
||
def timer(name, level: str = 'debug', logger__: logging.Logger = None):
|
||
"""Prints out the time elapsed during the call of some block.
|
||
|
||
Example:
|
||
with timer('sleepy'):
|
||
time.sleep(10)
|
||
|
||
"""
|
||
level = level.lower()
|
||
logger__ = logger__ or logger_
|
||
before = datetime.now()
|
||
try:
|
||
yield
|
||
finally:
|
||
elapsed = (datetime.now() - before).total_seconds()
|
||
msg = f'{name} elapsed {elapsed} seconds'
|
||
getattr(logger__, level)(msg)
|
||
|
||
|
||
@contextlib.contextmanager
|
||
def slow_logger(max_seconds: int, message: str, logger__=logger_, level=logging.WARNING):
|
||
"""Only logs when the context duration exceeds the `max_seconds`."""
|
||
before = now()
|
||
try:
|
||
yield
|
||
finally:
|
||
elapsed = (now() - before).total_seconds()
|
||
if elapsed >= max_seconds:
|
||
logger__.log(level, message % dict(elapsed=elapsed))
|
||
|
||
|
||
TIMERS = dict()
|
||
|
||
|
||
@contextlib.contextmanager
|
||
def cum_timer(name: str):
|
||
"""Track time usage within this context's name. The cumulative calls to this context will be printed on exit."""
|
||
before = now()
|
||
try:
|
||
yield
|
||
finally:
|
||
total_seconds = (now() - before).total_seconds()
|
||
existing_seconds, calls = TIMERS.get(name, (0, 0))
|
||
TIMERS[name] = (existing_seconds + total_seconds, calls + 1)
|
||
|
||
|
||
@atexit.register
|
||
def print_timer():
|
||
"""Print any cumulative timers that have been stored during exit."""
|
||
if not TIMERS:
|
||
return
|
||
|
||
total = sum(i[0] for i in TIMERS.values())
|
||
for name, (seconds, calls) in sorted(TIMERS.items(), key=lambda i: i[1]):
|
||
percent = int((seconds / total) * 100)
|
||
seconds = round(seconds, 5)
|
||
print(f'CUM_TIMER: {repr(name)} elapsed {seconds} cumulative seconds ({percent}%)'
|
||
f' {calls} calls', file=sys.stderr)
|
||
|
||
|
||
def limit_concurrent(limit: int, throw: bool = False):
|
||
"""Wrapper that limits the amount of concurrently running functions."""
|
||
sema = multiprocessing.Semaphore(value=limit)
|
||
|
||
def wrapper(func: callable):
|
||
if inspect.iscoroutinefunction(func):
|
||
@wraps(func)
|
||
async def wrapped(*a, **kw):
|
||
acquired = sema.acquire(block=False)
|
||
if not acquired:
|
||
if throw:
|
||
raise ValueError(f'Reached concurrent limit! {func=}')
|
||
return
|
||
try:
|
||
return await func(*a, **kw)
|
||
finally:
|
||
sema.release()
|
||
|
||
return wrapped
|
||
else:
|
||
@wraps(func)
|
||
def wrapped(*a, **kw):
|
||
acquired = sema.acquire(block=False)
|
||
if not acquired:
|
||
if throw:
|
||
raise ValueError(f'Reached concurrent limit! {func=}')
|
||
return
|
||
try:
|
||
return func(*a, **kw)
|
||
finally:
|
||
sema.release()
|
||
|
||
return wrapped
|
||
|
||
return wrapper
|
||
|
||
|
||
def truncate_object_bytes(obj: Union[List[str], str, None, Generator], maximum_bytes: int) -> Union[List[str], str]:
|
||
"""
|
||
Shorten an object. This is useful when inserting something into a tsvector.
|
||
|
||
>>> i = ['foo', 'foo', 'foo', 'foo', 'foo', 'foo', 'foo', 'foo', 'foo', 'foo']
|
||
>>> truncate_object_bytes(i, 100)
|
||
['foo', 'foo', 'foo', 'foo', 'foo']
|
||
"""
|
||
if not obj:
|
||
# Can't decrease an empty object.
|
||
return obj
|
||
size = sys.getsizeof(obj)
|
||
if size < maximum_bytes:
|
||
return obj
|
||
index = -1 * round(len(obj) * .2)
|
||
obj = obj[:index or -1]
|
||
return truncate_object_bytes(obj, maximum_bytes)
|
||
|
||
|
||
def truncate_generator_bytes(gen: Generator, maximum_bytes: int) -> Generator:
|
||
"""Yield from a generator until the chunk sizes go over the maximum_bytes."""
|
||
total_size = 0
|
||
while total_size < maximum_bytes:
|
||
try:
|
||
chunk = next(gen)
|
||
total_size += sys.getsizeof(chunk)
|
||
yield chunk
|
||
except StopIteration:
|
||
return
|
||
|
||
|
||
BACKGROUND_TASKS = set()
|
||
|
||
|
||
def add_background_task(task: Task):
|
||
BACKGROUND_TASKS.add(task)
|
||
task.add_done_callback(BACKGROUND_TASKS.discard)
|
||
|
||
|
||
def background_task(coro) -> Task:
|
||
"""Convenience function which creates an asyncio task for the provided coroutine.
|
||
|
||
The task is stored in a global set of background tasks so the task will not be discarded by the garbage collector.
|
||
"""
|
||
# Get the coroutine name for better logging (used in task name for debugging stuck tasks)
|
||
coro_name = getattr(coro, '__qualname__', None) or getattr(coro, '__name__', None) or str(coro)
|
||
|
||
async def error_logger():
|
||
try:
|
||
await coro
|
||
except Exception as e:
|
||
logger_.error(f'background_task: error in {coro_name}:', exc_info=e)
|
||
|
||
task = asyncio.create_task(error_logger(), name=f'bg:{coro_name}')
|
||
add_background_task(task)
|
||
return task
|
||
|
||
|
||
async def cancel_background_tasks():
|
||
"""Cancels any background async tasks, if any."""
|
||
if BACKGROUND_TASKS:
|
||
logger_.warning(f'Canceling {len(BACKGROUND_TASKS)} background tasks')
|
||
for task in BACKGROUND_TASKS:
|
||
task.cancel()
|
||
await asyncio.gather(*BACKGROUND_TASKS)
|
||
|
||
|
||
async def await_background_tasks(timeout: int = 10):
|
||
"""Awaits all background tasks, used only for testing.
|
||
|
||
Processes the FileWorker queue concurrently with background tasks to avoid
|
||
deadlocks when background tasks wait for file worker jobs to complete.
|
||
|
||
Args:
|
||
timeout: Maximum seconds to wait for background tasks (default 15s)
|
||
"""
|
||
if not PYTEST:
|
||
raise RuntimeError('await_background_tasks should only be used for testing')
|
||
|
||
from wrolpi.files.worker import file_worker
|
||
import time
|
||
|
||
iteration = 0
|
||
max_iterations = 1000 # Safety limit to prevent infinite loops
|
||
start_time = time.time()
|
||
|
||
while BACKGROUND_TASKS or file_worker.public_queue.qsize() > 0 or not file_worker.private_queue.empty():
|
||
iteration += 1
|
||
elapsed = time.time() - start_time
|
||
|
||
if elapsed > timeout:
|
||
if __debug__ and logger_.isEnabledFor(TRACE_LEVEL):
|
||
task_names = [t.get_name() if hasattr(t, 'get_name') else str(t) for t in BACKGROUND_TASKS]
|
||
logger_.trace(
|
||
f'await_background_tasks: TIMEOUT after {elapsed:.1f}s with {len(BACKGROUND_TASKS)} tasks: {task_names}')
|
||
# Cancel stuck tasks and clear to prevent cascading failures
|
||
for task in list(BACKGROUND_TASKS):
|
||
if not task.done():
|
||
task.cancel()
|
||
BACKGROUND_TASKS.clear()
|
||
break
|
||
|
||
if iteration > max_iterations:
|
||
if __debug__ and logger_.isEnabledFor(TRACE_LEVEL):
|
||
task_names = [t.get_name() if hasattr(t, 'get_name') else str(t) for t in BACKGROUND_TASKS]
|
||
logger_.trace(
|
||
f'await_background_tasks: SAFETY LIMIT after {iteration} iterations with {len(BACKGROUND_TASKS)} tasks: {task_names}')
|
||
break
|
||
|
||
# Process file worker queue
|
||
file_worker.transfer_queue()
|
||
while not file_worker.private_queue.empty():
|
||
await file_worker.process_queue()
|
||
|
||
if not BACKGROUND_TASKS:
|
||
continue
|
||
|
||
# Wait for any background task to complete (with timeout so we can keep processing queue)
|
||
done, pending = await asyncio.wait(
|
||
BACKGROUND_TASKS.copy(),
|
||
timeout=0.1,
|
||
return_when=asyncio.FIRST_COMPLETED
|
||
)
|
||
|
||
for task in done:
|
||
await task # Re-await to propagate exceptions
|
||
BACKGROUND_TASKS.discard(task)
|
||
|
||
|
||
def get_warn_once(message: str, logger__: logging.Logger, level=logging.ERROR):
|
||
"""Create a function that will report an error only once.
|
||
|
||
This is used when a function will be called many times, but the error is not important."""
|
||
|
||
def warn_once(exception: Exception):
|
||
if PYTEST:
|
||
# No need to fill pytest logs with these warnings.
|
||
return
|
||
|
||
from wrolpi.api_utils import api_app
|
||
if not api_app.shared_ctx.warn_once.get(message):
|
||
logger__.log(level, message, exc_info=exception)
|
||
api_app.shared_ctx.warn_once.update({message: True})
|
||
|
||
return warn_once
|
||
|
||
|
||
def get_global_statistics():
|
||
from wrolpi.db import get_db_file
|
||
db_size = 0
|
||
db_file = get_db_file()
|
||
if db_file.is_file():
|
||
db_size = db_file.stat().st_size
|
||
# Include the WAL sidecar; its contents are part of the database.
|
||
wal_file = db_file.with_name(db_file.name + '-wal')
|
||
if wal_file.is_file():
|
||
db_size += wal_file.stat().st_size
|
||
|
||
return dict(
|
||
db_size=db_size,
|
||
)
|
||
|
||
|
||
REFRESH_TASKS: List[Task] = []
|
||
|
||
|
||
async def cancel_refresh_tasks():
|
||
"""Cancel all refresh tasks, if any."""
|
||
if REFRESH_TASKS:
|
||
logger_.warning(f'Canceling {len(REFRESH_TASKS)} refreshes')
|
||
for task in REFRESH_TASKS:
|
||
task.cancel()
|
||
await asyncio.gather(*REFRESH_TASKS)
|
||
|
||
|
||
def cancelable_wrapper(func: callable):
|
||
"""Wraps an async function so that it will be canceled by `cancel_refresh_tasks`."""
|
||
|
||
@wraps(func)
|
||
async def wrapped(*args, **kwargs):
|
||
if PYTEST:
|
||
return await func(*args, **kwargs)
|
||
|
||
task = background_task(func(*args, **kwargs))
|
||
REFRESH_TASKS.append(task)
|
||
|
||
return wrapped
|
||
|
||
|
||
WHITESPACE_SPLITTER = re.compile(r'\s+')
|
||
TAB = re.compile(r'\t+')
|
||
|
||
|
||
def strip_surrogates(text: Optional[str]) -> Optional[str]:
|
||
"""Remove unpaired UTF-16 surrogate characters (e.g. broken emoji halves from scraped HTML).
|
||
|
||
sqlite3 refuses to store them: "UnicodeEncodeError: surrogates not allowed"."""
|
||
if text is None:
|
||
return None
|
||
return text.encode('utf-8', errors='ignore').decode('utf-8')
|
||
|
||
|
||
def split_lines_by_length(text: str, max_line_length: int = 38) -> str:
|
||
"""
|
||
Break up any long lines along word boundaries.
|
||
|
||
Newlines are preserved, but whitspace (spaces/tabs) is replaced with one space.
|
||
|
||
@param max_line_length: The most amount of characters a line will contain. 38 is default for mobile vertical.
|
||
|
||
>>> split_lines_by_length('Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do')
|
||
'Lorem ipsum dolor sit amet,\nconsectetur adipiscing elit, sed do'
|
||
"""
|
||
if not text:
|
||
return text
|
||
|
||
text = TAB.sub(' ', text)
|
||
new_text = ''
|
||
for line in text.splitlines():
|
||
if len(line) > max_line_length:
|
||
# This line is too long, break it up.
|
||
new_line = ''
|
||
words = WHITESPACE_SPLITTER.split(line)
|
||
while words:
|
||
word = words.pop(0)
|
||
if len(possible_line := f'{new_line} {word}') <= max_line_length:
|
||
new_line = possible_line
|
||
else:
|
||
# New word makes the line too long, start a new line.
|
||
new_text += f'\n{new_line.strip()}'
|
||
new_line = word
|
||
new_text += f'\n{new_line}'
|
||
else:
|
||
new_text += f'\n{line}'
|
||
return new_text.lstrip('\n')
|
||
|
||
|
||
def resolve_generators(obj: Union[Dict, List]) -> Any:
|
||
"""Recursively find generators within an object/list, resolve them.
|
||
|
||
Returns a new object/list without any generators."""
|
||
if isinstance(obj, str):
|
||
return obj
|
||
elif isinstance(obj, Tuple):
|
||
return tuple(resolve_generators(i) for i in obj)
|
||
elif isinstance(obj, Set):
|
||
return {resolve_generators(i) for i in obj}
|
||
elif isinstance(obj, Dict):
|
||
return {resolve_generators(i): resolve_generators(j) for i, j in obj.items()}
|
||
elif isinstance(obj, range):
|
||
return list(obj)
|
||
elif isinstance(obj, GeneratorType) or isinstance(obj, Iterable):
|
||
return [resolve_generators(i) for i in obj]
|
||
return obj
|
||
|
||
|
||
def url_strip_host(url: str) -> str:
|
||
"""Return a relative URL without the host or scheme.
|
||
|
||
>>> url_strip_host('https://example.com/foo')
|
||
'/foo'
|
||
"""
|
||
url = urlparse(url)
|
||
url = urlunsplit(('', '', url.path, url.query, url.fragment))
|
||
return url or '/'
|
||
|
||
|
||
def get_html_soup(html: Union[bytes, str]) -> bs4.BeautifulSoup:
|
||
soup = BeautifulSoup(html, features='html.parser')
|
||
return soup
|
||
|
||
|
||
def extract_html_text(html: str) -> str:
|
||
soup = get_html_soup(html)
|
||
|
||
# kill all script and style elements
|
||
for script in soup(["script", "style"]):
|
||
script.extract() # rip it out
|
||
|
||
# Use body if it exists, otherwise use the entire soup
|
||
text = soup.body.get_text() if soup.body else soup.get_text()
|
||
|
||
# break into lines and remove leading and trailing space on each
|
||
lines = (line.strip() for line in text.splitlines())
|
||
# break multi-headlines into a line each
|
||
chunks_ = (phrase.strip() for line in lines for phrase in line.split(" "))
|
||
# drop blank lines
|
||
text = '\n'.join(chunk for chunk in chunks_ if chunk)
|
||
|
||
return text
|
||
|
||
|
||
def extract_headlines(entries: List[str], search_str: str) -> List[Tuple[str, float]]:
|
||
"""Extract headlines (highlighted snippets) and ranks of the `search_str` from the provided entries.
|
||
|
||
Entries which do not match get the start of their text with a rank of 0."""
|
||
from wrolpi import fts
|
||
|
||
return fts.headline_texts(entries, search_str, tokens=8)
|
||
|
||
|
||
async def search_other_estimates(tag_names: List[str]) -> dict:
|
||
"""Estimate other things that are Tagged."""
|
||
from sqlalchemy import func
|
||
from wrolpi.db import get_db_session
|
||
from wrolpi.collections.models import Collection
|
||
from wrolpi.tags import Tag
|
||
from modules.videos.models import Channel
|
||
|
||
if not tag_names:
|
||
return dict(
|
||
channel_count=0,
|
||
)
|
||
|
||
with get_db_session() as session:
|
||
# TODO handle multiple tags
|
||
channel_count = session.query(func.count(Channel.id)) \
|
||
.join(Collection, Collection.id == Channel.collection_id) \
|
||
.join(Tag, Tag.id == Collection.tag_id) \
|
||
.filter(Tag.name == tag_names[0]) \
|
||
.scalar()
|
||
|
||
others = dict(
|
||
channel_count=channel_count,
|
||
)
|
||
return others
|
||
|
||
|
||
def replace_file(file: pathlib.Path | str, contents: str | bytes, missing_ok: bool = False):
|
||
"""Rename `file` to a temporary path, write contents to `file`, then, delete temporary file only if successful.
|
||
The original file will be restored if any of these steps fail."""
|
||
file = pathlib.Path(file)
|
||
|
||
if missing_ok is False and not file.is_file():
|
||
raise FileNotFoundError(f'Cannot replace non-existent file: {str(file)}')
|
||
|
||
temporary_path = pathlib.Path(str(file) + '.tmp')
|
||
if temporary_path.exists():
|
||
raise RuntimeError(f'Cannot replace file, temporary path already exists! {str(temporary_path)}')
|
||
|
||
if not contents:
|
||
raise RuntimeError('Refusing to replace file with empty contents')
|
||
|
||
try:
|
||
if file.exists():
|
||
file.rename(temporary_path)
|
||
mode = 'wb' if isinstance(contents, bytes) else 'wt'
|
||
with file.open(mode) as fp:
|
||
fp.write(contents)
|
||
os.fsync(fp.fileno())
|
||
if file.exists() and temporary_path.exists():
|
||
# New contents have been written, delete the old file.
|
||
temporary_path.unlink()
|
||
except Exception as e:
|
||
logger.error(f'Failed to replace file: {str(file)}', exc_info=e)
|
||
if temporary_path.exists():
|
||
# Move original file back.
|
||
if file.exists():
|
||
file.unlink()
|
||
temporary_path.rename(file)
|
||
elif file.exists():
|
||
logger.error(f'Replacing file, but original file still exists: {str(file)}')
|
||
else:
|
||
logger.critical('Neither original, nor temporary file exist. I am so sorry.')
|
||
raise
|
||
finally:
|
||
if file.is_file() and file.stat().st_size > 0 and temporary_path.is_file():
|
||
temporary_path.unlink()
|
||
|
||
|
||
def format_json_file(file: pathlib.Path, indent: int = 2):
|
||
"""Reformat the contents of a JSON file to include indentation to increase readability."""
|
||
if not indent or indent <= 0:
|
||
raise RuntimeError('Ident must be greater than 0')
|
||
if not file or not file.is_file():
|
||
raise RuntimeError(f'File does not exist: {file}')
|
||
|
||
content = json.loads(file.read_text())
|
||
content = json.dumps(content, indent=indent, sort_keys=True)
|
||
replace_file(file, content)
|
||
|
||
|
||
def format_html_string(html: str | bytes) -> str:
|
||
if not html.strip():
|
||
raise RuntimeError('Refusing to format empty HTMl string.')
|
||
|
||
soup = BeautifulSoup(html, features='lxml')
|
||
return soup.prettify()
|
||
|
||
|
||
def format_html_file(file: pathlib.Path):
|
||
"""Reformat the contents of an HTML file to include indentation to increase readability."""
|
||
if not file or not file.is_file():
|
||
raise RuntimeError(f'File does not exist: {file}')
|
||
|
||
pretty = format_html_string(file.read_text())
|
||
replace_file(file, pretty)
|
||
|
||
|
||
def html_file_screenshot(path: pathlib.Path | str) -> bytes:
|
||
"""Return a PNG screenshot of the HTML file at `path`.
|
||
|
||
Relative resources (images, fonts) next to the file are rendered too."""
|
||
from wrolpi import cmd # Avoid circular import.
|
||
|
||
# Set Chromium to headless. Use a wide window size so that screenshot will be the "desktop" version of the page.
|
||
options = webdriver.ChromeOptions()
|
||
options.add_argument('headless')
|
||
options.add_argument('disable-gpu')
|
||
options.add_argument('window-size=1280x720')
|
||
|
||
# Use the system chromedriver; Selenium Manager downloads drivers from the internet, and has no binary for
|
||
# linux/aarch64 (Raspberry Pi).
|
||
service = webdriver.ChromeService(executable_path=str(cmd.CHROMEDRIVER_BIN)) if cmd.CHROMEDRIVER_BIN else None
|
||
|
||
with webdriver.Chrome(options=options, service=service) as driver:
|
||
driver.get(f'file://{path}')
|
||
driver.set_window_size(1280, 720)
|
||
screenshot = driver.get_screenshot_as_png()
|
||
return screenshot
|
||
|
||
|
||
def html_screenshot(html: bytes | str) -> bytes:
|
||
"""Return a PNG screenshot of the provided HTML."""
|
||
with tempfile.NamedTemporaryFile('wb', suffix='.html') as fh:
|
||
fh.write(html.encode() if isinstance(html, str) else html)
|
||
fh.flush()
|
||
return html_file_screenshot(fh.name)
|
||
|
||
|
||
def chain(iterable: Union[List, Tuple], length: int) -> List:
|
||
"""
|
||
Steps through the provided list and generates slices of the provided length.
|
||
|
||
>>> list(chain([1, 2, 3, 4], 2))
|
||
[[1, 2], [2, 3], [3, 4]]
|
||
|
||
>>> list(chain([1, 2, 3, 4], 3))
|
||
[[1, 2, 3], [2, 3, 4]]
|
||
"""
|
||
if not iterable:
|
||
return
|
||
|
||
position = 0
|
||
maximum = 1 + len(iterable) - length
|
||
yielded = False
|
||
while position < maximum:
|
||
tick = iterable[position:position + length]
|
||
yield tick
|
||
yielded = True
|
||
position += 1
|
||
if not yielded:
|
||
# Not enough items to iterate. Yield the original iterable.
|
||
yield iterable
|
||
|
||
|
||
def get_title_from_html(html: str, url: str = None) -> str:
|
||
"""
|
||
Try and get the title from the
|
||
"""
|
||
soup = get_html_soup(html)
|
||
try:
|
||
return soup.title.string.strip()
|
||
except Exception: # noqa
|
||
logger_.debug(f'Unable to extract title {url}')
|
||
|
||
|
||
@wrol_mode_check
|
||
def can_connect_to_server(hostname: str, port: int = 80) -> bool:
|
||
"""Return True only if `hostname` responds."""
|
||
# Thanks https://stackoverflow.com/questions/20913411/test-if-an-internet-connection-is-present-in-python
|
||
try:
|
||
# connect to the host -- tells us if the host is actually
|
||
# reachable
|
||
conn = socket.create_connection((hostname, port))
|
||
conn.close()
|
||
return True
|
||
except Exception as e:
|
||
logger_.debug('can_connect_to_server encountered error', exc_info=e)
|
||
return False
|
||
|
||
|
||
def is_valid_hex_color(hex_color: str) -> bool:
|
||
# Regular expression to match valid HTML hex color codes
|
||
hex_color_regex = re.compile(r'^#([A-F0-9]{6}|[A-F0-9]{3})$', re.IGNORECASE)
|
||
return bool(hex_color_regex.match(hex_color))
|
||
|
||
|
||
def is_hardlinked(path: pathlib.Path) -> bool:
|
||
"""Returns True only if `path` has more than one link."""
|
||
return path.stat().st_nlink > 1
|
||
|
||
|
||
def unique_by_predicate(
|
||
iterable: list | tuple | set | Generator,
|
||
predicate: callable = None,
|
||
) -> list | tuple | set:
|
||
"""Finds the first unique items in the provided iterable based on the predicate function. Order is preserved.
|
||
|
||
`predicate` defaults to the items themselves.
|
||
|
||
>>> unique_by_predicate([['apple', 'banana', 'pear', 'kiwi', 'plum', 'peach']], lambda i: len(i))
|
||
# [ 'apple', 'banana', 'pear' ]
|
||
>>> unique_by_predicate([Path('foo.txt'), Path('foo.mp4'), Path('bar.txt')], lambda i: i.stem)
|
||
# [ Path('foo.txt'), Path('bar.txt') ]
|
||
>>> unique_by_predicate([1, 1, 2, 3, 3, 3, 3, 4], None)
|
||
# [ 1, 2, 3, 4 ]
|
||
"""
|
||
predicate = predicate or (lambda i: i) # Return an iterable based off the object itself by default.
|
||
seen = set()
|
||
unique_items = [i for i in iterable if not ((key := predicate(i)) in seen or seen.add(key))]
|
||
if isinstance(iterable, Generator):
|
||
# Return a list when provided iterable is a Generator.
|
||
return unique_items
|
||
# Return the same type as the provided iterable
|
||
return iterable.__class__(unique_items)
|
||
|
||
|
||
def cached_multiprocessing_result(func: callable):
|
||
"""Simple multiprocessing results cacher which uses the args/kwargs of the wrapped function as the caching key.
|
||
|
||
@warning: Asumes the wrapped function is async.
|
||
"""
|
||
|
||
@functools.wraps(func)
|
||
async def wrapped(*args, **kwargs):
|
||
from wrolpi.api_utils import api_app
|
||
key = (func.__name__, *args, *tuple(kwargs.items()))
|
||
if result := api_app.shared_ctx.cache.get(key):
|
||
return result
|
||
|
||
result = await func(*args, **kwargs)
|
||
api_app.shared_ctx.cache[key] = result
|
||
return result
|
||
|
||
return wrapped
|
||
|
||
|
||
def create_empty_config_files() -> list[str]:
|
||
"""Creates config files only if they do not exist and will not conflict with the DB."""
|
||
from modules.videos.lib import get_channels_config, get_videos_downloader_config
|
||
from wrolpi.downloader import get_download_manager_config
|
||
from wrolpi.tags import get_tags_config
|
||
from wrolpi.db import get_db_session
|
||
|
||
created = [] # The names of the configs that are created.
|
||
|
||
if not get_wrolpi_config().get_file().is_file():
|
||
get_wrolpi_config().save(overwrite=True)
|
||
created.append(get_wrolpi_config().get_file().name)
|
||
|
||
with get_db_session() as session:
|
||
if not get_channels_config().get_file().is_file():
|
||
from modules.videos.models import Channel
|
||
if session.query(Channel).count() == 0:
|
||
# No Channels will be deleted, create the empty Channels config file.
|
||
get_channels_config().save(overwrite=True)
|
||
created.append(get_channels_config().get_file().name)
|
||
|
||
if not get_tags_config().get_file().is_file():
|
||
from wrolpi.tags import Tag
|
||
if session.query(Tag).count() == 0:
|
||
# No Tags will be deleted, create the empty Tags config file.
|
||
get_tags_config().save(overwrite=True)
|
||
created.append(get_tags_config().get_file().name)
|
||
|
||
if not get_download_manager_config().get_file().is_file():
|
||
from wrolpi.downloader import Download
|
||
if session.query(Download).count() == 0:
|
||
# No Downloads will be deleted, create the empty Download Manager config file.
|
||
get_download_manager_config().save(overwrite=True)
|
||
created.append(get_download_manager_config().get_file().name)
|
||
|
||
if not get_videos_downloader_config().get_file().is_file():
|
||
get_videos_downloader_config().save(overwrite=True)
|
||
created.append(get_videos_downloader_config().get_file().name)
|
||
|
||
from modules.archive.lib import get_archive_downloader_config
|
||
if not get_archive_downloader_config().get_file().is_file():
|
||
get_archive_downloader_config().save(overwrite=True)
|
||
created.append(get_archive_downloader_config().get_file().name)
|
||
|
||
# NOTE: Inventory is config-only (no DB) and manages its own per-inventory files under config/inventory/;
|
||
# it is seeded by migrate_legacy_inventory()/import on startup, not here.
|
||
|
||
return created
|
||
|
||
|
||
GPG_PUBLIC_KEY = Path(__file__).parent / 'roland@learningselfreliance.com.gpg'
|
||
|
||
|
||
async def verify_gpg_signature(data_path: Path, signature_path: Path) -> bool:
|
||
"""Verify a detached GPG signature against the shipped WROLPi public key.
|
||
|
||
Accepts file paths so that large files are never loaded into memory.
|
||
Uses a temporary GNUPGHOME to avoid polluting the system keyring.
|
||
Returns True if the signature is valid, False otherwise."""
|
||
from wrolpi.cmd import run_command
|
||
|
||
with tempfile.TemporaryDirectory() as gnupghome:
|
||
env = {**os.environ, 'GNUPGHOME': gnupghome}
|
||
|
||
# Import the trusted public key.
|
||
result = await run_command(
|
||
('gpg', '--batch', '--quiet', '--import', str(GPG_PUBLIC_KEY)),
|
||
timeout=10, env=env, log_command=False,
|
||
)
|
||
if result.return_code != 0:
|
||
logger.error(f'Failed to import GPG key: {result.stderr.decode()}')
|
||
return False
|
||
|
||
result = await run_command(
|
||
('gpg', '--batch', '--verify', str(signature_path), str(data_path)),
|
||
timeout=10, env=env, log_command=False,
|
||
)
|
||
return result.return_code == 0
|