#!/usr/bin/env python3 # -*- coding: UTF-8 -*- # Copyright (c) 2019 The ungoogled-chromium Authors. All rights reserved. # Use of this source code is governed by a BSD-style license that can be # found in the LICENSE file. """ Substitute domain names in the source tree with blockable strings. """ from pathlib import Path import argparse import collections import contextlib import io import os import stat import re import tarfile import tempfile import zlib from _extraction import extract_tar_file from _common import ENCODING, get_logger, add_common_params # Encodings to try on source tree files TREE_ENCODINGS = ('UTF-8', 'ISO-8859-1') # Constants for domain substitution cache _INDEX_LIST = 'cache_index.list' _INDEX_HASH_DELIMITER = '|' _ORIG_DIR = 'orig' # Constants for timestamp manipulation # Delta between all file timestamps in nanoseconds _TIMESTAMP_DELTA = 1 * 10**9 class DomainRegexList: """Representation of a domain_regex.list file""" _regex_pair_tuple = collections.namedtuple('DomainRegexPair', ('pattern', 'replacement')) # Constants for format: _PATTERN_REPLACE_DELIM = '#' def __init__(self, path): self._data = tuple(filter(len, path.read_text().splitlines())) # Cache of compiled regex pairs self._compiled_regex = None def _compile_regex(self, line): """Generates a regex pair tuple for the given line""" pattern, replacement = line.split(self._PATTERN_REPLACE_DELIM) return self._regex_pair_tuple(re.compile(pattern), replacement) @property def regex_pairs(self): """ Returns a tuple of compiled regex pairs """ if not self._compiled_regex: self._compiled_regex = tuple(map(self._compile_regex, self._data)) return self._compiled_regex @property def search_regex(self): """ Returns a single expression to search for domains """ return re.compile('|'.join( map(lambda x: x.split(self._PATTERN_REPLACE_DELIM, 1)[0], self._data))) # Private Methods def _substitute_path(path, regex_iter): """ Perform domain substitution on path and add it to the domain substitution cache. path is a pathlib.Path to the file to be domain substituted. regex_iter is an iterable of regular expression namedtuple like from config.DomainRegexList.regex_pairs() Returns a tuple of the CRC32 hash of the substituted raw content and the original raw content; None for both entries if no substitutions were made. Raises FileNotFoundError if path does not exist. Raises UnicodeDecodeError if path's contents cannot be decoded. """ if not os.access(path, os.W_OK): # If the patch cannot be written to, it cannot be opened for updating print(str(path) + " cannot be opened for writing! Adding write permission...") path.chmod(path.stat().st_mode | stat.S_IWUSR) with path.open('r+b') as input_file: original_content = input_file.read() if not original_content: return (None, None) content = None encoding = None for encoding in TREE_ENCODINGS: try: content = original_content.decode(encoding) break except UnicodeDecodeError: continue if not content: raise UnicodeDecodeError('Unable to decode with any encoding: %s' % path) file_subs = 0 for regex_pair in regex_iter: content, sub_count = regex_pair.pattern.subn(regex_pair.replacement, content) file_subs += sub_count if file_subs > 0: substituted_content = content.encode(encoding) input_file.seek(0) input_file.write(content.encode(encoding)) input_file.truncate() return (zlib.crc32(substituted_content), original_content) return (None, None) def _validate_file_index(index_file, resolved_tree, cache_index_files): """ Validation of file index and hashes against the source tree. Updates cache_index_files Returns True if the file index is valid; False otherwise """ all_hashes_valid = True crc32_regex = re.compile(r'^[a-zA-Z0-9]{8}$') for entry in index_file.read().decode(ENCODING).splitlines(): try: relative_path, file_hash = entry.split(_INDEX_HASH_DELIMITER) except ValueError as exc: get_logger().error('Could not split entry "%s": %s', entry, exc) continue if not relative_path or not file_hash: get_logger().error('Entry %s of domain substitution cache file index is not valid', _INDEX_HASH_DELIMITER.join((relative_path, file_hash))) all_hashes_valid = False continue if not crc32_regex.match(file_hash): get_logger().error('File index hash for %s does not appear to be a CRC32 hash', relative_path) all_hashes_valid = False continue if zlib.crc32((resolved_tree / relative_path).read_bytes()) != int(file_hash, 16): get_logger().error('Hashes do not match for: %s', relative_path) all_hashes_valid = False continue if relative_path in cache_index_files: get_logger().error('File %s shows up at least twice in the file index', relative_path) all_hashes_valid = False continue cache_index_files.add(relative_path) return all_hashes_valid @contextlib.contextmanager def _update_timestamp(path: os.PathLike, set_new: bool) -> None: """ Context manager to set the timestamp of the path to plus or minus a fixed delta, regardless of modifications within the context. if set_new is True, the delta is added. Otherwise, the delta is subtracted. """ stats = os.stat(path) if set_new: new_timestamp = (stats.st_atime_ns + _TIMESTAMP_DELTA, stats.st_mtime_ns + _TIMESTAMP_DELTA) else: new_timestamp = (stats.st_atime_ns - _TIMESTAMP_DELTA, stats.st_mtime_ns - _TIMESTAMP_DELTA) try: yield finally: os.utime(path, ns=new_timestamp) # Public Methods def apply_substitution(regex_path, files_path, source_tree, domainsub_cache): """ Substitute domains in source_tree with files and substitutions, and save the pre-domain substitution archive to presubdom_archive. regex_path is a pathlib.Path to domain_regex.list files_path is a pathlib.Path to domain_substitution.list source_tree is a pathlib.Path to the source tree. domainsub_cache is a pathlib.Path to the domain substitution cache. Raises NotADirectoryError if the patches directory is not a directory or does not exist Raises FileNotFoundError if the source tree or required directory does not exist. Raises FileExistsError if the domain substitution cache already exists. Raises ValueError if an entry in the domain substitution list contains the file index hash delimiter. """ if not source_tree.exists(): raise FileNotFoundError(source_tree) if not regex_path.exists(): raise FileNotFoundError(regex_path) if not files_path.exists(): raise FileNotFoundError(files_path) if domainsub_cache and domainsub_cache.exists(): raise FileExistsError(domainsub_cache) resolved_tree = source_tree.resolve() regex_pairs = DomainRegexList(regex_path).regex_pairs fileindex_content = io.BytesIO() with tarfile.open(str(domainsub_cache), 'w:%s' % domainsub_cache.suffix[1:], compresslevel=1) if domainsub_cache else open(os.devnull, 'w') as cache_tar: for relative_path in filter(len, files_path.read_text().splitlines()): if _INDEX_HASH_DELIMITER in relative_path: if domainsub_cache: # Cache tar will be incomplete; remove it for convenience cache_tar.close() domainsub_cache.unlink() raise ValueError( 'Path "%s" contains the file index hash delimiter "%s"' % relative_path, _INDEX_HASH_DELIMITER) path = resolved_tree / relative_path if not path.exists(): get_logger().warning('Skipping non-existant path: %s', path) continue if path.is_symlink(): get_logger().warning('Skipping path that has become a symlink: %s', path) continue with _update_timestamp(path, set_new=True): crc32_hash, orig_content = _substitute_path(path, regex_pairs) if crc32_hash is None: get_logger().info('Path has no substitutions: %s', relative_path) continue if domainsub_cache: fileindex_content.write('{}{}{:08x}\n'.format(relative_path, _INDEX_HASH_DELIMITER, crc32_hash).encode(ENCODING)) orig_tarinfo = tarfile.TarInfo(str(Path(_ORIG_DIR) / relative_path)) orig_tarinfo.size = len(orig_content) with io.BytesIO(orig_content) as orig_file: cache_tar.addfile(orig_tarinfo, orig_file) if domainsub_cache: fileindex_tarinfo = tarfile.TarInfo(_INDEX_LIST) fileindex_tarinfo.size = fileindex_content.tell() fileindex_content.seek(0) cache_tar.addfile(fileindex_tarinfo, fileindex_content) def revert_substitution(domainsub_cache, source_tree): """ Revert domain substitution on source_tree using the pre-domain substitution archive presubdom_archive. It first checks if the hashes of the substituted files match the hashes computed during the creation of the domain substitution cache, raising KeyError if there are any mismatches. Then, it proceeds to reverting files in the source_tree. domainsub_cache is removed only if all the files from the domain substitution cache were relocated to the source tree. domainsub_cache is a pathlib.Path to the domain substitution cache. source_tree is a pathlib.Path to the source tree. Raises KeyError if: * There is a hash mismatch while validating the cache * The cache's file index is corrupt or missing * The cache is corrupt or is not consistent with the file index Raises FileNotFoundError if the source tree or domain substitution cache do not exist. """ # This implementation trades disk space/wear for performance (unless a ramdisk is used # for the source tree) # Assumptions made for this process: # * The correct tar file was provided (so no huge amount of space is wasted) # * The tar file is well-behaved (e.g. no files extracted outside of destination path) # * Cache file index and cache contents are already consistent (i.e. no files exclusive to # one or the other) if not domainsub_cache: get_logger().error('Cache file must be specified.') if not domainsub_cache.exists(): raise FileNotFoundError(domainsub_cache) if not source_tree.exists(): raise FileNotFoundError(source_tree) resolved_tree = source_tree.resolve() cache_index_files = set() # All files in the file index with tempfile.TemporaryDirectory(prefix='domsubcache_files', dir=str(resolved_tree)) as tmp_extract_name: extract_path = Path(tmp_extract_name) get_logger().debug('Extracting domain substitution cache...') extract_tar_file(domainsub_cache, extract_path, None, False) # Validate source tree file hashes match get_logger().debug('Validating substituted files in source tree...') with (extract_path / _INDEX_LIST).open('rb') as index_file: #pylint: disable=no-member if not _validate_file_index(index_file, resolved_tree, cache_index_files): raise KeyError('Domain substitution cache file index is corrupt or hashes mismatch ' 'the source tree.') # Move original files over substituted ones get_logger().debug('Moving original files over substituted ones...') for relative_path in cache_index_files: with _update_timestamp(resolved_tree / relative_path, set_new=False): (extract_path / _ORIG_DIR / relative_path).replace(resolved_tree / relative_path) # Quick check for unused files in cache orig_has_unused = False for orig_path in (extract_path / _ORIG_DIR).rglob('*'): #pylint: disable=no-member if orig_path.is_file(): get_logger().warning('Unused file from cache: %s', orig_path) orig_has_unused = True if orig_has_unused: get_logger().warning('Cache contains unused files. Not removing.') else: domainsub_cache.unlink() def _callback(args): """CLI Callback""" if args.reverting: revert_substitution(args.cache, args.directory) else: apply_substitution(args.regex, args.files, args.directory, args.cache) def main(): """CLI Entrypoint""" parser = argparse.ArgumentParser() add_common_params(parser) parser.set_defaults(callback=_callback) subparsers = parser.add_subparsers(title='', dest='packaging') # apply apply_parser = subparsers.add_parser( 'apply', help='Apply domain substitution', description='Applies domain substitution and creates the domain substitution cache.') apply_parser.add_argument('-r', '--regex', type=Path, required=True, help='Path to domain_regex.list') apply_parser.add_argument('-f', '--files', type=Path, required=True, help='Path to domain_substitution.list') apply_parser.add_argument( '-c', '--cache', type=Path, help='The path to the domain substitution cache. The path must not already exist.') apply_parser.add_argument('directory', type=Path, help='The directory to apply domain substitution') apply_parser.set_defaults(reverting=False) # revert revert_parser = subparsers.add_parser( 'revert', help='Revert domain substitution', description='Reverts domain substitution based only on the domain substitution cache.') revert_parser.add_argument('directory', type=Path, help='The directory to reverse domain substitution') revert_parser.add_argument('-c', '--cache', type=Path, required=True, help=('The path to the domain substitution cache. ' 'The path must exist and will be removed if successful.')) revert_parser.set_defaults(reverting=True) args = parser.parse_args() args.callback(args) if __name__ == '__main__': main()