def _reset_service_updates(signature_type): service_updates = Hash( 'service-updates', get_client( host=config.core.redis.persistent.host, port=config.core.redis.persistent.port, private=False, )) for svc in service_updates.items(): if svc.lower() == signature_type.lower(): update_data = service_updates.get(svc) update_data['next_update'] = now_as_iso(120) update_data['previous_update'] = now_as_iso(-10**10) service_updates.set(svc, update_data) break
class ServiceUpdater(ThreadedCoreBase): def __init__(self, logger: logging.Logger = None, shutdown_timeout: float = None, config: Config = None, datastore: AssemblylineDatastore = None, redis: RedisType = None, redis_persist: RedisType = None, default_pattern=".*"): self.updater_type = os.environ['SERVICE_PATH'].split('.')[-1].lower() self.default_pattern = default_pattern if not logger: al_log.init_logging(f'updater.{self.updater_type}', log_level=os.environ.get('LOG_LEVEL', "WARNING")) logger = logging.getLogger(f'assemblyline.updater.{self.updater_type}') super().__init__(f'assemblyline.{SERVICE_NAME}_updater', logger=logger, shutdown_timeout=shutdown_timeout, config=config, datastore=datastore, redis=redis, redis_persist=redis_persist) self.update_data_hash = Hash(f'service-updates-{SERVICE_NAME}', self.redis_persist) self._update_dir = None self._update_tar = None self._time_keeper = None self._service: Optional[Service] = None self.event_sender = EventSender('changes.services', host=self.config.core.redis.nonpersistent.host, port=self.config.core.redis.nonpersistent.port) self.service_change_watcher = EventWatcher(self.redis, deserializer=ServiceChange.deserialize) self.service_change_watcher.register(f'changes.services.{SERVICE_NAME}', self._handle_service_change_event) self.signature_change_watcher = EventWatcher(self.redis, deserializer=SignatureChange.deserialize) self.signature_change_watcher.register(f'changes.signatures.{SERVICE_NAME.lower()}', self._handle_signature_change_event) # A event flag that gets set when an update should be run for # reasons other than it being the regular interval (eg, change in signatures) self.source_update_flag = threading.Event() self.local_update_flag = threading.Event() self.local_update_start = threading.Event() # Load threads self._internal_server = None self.expected_threads = { 'Sync Service Settings': self._sync_settings, 'Outward HTTP Server': self._run_http, 'Internal HTTP Server': self._run_internal_http, 'Run source updates': self._run_source_updates, 'Run local updates': self._run_local_updates, } # Only used by updater with 'generates_signatures: false' self.latest_updates_dir = os.path.join(UPDATER_DIR, 'latest_updates') if not os.path.exists(self.latest_updates_dir): os.makedirs(self.latest_updates_dir) def trigger_update(self): self.source_update_flag.set() def update_directory(self): return self._update_dir def update_tar(self): return self._update_tar def get_active_config_hash(self) -> int: return self.update_data_hash.get(CONFIG_HASH_KEY) or 0 def set_active_config_hash(self, config_hash: int): self.update_data_hash.set(CONFIG_HASH_KEY, config_hash) def get_source_update_time(self) -> float: return self.update_data_hash.get(SOURCE_UPDATE_TIME_KEY) or 0 def set_source_update_time(self, update_time: float): self.update_data_hash.set(SOURCE_UPDATE_TIME_KEY, update_time) def get_source_extra(self) -> dict[str, Any]: return self.update_data_hash.get(SOURCE_EXTRA_KEY) or {} def set_source_extra(self, extra_data: dict[str, Any]): self.update_data_hash.set(SOURCE_EXTRA_KEY, extra_data) def get_local_update_time(self) -> float: if self._time_keeper: return os.path.getctime(self._time_keeper) return 0 def status(self): return { 'local_update_time': self.get_local_update_time(), 'download_available': self._update_dir is not None, '_directory': self._update_dir, '_tar': self._update_tar, } def stop(self): super().stop() self.signature_change_watcher.stop() self.service_change_watcher.stop() self.source_update_flag.set() self.local_update_flag.set() self.local_update_start.set() if self._internal_server: self._internal_server.shutdown() def try_run(self): self.signature_change_watcher.start() self.service_change_watcher.start() self.maintain_threads(self.expected_threads) def _run_internal_http(self): """run backend insecure http server A small inprocess server to syncronize info between gunicorn and the updater daemon. This HTTP server is not safe for exposing externally, but fine for IPC. """ them = self class Handler(BaseHTTPRequestHandler): def do_GET(self): self.send_response(200) self.send_header("Content-type", "application/json") self.end_headers() self.wfile.write(json.dumps(them.status()).encode()) def log_error(self, format: str, *args: Any): them.log.info(format % args) def log_message(self, format: str, *args: Any): them.log.debug(format % args) self._internal_server = ThreadingHTTPServer(('0.0.0.0', 9999), Handler) self._internal_server.serve_forever() def _run_http(self): # Start a server for our http interface in a separate process my_env = os.environ.copy() proc = subprocess.Popen(["gunicorn", "assemblyline_v4_service.updater.app:app", "--config=python:assemblyline_v4_service.updater.gunicorn_config"], env=my_env) while self.sleep(1): if proc.poll() is not None: break # If we have left the loop and the process is still alive, stop it. if proc.poll() is not None: proc.terminate() proc.wait() @staticmethod def config_hash(service: Service) -> int: if service is None: return 0 return hash(json.dumps(service.update_config.as_primitives())) def _handle_signature_change_event(self, data: SignatureChange): self.local_update_flag.set() def _handle_service_change_event(self, data: ServiceChange): if data.operation == Operation.Modified: self._pull_settings() def _sync_settings(self): # Download the service object from datastore self._service = self.datastore.get_service_with_delta(SERVICE_NAME) while self.sleep(SERVICE_PULL_INTERVAL): self._pull_settings() def _pull_settings(self): # Download the service object from datastore self._service = self.datastore.get_service_with_delta(SERVICE_NAME) # If the update configuration for the service has changed, trigger an update if self.config_hash(self._service) != self.get_active_config_hash(): self.source_update_flag.set() def do_local_update(self) -> None: old_update_time = self.get_local_update_time() if not os.path.exists(UPDATER_DIR): os.makedirs(UPDATER_DIR) _, time_keeper = tempfile.mkstemp(prefix="time_keeper_", dir=UPDATER_DIR) if self._service.update_config.generates_signatures: output_directory = tempfile.mkdtemp(prefix="update_dir_", dir=UPDATER_DIR) self.log.info("Setup service account.") username = self.ensure_service_account() self.log.info("Create temporary API key.") with temporary_api_key(self.datastore, username) as api_key: self.log.info(f"Connecting to Assemblyline API: {UI_SERVER}") al_client = get_client(UI_SERVER, apikey=(username, api_key), verify=False) # Check if new signatures have been added self.log.info("Check for new signatures.") if al_client.signature.update_available( since=epoch_to_iso(old_update_time) or '', sig_type=self.updater_type)['update_available']: self.log.info("An update is available for download from the datastore") self.log.debug(f"{self.updater_type} update available since {epoch_to_iso(old_update_time) or ''}") extracted_zip = False attempt = 0 # Sometimes a zip file isn't always returned, will affect service's use of signature source. Patience.. while not extracted_zip and attempt < 5: temp_zip_file = os.path.join(output_directory, 'temp.zip') al_client.signature.download( output=temp_zip_file, query=f"type:{self.updater_type} AND (status:NOISY OR status:DEPLOYED)") self.log.debug(f"Downloading update to {temp_zip_file}") if os.path.exists(temp_zip_file) and os.path.getsize(temp_zip_file) > 0: self.log.debug(f"File type ({os.path.getsize(temp_zip_file)}B): {zip_ident(temp_zip_file, 'unknown')}") try: with ZipFile(temp_zip_file, 'r') as zip_f: zip_f.extractall(output_directory) extracted_zip = True self.log.info("Zip extracted.") except BadZipFile: attempt += 1 self.log.warning(f"[{attempt}/5] Bad zip. Trying again after 30s...") time.sleep(30) except Exception as e: self.log.error(f'Problem while extracting signatures to disk: {e}') break os.remove(temp_zip_file) if extracted_zip: self.log.info("New ruleset successfully downloaded and ready to use") self.serve_directory(output_directory, time_keeper) else: self.log.error("Signatures aren't saved to disk.") shutil.rmtree(output_directory, ignore_errors=True) if os.path.exists(time_keeper): os.unlink(time_keeper) else: self.log.info("No signature updates available.") shutil.rmtree(output_directory, ignore_errors=True) if os.path.exists(time_keeper): os.unlink(time_keeper) else: output_directory = self.prepare_output_directory() self.serve_directory(output_directory, time_keeper) def do_source_update(self, service: Service) -> None: self.log.info(f"Connecting to Assemblyline API: {UI_SERVER}...") run_time = time.time() username = self.ensure_service_account() with temporary_api_key(self.datastore, username) as api_key: with tempfile.TemporaryDirectory() as update_dir: al_client = get_client(UI_SERVER, apikey=(username, api_key), verify=False) old_update_time = self.get_source_update_time() self.log.info("Connected!") # Parse updater configuration previous_hashes: dict[str, dict[str, str]] = self.get_source_extra() sources: dict[str, UpdateSource] = {_s['name']: _s for _s in service.update_config.sources} files_sha256: dict[str, dict[str, str]] = {} # Go through each source and download file for source_name, source_obj in sources.items(): source = source_obj.as_primitives() uri: str = source['uri'] default_classification = source.get('default_classification', classification.UNRESTRICTED) try: # Pull sources from external locations (method depends on the URL) files = git_clone_repo(source, old_update_time, self.default_pattern, self.log, update_dir) \ if uri.endswith('.git') else url_download(source, old_update_time, self.log, update_dir) # Add to collection of sources for caching purposes self.log.info(f"Found new {self.updater_type} rule files to process for {source_name}!") validated_files = list() for file, sha256 in files: files_sha256.setdefault(source_name, {}) if previous_hashes.get(source_name, {}).get(file, None) != sha256 and self.is_valid(file): files_sha256[source_name][file] = sha256 validated_files.append((file, sha256)) # Import into Assemblyline self.import_update(validated_files, al_client, source_name, default_classification) except SkipSource: # This source hasn't changed, no need to re-import into Assemblyline self.log.info(f'No new {self.updater_type} rule files to process for {source_name}') if source_name in previous_hashes: files_sha256[source_name] = previous_hashes[source_name] continue self.set_source_update_time(run_time) self.set_source_extra(files_sha256) self.set_active_config_hash(self.config_hash(service)) self.local_update_flag.set() # Define to determine if file is a valid signature file def is_valid(self, file_path) -> bool: return True # Define how your source update gets imported into Assemblyline def import_update(self, files_sha256: List[Tuple[str, str]], client: Client, source_name: str, default_classification=None): raise NotImplementedError() # Define how to prepare the output directory before being served, must return the path of the directory to serve. def prepare_output_directory(self) -> str: output_directory = tempfile.mkdtemp() shutil.copytree(self.latest_updates_dir, output_directory, dirs_exist_ok=True) return output_directory def _run_source_updates(self): # Wait until basic data is loaded while self._service is None and self.sleep(1): pass if not self._service: return self.log.info("Service info loaded") try: self.log.info("Checking for in cluster update cache") self.do_local_update() self._service_stage_hash.set(SERVICE_NAME, ServiceStage.Running) self.event_sender.send(SERVICE_NAME, {'operation': Operation.Modified, 'name': SERVICE_NAME}) except Exception: self.log.exception('An error occurred loading cached update files. Continuing.') self.local_update_start.set() # Go into a loop running the update whenever triggered or its time to while self.running: # Stringify and hash the the current update configuration service = self._service update_interval = service.update_config.update_interval_seconds # Is it time to update yet? if time.time() - self.get_source_update_time() < update_interval and not self.source_update_flag.is_set(): self.source_update_flag.wait(60) continue if not self.running: return # With temp directory self.source_update_flag.clear() self.log.info('Calling update function...') # Run update function # noinspection PyBroadException try: self.do_source_update(service=service) except Exception: self.log.exception('An error occurred running the update. Will retry...') self.source_update_flag.set() self.sleep(60) continue def serve_directory(self, new_directory: str, new_time: str): self.log.info("Update finished with new data.") new_tar = '' try: # Tar update directory _, new_tar = tempfile.mkstemp(prefix="signatures_", dir=UPDATER_DIR, suffix='.tar.bz2') tar_handle = tarfile.open(new_tar, 'w:bz2') tar_handle.add(new_directory, '/') tar_handle.close() # swap update directory with old one self._update_dir, new_directory = new_directory, self._update_dir self._update_tar, new_tar = new_tar, self._update_tar self._time_keeper, new_time = new_time, self._time_keeper self.log.info(f"Now serving: {self._update_dir} and {self._update_tar} ({self.get_local_update_time()})") finally: if new_tar and os.path.exists(new_tar): self.log.info(f"Remove old tar file: {new_tar}") time.sleep(3) os.unlink(new_tar) if new_directory and os.path.exists(new_directory): self.log.info(f"Remove old directory: {new_directory}") shutil.rmtree(new_directory, ignore_errors=True) if new_time and os.path.exists(new_time): self.log.info(f"Remove old time keeper file: {new_time}") os.unlink(new_time) def _run_local_updates(self): # Wait until basic data is loaded while self._service is None and self.sleep(1): pass if not self._service: return self.local_update_start.wait() # Go into a loop running the update whenever triggered or its time to while self.running: # Is it time to update yet? if not self.local_update_flag.is_set(): self.local_update_flag.wait(60) continue if not self.running: return self.local_update_flag.clear() # With temp directory self.log.info('Updating local files...') # Run update function # noinspection PyBroadException try: self.do_local_update() if self._service_stage_hash.get(SERVICE_NAME) == ServiceStage.Update: self._service_stage_hash.set(SERVICE_NAME, ServiceStage.Running) self.event_sender.send(SERVICE_NAME, {'operation': Operation.Modified, 'name': SERVICE_NAME}) except Exception: self.log.exception('An error occurred finding new local files. Will retry...') self.local_update_flag.set() self.sleep(60) continue def ensure_service_account(self): """Check that the update service account exists, if it doesn't, create it.""" uname = 'update_service_account' if self.datastore.user.get_if_exists(uname): return uname user_data = User({ "agrees_with_tos": "NOW", "classification": "RESTRICTED", "name": "Update Account", "password": get_password_hash(''.join(random.choices(string.ascii_letters, k=20))), "uname": uname, "type": ["signature_importer"] }) self.datastore.user.save(uname, user_data) self.datastore.user_settings.save(uname, UserSettings()) return uname
class ServiceUpdater(CoreBase): def __init__(self, redis_persist=None, redis=None, logger=None, datastore=None): super().__init__('assemblyline.service.updater', logger=logger, datastore=datastore, redis_persist=redis_persist, redis=redis) if not FILE_UPDATE_DIRECTORY: raise RuntimeError( "The updater process must be run within the orchestration environment, " "the update volume must be mounted, and the path to the volume must be " "set in the environment variable FILE_UPDATE_DIRECTORY. Setting " "FILE_UPDATE_DIRECTORY directly may be done for testing.") # The directory where we want working temporary directories to be created. # Building our temporary directories in the persistent update volume may # have some performance down sides, but may help us run into fewer docker FS overlay # cleanup issues. Try to flush it out every time we start. This service should # be a singleton anyway. self.temporary_directory = os.path.join(FILE_UPDATE_DIRECTORY, '.tmp') shutil.rmtree(self.temporary_directory, ignore_errors=True) os.makedirs(self.temporary_directory) self.container_update = Hash('container-update', self.redis_persist) self.services = Hash('service-updates', self.redis_persist) self.latest_service_tags = Hash('service-tags', self.redis_persist) self.running_updates: Dict[str, Thread] = {} # Prepare a single threaded scheduler self.scheduler = sched.scheduler() # if 'KUBERNETES_SERVICE_HOST' in os.environ and NAMESPACE: self.controller = KubernetesUpdateInterface( prefix='alsvc_', namespace=NAMESPACE, priority_class='al-core-priority') else: self.controller = DockerUpdateInterface() def sync_services(self): """Download the service list and make sure our settings are up to date""" self.scheduler.enter(SERVICE_SYNC_INTERVAL, 0, self.sync_services) existing_services = (set(self.services.keys()) | set(self.container_update.keys()) | set(self.latest_service_tags.keys())) discovered_services = [] # Get all the service data for service in self.datastore.list_all_services(full=True): discovered_services.append(service.name) # Ensure that any disabled services are not being updated if not service.enabled and self.services.exists(service.name): self.log.info(f"Service updates disabled for {service.name}") self.services.pop(service.name) if not service.enabled: continue # Ensure that any enabled services with an update config are being updated stage = self.get_service_stage(service.name) record = self.services.get(service.name) if stage in UPDATE_STAGES and service.update_config: # Stringify and hash the the current update configuration config_hash = hash( json.dumps(service.update_config.as_primitives())) # If we can update, but there is no record, create one if not record: self.log.info( f"Service updates enabled for {service.name}") self.services.add( service.name, dict( next_update=now_as_iso(), previous_update=now_as_iso(-10**10), config_hash=config_hash, sha256=None, )) else: # If there is a record, check that its configuration hash is still good # If an update is in progress, it may overwrite this, but we will just come back # and reapply this again in the iteration after that if record.get('config_hash', None) != config_hash: record['next_update'] = now_as_iso() record['config_hash'] = config_hash self.services.set(service.name, record) if stage == ServiceStage.Update: if (record and record.get('sha256', None) is not None) or not service.update_config: self._service_stage_hash.set(service.name, ServiceStage.Running) # Remove services we have locally or in redis that have been deleted from the database for stray_service in existing_services - set(discovered_services): self.log.info(f"Service updates disabled for {stray_service}") self.services.pop(stray_service) self._service_stage_hash.pop(stray_service) self.container_update.pop(stray_service) self.latest_service_tags.pop(stray_service) def container_updates(self): """Go through the list of services and check what are the latest tags for it""" self.scheduler.enter(UPDATE_CHECK_INTERVAL, 0, self.container_updates) for service_name, update_data in self.container_update.items().items(): self.log.info( f"Service {service_name} is being updated to version {update_data['latest_tag']}..." ) # Load authentication params username = None password = None auth = update_data['auth'] or {} if auth: username = auth.get('username', None) password = auth.get('password', None) try: self.controller.launch( name=service_name, docker_config=DockerConfig( dict(allow_internet_access=True, registry_username=username, registry_password=password, cpu_cores=1, environment=[], image=update_data['image'], ports=[])), mounts=[], env={ "SERVICE_TAG": update_data['latest_tag'], "SERVICE_API_HOST": os.environ.get('SERVICE_API_HOST', "http://al_service_server:5003"), "REGISTER_ONLY": 'true' }, network='al_registration', blocking=True) latest_tag = update_data['latest_tag'].replace('stable', '') service_key = f"{service_name}_{latest_tag}" if self.datastore.service.get_if_exists(service_key): operations = [(self.datastore.service_delta.UPDATE_SET, 'version', latest_tag)] if self.datastore.service_delta.update( service_name, operations): # Update completed, cleanup self.log.info( f"Service {service_name} update successful!") else: self.log.error( f"Service {service_name} has failed to update because it cannot set " f"{latest_tag} as the new version. Update procedure cancelled..." ) else: self.log.error( f"Service {service_name} has failed to update because resulting " f"service key ({service_key}) does not exist. Update procedure cancelled..." ) except Exception as e: self.log.error( f"Service {service_name} has failed to update. Update procedure cancelled... [{str(e)}]" ) self.container_update.pop(service_name) def container_versions(self): """Go through the list of services and check what are the latest tags for it""" self.scheduler.enter(CONTAINER_CHECK_INTERVAL, 0, self.container_versions) for service in self.datastore.list_all_services(full=True): if not service.enabled: continue image_name, tag_name, auth = get_latest_tag_for_service( service, self.config, self.log) self.latest_service_tags.set( service.name, { 'auth': auth, 'image': image_name, service.update_channel: tag_name }) def try_run(self): """Run the scheduler loop until told to stop.""" # Do an initial call to the main methods, who will then be registered with the scheduler self.sync_services() self.update_services() self.container_versions() self.container_updates() self.heartbeat() # Run as long as we need to while self.running: delay = self.scheduler.run(False) time.sleep(min(delay, 0.1)) def heartbeat(self): """Periodically touch a file on disk. Since tasks are run serially, the delay between touches will be the maximum of HEARTBEAT_INTERVAL and the longest running task. """ if self.config.logging.heartbeat_file: self.scheduler.enter(HEARTBEAT_INTERVAL, 0, self.heartbeat) super().heartbeat() def update_services(self): """Check if we need to update any services. Spin off a thread to actually perform any updates. Don't allow multiple threads per service. """ self.scheduler.enter(UPDATE_CHECK_INTERVAL, 0, self.update_services) # Check for finished update threads self.running_updates = { name: thread for name, thread in self.running_updates.items() if thread.is_alive() } # Check if its time to try to update the service for service_name, data in self.services.items().items(): if data['next_update'] <= now_as_iso( ) and service_name not in self.running_updates: self.log.info(f"Time to update {service_name}") self.running_updates[service_name] = Thread( target=self.run_update, kwargs=dict(service_name=service_name)) self.running_updates[service_name].start() def run_update(self, service_name): """Common setup and tear down for all update types.""" # noinspection PyBroadException try: # Check for new update with service specified update method service = self.datastore.get_service_with_delta(service_name) update_method = service.update_config.method update_data = self.services.get(service_name) update_hash = None try: # Actually run the update method if update_method == 'run': update_hash = self.do_file_update( service=service, previous_hash=update_data['sha256'], previous_update=update_data['previous_update']) elif update_method == 'build': update_hash = self.do_build_update() # If we have performed an update, write that data if update_hash is not None and update_hash != update_data[ 'sha256']: update_data['sha256'] = update_hash update_data['previous_update'] = now_as_iso() else: update_hash = None finally: # Update the next service update check time, don't update the config_hash, # as we don't want to disrupt being re-run if our config has changed during this run update_data['next_update'] = now_as_iso( service.update_config.update_interval_seconds) self.services.set(service_name, update_data) if update_hash: self.log.info( f"New update applied for {service_name}. Restarting service." ) self.controller.restart(service_name=service_name) except BaseException: self.log.exception( "An error occurred while running an update for: " + service_name) def do_build_update(self): """Update a service by building a new container to run.""" raise NotImplementedError() def do_file_update(self, service, previous_hash, previous_update): """Update a service by running a container to get new files.""" temp_directory = tempfile.mkdtemp(dir=self.temporary_directory) chmod(temp_directory, 0o777) input_directory = os.path.join(temp_directory, 'input_directory') output_directory = os.path.join(temp_directory, 'output_directory') service_dir = os.path.join(FILE_UPDATE_DIRECTORY, service.name) image_variables = defaultdict(str) image_variables.update(self.config.services.image_variables) try: # Use chmod directly to avoid effects of umask os.makedirs(input_directory) chmod(input_directory, 0o755) os.makedirs(output_directory) chmod(output_directory, 0o777) username = self.ensure_service_account() with temporary_api_key(self.datastore, username) as api_key: # Write out the parameters we want to pass to the update container with open(os.path.join(input_directory, 'config.yaml'), 'w') as fh: yaml.safe_dump( { 'previous_update': previous_update, 'previous_hash': previous_hash, 'sources': [ x.as_primitives() for x in service.update_config.sources ], 'api_user': username, 'api_key': api_key, 'ui_server': UI_SERVER }, fh) # Run the update container run_options = service.update_config.run_options run_options.image = string.Template( run_options.image).safe_substitute(image_variables) self.controller.launch( name=service.name, docker_config=run_options, mounts=[ { 'volume': FILE_UPDATE_VOLUME, 'source_path': os.path.relpath(temp_directory, start=FILE_UPDATE_DIRECTORY), 'dest_path': '/mount/' }, ], env={ 'UPDATE_CONFIGURATION_PATH': '/mount/input_directory/config.yaml', 'UPDATE_OUTPUT_PATH': '/mount/output_directory/' }, network=f'service-net-{service.name}', blocking=True, ) # Read out the results from the output container results_meta_file = os.path.join(output_directory, 'response.yaml') if not os.path.exists(results_meta_file) or not os.path.isfile( results_meta_file): self.log.warning( f"Update produced no output for {service.name}") return None with open(results_meta_file) as rf: results_meta = yaml.safe_load(rf) update_hash = results_meta.get('hash', None) # Erase the results meta file os.unlink(results_meta_file) # Get a timestamp for now, and switch it to basic format representation of time # Still valid iso 8601, and : is sometimes a restricted character timestamp = now_as_iso().replace(":", "") # FILE_UPDATE_DIRECTORY/{service_name} is the directory mounted to the service, # the service sees multiple directories in that directory, each with a timestamp destination_dir = os.path.join(service_dir, service.name + '_' + timestamp) shutil.move(output_directory, destination_dir) # Remove older update files, due to the naming scheme, older ones will sort first lexically existing_folders = [] for folder_name in os.listdir(service_dir): folder_path = os.path.join(service_dir, folder_name) if os.path.isdir(folder_path) and folder_name.startswith( service.name): existing_folders.append(folder_name) existing_folders.sort() self.log.info( f'There are {len(existing_folders)} update folders for {service.name} in cache.' ) if len(existing_folders) > UPDATE_FOLDER_LIMIT: extra_count = len(existing_folders) - UPDATE_FOLDER_LIMIT self.log.info( f'We will only keep {UPDATE_FOLDER_LIMIT} updates, deleting {extra_count}.' ) for extra_folder in existing_folders[:extra_count]: # noinspection PyBroadException try: shutil.rmtree( os.path.join(service_dir, extra_folder)) except Exception: self.log.exception( 'Failed to delete update folder') return update_hash finally: # If the working directory is still there for any reason erase it shutil.rmtree(temp_directory, ignore_errors=True) def ensure_service_account(self): """Check that the update service account exists, if it doesn't, create it.""" uname = 'update_service_account' if self.datastore.user.get_if_exists(uname): return uname user_data = User({ "agrees_with_tos": "NOW", "classification": "RESTRICTED", "name": "Update Account", "password": get_password_hash(''.join( random.choices(string.ascii_letters, k=20))), "uname": uname, "type": ["signature_importer"] }) self.datastore.user.save(uname, user_data) self.datastore.user_settings.save(uname, UserSettings()) return uname