import gzip import logging import os import shlex import shutil import signal import subprocess from django.conf import settings class SingeFileError(Exception): pass logger = logging.getLogger(__name__) def create_snapshot(url: str, filepath: str): singlefile_path = settings.LD_SINGLEFILE_PATH # parse options to list of arguments ublock_options = shlex.split(settings.LD_SINGLEFILE_UBLOCK_OPTIONS) custom_options = shlex.split(settings.LD_SINGLEFILE_OPTIONS) temp_filepath = filepath + ".tmp" # concat lists args = [singlefile_path] + ublock_options + custom_options + [url, temp_filepath] try: # Use start_new_session=True to create a new process group process = subprocess.Popen(args, start_new_session=True) process.wait(timeout=settings.LD_SINGLEFILE_TIMEOUT_SEC) # check if the file was created if not os.path.exists(temp_filepath): raise SingeFileError("Failed to create snapshot") with open(temp_filepath, "rb") as raw_file, gzip.open( filepath, "wb" ) as gz_file: shutil.copyfileobj(raw_file, gz_file) os.remove(temp_filepath) except subprocess.TimeoutExpired: # First try to terminate properly try: logger.error( "Timeout expired while creating snapshot. Terminating process..." ) process.terminate() process.wait(timeout=20) raise SingeFileError("Timeout expired while creating snapshot") except subprocess.TimeoutExpired: # Kill the whole process group, which should also clean up any chromium # processes spawned by single-file logger.error("Timeout expired while terminating. Killing process...") os.killpg(os.getpgid(process.pid), signal.SIGTERM) raise SingeFileError("Timeout expired while creating snapshot") except subprocess.CalledProcessError as error: raise SingeFileError(f"Failed to create snapshot: {error.stderr}")