Merge pull request #111 from CybercentreCanada/persistent-service-update

Persistent service update

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:

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)

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'^(?P<scheme>https?://)', fr'\g<scheme>{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: