diff --git a/host_modules/image_service.py b/host_modules/image_service.py new file mode 100644 index 00000000..e30e1293 --- /dev/null +++ b/host_modules/image_service.py @@ -0,0 +1,135 @@ +""" +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 os +import requests +import stat +import subprocess + +from host_modules import host_service +import tempfile + +MOD_NAME = "image_service" + +DEFAULT_IMAGE_SAVE_AS = "/tmp/downloaded-sonic.bin" + +logger = logging.getLogger(__name__) + + +class ImageService(host_service.HostModule): + """DBus endpoint that handles downloading and installing SONiC images""" + + @host_service.method( + host_service.bus_name(MOD_NAME), in_signature="ss", out_signature="is" + ) + def download(self, image_url, save_as): + """ + Download a SONiC image. + + Args: + image_url: url for remote image. + save_as: local path for the downloaded image. The directory must exist and be *all* writable. + """ + logger.info("Download new sonic image from {} as {}".format(image_url, save_as)) + # Check if the directory exists, is absolute and has write permission. + if not os.path.isabs(save_as): + logger.error("The path {} is not an absolute path".format(save_as)) + return errno.EINVAL, "Path is not absolute" + dir = os.path.dirname(save_as) + if not os.path.isdir(dir): + 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) + ): + 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) + if response.status_code != 200: + logger.error( + "Failed to download image: HTTP status code {}".format( + response.status_code + ) + ) + return errno.EIO, "HTTP error: {}".format(response.status_code) + + with tempfile.NamedTemporaryFile(dir="/tmp", delete=False) as tmp_file: + for chunk in response.iter_content(chunk_size=8192): + tmp_file.write(chunk) + temp_file_path = tmp_file.name + os.replace(temp_file_path, save_as) + return 0, "Download successful" + except Exception as e: + logger.error("Failed to write downloaded image to disk: {}".format(e)) + return errno.EIO, str(e) + + @host_service.method( + host_service.bus_name(MOD_NAME), in_signature="s", out_signature="is" + ) + def install(self, where): + """ + Install a a sonic image: + + Args: + where: either a local path or a remote url pointing to the image. + """ + logger.info("Using sonic-installer to install the image at {}.".format(where)) + cmd = ["/usr/local/bin/sonic-installer", "install", "-y", where] + result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE) + msg = "" + if result.returncode: + lines = result.stderr.decode().split("\n") + for line in lines: + if "Error" in line: + 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) diff --git a/tests/host_modules/image_service_test.py b/tests/host_modules/image_service_test.py new file mode 100644 index 00000000..22e70e82 --- /dev/null +++ b/tests/host_modules/image_service_test.py @@ -0,0 +1,372 @@ +import hashlib +import subprocess +import sys +import os +import stat +import pytest +from unittest import mock +from host_modules.image_service import ImageService + + +class TestImageService(object): + @mock.patch("dbus.SystemBus") + @mock.patch("dbus.service.BusName") + @mock.patch("dbus.service.Object.__init__") + @mock.patch("os.path.isdir") + @mock.patch("os.stat") + @mock.patch("requests.get") + def test_download_success( + self, mock_get, mock_stat, mock_isdir, MockInit, MockBusName, MockSystemBus + ): + """ + Test that the `download` method successfully downloads an image when the directory exists and is writable. + """ + # Arrange + image_service = ImageService(mod_name="image_service") + image_url = "http://example.com/sonic_image.img" + save_as = "/tmp/sonic_image.img" + mock_isdir.return_value = True + mock_stat.return_value.st_mode = stat.S_IWUSR | stat.S_IWGRP | stat.S_IWOTH + mock_response = mock.Mock() + mock_response.iter_content = lambda chunk_size: [b"data"] + mock_response.status_code = 200 + mock_response.iter_content = lambda chunk_size: iter([b"data"]) + mock_get.return_value = mock_response + + # Act + rc, msg = image_service.download(image_url, save_as) + + # Assert + assert rc == 0, "wrong return value" + assert ( + "download" in msg.lower() and "successful" in msg.lower() + ), "message should contains 'download' and 'successful'" + mock_get.assert_called_once_with(image_url, stream=True) + + @mock.patch("dbus.SystemBus") + @mock.patch("dbus.service.BusName") + @mock.patch("dbus.service.Object.__init__") + @mock.patch("os.path.isdir") + def test_download_fail_no_dir( + self, mock_isdir, MockInit, MockBusName, MockSystemBus + ): + """ + Test that the `download` method fails when the directory does not exist. + """ + # Arrange + image_service = ImageService(mod_name="image_service") + image_url = "http://example.com/sonic_image.img" + save_as = "/nonexistent_dir/sonic_image.img" + mock_isdir.return_value = False + + # Act + rc, msg = image_service.download(image_url, save_as) + + # Assert + assert rc != 0, "wrong return value" + assert ( + "not" in msg.lower() and "exist" in msg.lower() + ), "message should contains 'not' and 'exist'" + + @mock.patch("dbus.SystemBus") + @mock.patch("dbus.service.BusName") + @mock.patch("dbus.service.Object.__init__") + @mock.patch("os.path.isdir") + @mock.patch("os.stat") + def test_download_fail_missing_other_write( + self, mock_stat, mock_isdir, MockInit, MockBusName, MockSystemBus + ): + """ + Test that the `download` method fails when the directory is not writable by others. + """ + # Arrange + image_service = ImageService(mod_name="image_service") + image_url = "http://example.com/sonic_image.img" + save_as = "/tmp/sonic_image.img" + mock_isdir.return_value = True + mock_stat.return_value.st_mode = ( + stat.S_IWUSR | stat.S_IWGRP + ) # Missing write permission for others + + # Act + rc, msg = image_service.download(image_url, save_as) + + # Assert + assert rc != 0, "wrong return value" + assert ( + "permission" in msg.lower() or "writable" in msg.lower() + ), "message should contain 'permission' or 'writable'" + + @mock.patch("dbus.SystemBus") + @mock.patch("dbus.service.BusName") + @mock.patch("dbus.service.Object.__init__") + def test_download_failed_relative_path(self, MockInit, MockBusName, MockSystemBus): + """ + Test that the `download` method fails when the save_as path is not absolute. + """ + # Arrange + image_service = ImageService(mod_name="image_service") + image_url = "http://example.com/sonic_image.img" + save_as = "relative/path/sonic_image.img" + + # Act + rc, msg = image_service.download(image_url, save_as) + + # Assert + assert rc != 0, "wrong return value" + assert "absolute" in msg.lower(), "message should contain 'absolute'" + + @mock.patch("dbus.SystemBus") + @mock.patch("dbus.service.BusName") + @mock.patch("dbus.service.Object.__init__") + @mock.patch("os.path.isdir") + @mock.patch("os.stat") + @mock.patch("requests.get") + def test_download_failed_not_found( + self, mock_get, mock_stat, mock_isdir, MockInit, MockBusName, MockSystemBus + ): + """ + Test that the `download` method fails when the image URL is not found (404 error). + """ + # Arrange + image_service = ImageService(mod_name="image_service") + image_url = "http://example.com/nonexistent_image.img" + save_as = "/tmp/sonic_image.img" + mock_isdir.return_value = True + mock_stat.return_value.st_mode = stat.S_IWUSR | stat.S_IWGRP | stat.S_IWOTH + mock_response = mock.Mock() + mock_response.status_code = 404 + mock_get.return_value = mock_response + + # Act + rc, msg = image_service.download(image_url, save_as) + + # Assert + assert rc != 0, "wrong return value" + assert ( + "404" in msg and "error" in msg.lower() + ), "message should contain '404' and 'error'" + mock_get.assert_called_once_with(image_url, stream=True) + + @mock.patch("dbus.SystemBus") + @mock.patch("dbus.service.BusName") + @mock.patch("dbus.service.Object.__init__") + @mock.patch("os.path.isdir") + @mock.patch("os.stat") + @mock.patch("requests.get") + @mock.patch("tempfile.NamedTemporaryFile") + def test_download_fail_write_io_exception( + self, + mock_tempfile, + mock_get, + mock_stat, + mock_isdir, + MockInit, + MockBusName, + MockSystemBus, + ): + """ + Test that the `download` method fails when there is an IOError while writing the file. + """ + # Arrange + image_service = ImageService(mod_name="image_service") + image_url = "http://example.com/sonic_image.img" + save_as = "/tmp/sonic_image.img" + mock_isdir.return_value = True + mock_stat.return_value.st_mode = stat.S_IWUSR | stat.S_IWGRP | stat.S_IWOTH + mock_response = mock.Mock() + mock_response.iter_content = lambda chunk_size: [b"data"] + mock_response.status_code = 200 + mock_get.return_value = mock_response + mock_tempfile.side_effect = IOError("Disk write error") + + # Act + rc, msg = image_service.download(image_url, save_as) + + # Assert + assert rc != 0, "wrong return value" + assert ( + "disk write error" in msg.lower() + ), "message should contain 'disk write error'" + mock_get.assert_called_once_with(image_url, stream=True) + mock_tempfile.assert_called_once_with( + delete=False, dir=os.path.dirname(save_as) + ) + + @mock.patch("dbus.SystemBus") + @mock.patch("dbus.service.BusName") + @mock.patch("dbus.service.Object.__init__") + @mock.patch("subprocess.run") + def test_install_success(self, mock_run, MockInit, MockBusName, MockSystemBus): + """ + Test that the `install` method successfully installs an image. + """ + # Arrange + image_service = ImageService(mod_name="image_service") + where = "/tmp/sonic_image.img" + mock_result = mock.Mock() + mock_result.returncode = 0 + mock_result.stderr = b"" + mock_run.return_value = mock_result + + # Act + rc, msg = image_service.install(where) + + # Assert + assert rc == 0, "wrong return value" + assert msg == "", "message should be empty on success" + mock_run.assert_called_once_with( + ["/usr/local/bin/sonic-installer", "install", "-y", where], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + + @mock.patch("dbus.SystemBus") + @mock.patch("dbus.service.BusName") + @mock.patch("dbus.service.Object.__init__") + @mock.patch("subprocess.run") + def test_install_fail(self, mock_run, MockInit, MockBusName, MockSystemBus): + """ + Test that the `install` method fails when the installation command returns a non-zero exit code. + """ + # Arrange + image_service = ImageService(mod_name="image_service") + where = "/tmp/sonic_image.img" + mock_result = mock.Mock() + mock_result.returncode = 1 + mock_result.stderr = b"Error: Installation failed" + mock_run.return_value = mock_result + + # Act + rc, msg = image_service.install(where) + + # Assert + assert rc != 0, "wrong return value" + assert "Error" in msg, "message should contain 'Error'" + mock_run.assert_called_once_with( + ["/usr/local/bin/sonic-installer", "install", "-y", where], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + + @pytest.mark.parametrize( + "algorithm, expected_checksum", + [ + ("sha256", hashlib.sha256(b"test data").hexdigest()), + ("sha512", hashlib.sha512(b"test data").hexdigest()), + ("md5", hashlib.md5(b"test data").hexdigest()), + ], + ) + @mock.patch("dbus.SystemBus") + @mock.patch("dbus.service.BusName") + @mock.patch("dbus.service.Object.__init__") + @mock.patch("os.path.isfile") + @mock.patch("builtins.open", new_callable=mock.mock_open, read_data=b"test data") + def test_checksum( + self, + mock_open, + mock_isfile, + MockInit, + MockBusName, + MockSystemBus, + algorithm, + expected_checksum, + ): + """ + Test that the `checksum` method correctly calculates the checksum of a file for different algorithms. + """ + # Arrange + image_service = ImageService(mod_name="image_service") + file_path = "/tmp/test_file.img" + mock_isfile.return_value = True + + # Act + rc, checksum = image_service.checksum(file_path, algorithm) + + # Assert + assert rc == 0, "wrong return value" + assert checksum == expected_checksum, "checksum does not match expected value" + mock_isfile.assert_called_once_with(file_path) + mock_open.assert_called_once_with(file_path, "rb") + + @mock.patch("dbus.SystemBus") + @mock.patch("dbus.service.BusName") + @mock.patch("dbus.service.Object.__init__") + @mock.patch("os.path.isfile") + def test_checksum_no_such_file( + self, mock_isfile, MockInit, MockBusName, MockSystemBus + ): + """ + Test that the `checksum` method fails when the file does not exist. + """ + # Arrange + image_service = ImageService(mod_name="image_service") + file_path = "/nonexistent_dir/test_file.img" + algorithm = "sha256" + mock_isfile.return_value = False + + # Act + rc, msg = image_service.checksum(file_path, algorithm) + + # Assert + assert rc != 0, "wrong return value" + assert "not exist" in msg.lower(), "message should contain 'not exist'" + mock_isfile.assert_called_once_with(file_path) + + @mock.patch("dbus.SystemBus") + @mock.patch("dbus.service.BusName") + @mock.patch("dbus.service.Object.__init__") + @mock.patch("os.path.isfile") + def test_checksum_unsupported_algorithm( + self, mock_isfile, MockInit, MockBusName, MockSystemBus + ): + """ + Test that the `checksum` method fails when an unsupported algorithm is provided. + """ + # Arrange + image_service = ImageService(mod_name="image_service") + file_path = "/tmp/test_file.img" + algorithm = "unsupported_algo" + mock_isfile.return_value = True + + # Act + rc, msg = image_service.checksum(file_path, algorithm) + + # Assert + assert rc != 0, "wrong return value" + assert ( + "unsupported algorithm" in msg.lower() + ), "message should contain 'unsupported algorithm'" + mock_isfile.assert_called_once_with(file_path) + + @mock.patch("dbus.SystemBus") + @mock.patch("dbus.service.BusName") + @mock.patch("dbus.service.Object.__init__") + @mock.patch("os.path.isfile") + @mock.patch("builtins.open", new_callable=mock.mock_open, read_data=b"test data") + def test_checksum_general_exception( + self, mock_open, mock_isfile, MockInit, MockBusName, MockSystemBus + ): + """ + Test that the `checksum` method handles general exceptions during file reading. + """ + # Arrange + image_service = ImageService(mod_name="image_service") + file_path = "/tmp/test_file.img" + algorithm = "sha256" + mock_isfile.return_value = True + + with mock.patch.object(hashlib, algorithm) as mock_hash_func: + mock_hash_instance = mock_hash_func.return_value + mock_hash_instance.update.side_effect = Exception("General error") + + # Act + rc, msg = image_service.checksum(file_path, algorithm) + + # Assert + assert rc != 0, "wrong return value" + assert ( + "general error" in msg.lower() + ), "message should contain 'general error'" + mock_isfile.assert_called_once_with(file_path) + mock_open.assert_called_once_with(file_path, "rb")