diff --git a/assemblyline_result_sample_service/service_manifest.yml b/assemblyline_result_sample_service/service_manifest.yml index add7df2e..807c81de 100644 --- a/assemblyline_result_sample_service/service_manifest.yml +++ b/assemblyline_result_sample_service/service_manifest.yml @@ -104,17 +104,17 @@ docker_config: ram_mb: 256 # Dependencies configuration block -dependencies: - updates: - container: - allow_internet_access: true - command: ["python", "-m", "update_server"] - image: ${REGISTRY}cccs/assemblyline-service-resultsample:$SERVICE_TAG - ports: ["5003"] - # Upper-bound, adjust based on sources - # cpu_cores: 2 - # ram_mb: 4096 - run_as_core: True +# dependencies: +# updates: +# container: +# allow_internet_access: true +# command: ["python", "-m", "update_server"] +# image: ${REGISTRY}cccs/assemblyline-service-resultsample:$SERVICE_TAG +# ports: ["5003"] +# # Upper-bound, adjust based on sources +# # cpu_cores: 2 +# # ram_mb: 4096 +# run_as_core: True # Update configuration block update_config: diff --git a/assemblyline_result_sample_service/update_server.py b/assemblyline_result_sample_service/update_server.py index 3e13a5ea..30721019 100644 --- a/assemblyline_result_sample_service/update_server.py +++ b/assemblyline_result_sample_service/update_server.py @@ -3,15 +3,13 @@ import tempfile import time -import certifi -import requests from assemblyline.common import forge from assemblyline.common import log as al_log -from assemblyline.common.isotime import iso_to_epoch from assemblyline.odm.models.service import Service, UpdateSource from assemblyline_client import get_client from assemblyline_v4_service.updater.updater import ServiceUpdater, temporary_api_key +from assemblyline_v4_service.updater.helper import url_download, SkipSource al_log.init_logging('updater.sample') classification = forge.get_classification() @@ -28,94 +26,6 @@ HASH_LEN = 1000 -class SkipSource(RuntimeError): - pass - - -def add_cacert(cert: str): - # Add certificate to requests - cafile = certifi.where() - with open(cafile, 'a') as ca_editor: - ca_editor.write(f"\n{cert}") - - -def url_download(source, target_path, logger, previous_update=None): - uri = source['uri'] - username = source.get('username', None) - password = source.get('password', None) - ca_cert = source.get('ca_cert', None) - ignore_ssl_errors = source.get('ssl_ignore_errors', False) - auth = (username, password) if username and password else None - - proxy = source.get('proxy', None) - headers = source.get('headers', None) - - logger.info(f"This source is configured to {'ignore SSL errors' if ignore_ssl_errors else 'verify SSL'}.") - if ca_cert: - logger.info("A CA certificate has been provided with this source.") - add_cacert(ca_cert) - - # Create a requests session - session = requests.Session() - session.verify = not ignore_ssl_errors - - # Let https requests go through proxy - if proxy: - os.environ['https_proxy'] = proxy - - try: - if isinstance(previous_update, str): - previous_update = iso_to_epoch(previous_update) - - # Check the response header for the last modified date - response = session.head(uri, auth=auth, headers=headers) - last_modified = response.headers.get('Last-Modified', None) - if last_modified: - # Convert the last modified time to epoch - last_modified = time.mktime(time.strptime(last_modified, "%a, %d %b %Y %H:%M:%S %Z")) - - # Compare the last modified time with the last updated time - if previous_update and last_modified <= previous_update: - # File has not been modified since last update, do nothing - logger.info("The file has not been modified since last run, skipping...") - return False - - if previous_update: - previous_update = time.strftime("%a, %d %b %Y %H:%M:%S %Z", time.gmtime(previous_update)) - if headers: - headers['If-Modified-Since'] = previous_update - else: - headers = {'If-Modified-Since': previous_update} - - logger.info(f"Downloading file from: {source['uri']}") - with session.get(uri, auth=auth, headers=headers, stream=True) as response: - # Check the response code - if response.status_code == requests.codes['not_modified']: - # File has not been modified since last update, do nothing - logger.info("The file has not been modified since last run, skipping...") - raise SkipSource - elif response.ok: - with open(target_path, 'wb') as f: - for content in response.iter_content(BLOCK_SIZE): - f.write(content) - - # Clear proxy setting - if proxy: - del os.environ['https_proxy'] - - # Return file_path - return True - except requests.Timeout: - pass - except Exception as e: - # Catch all other types of exceptions such as ConnectionError, ProxyError, etc. - logger.warning(str(e)) - return False - finally: - # Close the requests session - session.close() - - class SampleUpdateServer(ServiceUpdater): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) diff --git a/assemblyline_v4_service/updater/helper.py b/assemblyline_v4_service/updater/helper.py index 9bfe8f9b..4c1e12c8 100644 --- a/assemblyline_v4_service/updater/helper.py +++ b/assemblyline_v4_service/updater/helper.py @@ -164,17 +164,6 @@ def git_clone_repo(source: Dict[str, Any], previous_update: int = None, default_ add_cacert(ca_cert) git_env['GIT_SSL_CAINFO'] = certifi.where() - if key: - logger.info(f"key found for {url}") - # Save the key to a file - git_ssh_identity_file = os.path.join(tempfile.gettempdir(), 'id_rsa') - with open(git_ssh_identity_file, 'w') as key_fh: - key_fh.write(key) - os.chmod(git_ssh_identity_file, 0o0400) - - git_ssh_cmd = f"ssh -oStrictHostKeyChecking=no -i {git_ssh_identity_file}" - git_env['GIT_SSH_COMMAND'] = git_ssh_cmd - if auth: logger.info("Credentials provided for auth..") url = re.sub(r'^(?Phttps?://)', fr'\g{auth}', url) @@ -183,16 +172,27 @@ def git_clone_repo(source: Dict[str, Any], previous_update: int = None, default_ if os.path.exists(clone_dir): shutil.rmtree(clone_dir) - repo = Repo.clone_from(url, clone_dir, env=git_env, git_config=git_config) + with tempfile.NamedTemporaryFile() as git_ssh_identity_file: + if key: + logger.info(f"key found for {url}") + # Save the key to a file + git_ssh_identity_file.write(key.encode()) + git_ssh_identity_file.seek(0) + os.chmod(git_ssh_identity_file.name, 0o0400) - # Check repo last commit - if previous_update: - if isinstance(previous_update, str): - previous_update = iso_to_epoch(previous_update) - for c in repo.iter_commits(): - if c.committed_date < previous_update: - raise SkipSource() - break + git_ssh_cmd = f"ssh -oStrictHostKeyChecking=no -i {git_ssh_identity_file.name}" + git_env['GIT_SSH_COMMAND'] = git_ssh_cmd + + repo = Repo.clone_from(url, clone_dir, env=git_env, git_config=git_config) + + # Check repo last commit + if previous_update: + if isinstance(previous_update, str): + previous_update = iso_to_epoch(previous_update) + for c in repo.iter_commits(): + if c.committed_date < previous_update: + raise SkipSource() + break # Clear proxy setting if proxy: