diff --git a/host_modules/image_service.py b/host_modules/image_service.py index 0da91e3d..5525d0d1 100644 --- a/host_modules/image_service.py +++ b/host_modules/image_service.py @@ -1,15 +1,17 @@ """ -Services related to SONiC images, such as: -1) Download -2) Install +This module provides services related to SONiC images, including: +1) Downloading images +2) Installing images +3) Calculating checksums for images """ +import errno +import hashlib import logging -import subprocess import os -import stat import requests -import errno +import stat +import subprocess from host_modules import host_service @@ -41,13 +43,17 @@ def download(self, image_url, save_as): logger.error("Directory {} does not exist".format(dir)) return errno.ENOENT, "Directory does not exist" st_mode = os.stat(dir).st_mode - if not (st_mode & stat.S_IWUSR) or not (st_mode & stat.S_IWGRP) or not (st_mode & stat.S_IWOTH): + if ( + not (st_mode & stat.S_IWUSR) + or not (st_mode & stat.S_IWGRP) + or not (st_mode & stat.S_IWOTH) + ): logger.error("Directory {} is not all writable {}".format(dir, st_mode)) return errno.EACCES, "Directory is not all writable" try: response = requests.get(image_url, stream=True) response.raise_for_status() - with open(save_as, 'wb') as f: + with open(save_as, "wb") as f: for chunk in response.iter_content(chunk_size=8192): f.write(chunk) return 0, "Download successful" @@ -76,3 +82,41 @@ def install(self, where): msg = line break return result.returncode, msg + + @host_service.method( + host_service.bus_name(MOD_NAME), in_signature="ss", out_signature="is" + ) + def checksum(self, file_path, algorithm): + """ + Calculate the checksum of a file. + + Args: + file_path: path to the file. + algorithm: checksum algorithm to use (sha256, sha512, md5). + """ + + logger.info("Calculating {} checksum for file {}".format(algorithm, file_path)) + + if not os.path.isfile(file_path): + logger.error("File {} does not exist".format(file_path)) + return errno.ENOENT, "File does not exist" + + hash_func = None + if algorithm == "sha256": + hash_func = hashlib.sha256() + elif algorithm == "sha512": + hash_func = hashlib.sha512() + elif algorithm == "md5": + hash_func = hashlib.md5() + else: + logger.error("Unsupported algorithm: {}".format(algorithm)) + return errno.EINVAL, "Unsupported algorithm" + + try: + with open(file_path, "rb") as f: + for chunk in iter(lambda: f.read(4096), b""): + hash_func.update(chunk) + return 0, hash_func.hexdigest() + except Exception as e: + logger.error("Failed to calculate checksum: {}".format(e)) + return errno.EIO, str(e)