diff --git a/gns3server/api/routes/controller/images.py b/gns3server/api/routes/controller/images.py index 6141e4152..0496e29ee 100644 --- a/gns3server/api/routes/controller/images.py +++ b/gns3server/api/routes/controller/images.py @@ -110,7 +110,7 @@ async def create_qemu_image( image_path: str, image_data: schemas.QemuDiskImageCreate, images_repo: ImagesRepository = Depends(get_repository(ImagesRepository)), -) -> models.Image: +) -> Optional[models.Image]: """ Create a new blank Qemu image. diff --git a/gns3server/db/repositories/images.py b/gns3server/db/repositories/images.py index e2d0f2401..1d52295e7 100644 --- a/gns3server/db/repositories/images.py +++ b/gns3server/db/repositories/images.py @@ -154,7 +154,7 @@ class ImagesRepository(BaseRepository): await self._db_session.refresh(image_db) # force refresh of updated_at value return image_db - async def save_verified_image(self, info: dict) -> models.Image: + async def save_verified_image(self, info: dict) -> Optional[models.Image]: """Upsert an exact path, preserving template associations and the image ID. Callers coordinate publication/inspection with image_lock(). Each commit @@ -217,11 +217,12 @@ class ImagesRepository(BaseRepository): except Exception: await self._db_session.rollback() raise + return None async def delete_image_exact(self, image_id: int) -> bool: result = await self._db_session.execute(delete(models.Image).where(models.Image.image_id == image_id)) await self._db_session.commit() - return result.rowcount > 0 + return cast(CursorResult, result).rowcount > 0 async def delete_image(self, image_path: str) -> bool: """ @@ -246,7 +247,7 @@ class ImagesRepository(BaseRepository): await self._db_session.commit() return cast(CursorResult, result).rowcount > 0 - async def prune_images(self, skip_images: Optional[list[str]] = None, is_in_use: Optional[Callable] = None) -> int: + async def prune_images(self, skip_images: Optional[list[str]] = None, is_in_use: Optional[Callable] = None) -> int: """ Prune images not attached to any template. """ diff --git a/gns3server/db/tasks.py b/gns3server/db/tasks.py index b3f762021..64f3f2673 100644 --- a/gns3server/db/tasks.py +++ b/gns3server/db/tasks.py @@ -203,7 +203,7 @@ async def update_disk_checksums(updated_disks: List[str]) -> None: repository = ImagesRepository(db_session) image = await repository.get_image(path) if image: - info = await read_image_info(path, image.image_type, allow_raw_image=True) + info = await read_image_info(path, str(image.image_type), allow_raw_image=True) try: os.unlink(path + ".md5sum") except FileNotFoundError: diff --git a/gns3server/utils/images.py b/gns3server/utils/images.py index fa4be5cb7..e2c458e1b 100644 --- a/gns3server/utils/images.py +++ b/gns3server/utils/images.py @@ -164,7 +164,7 @@ def inspect_image_file(path, expected_image_type=None, allow_raw_image=False, st ) -async def read_image_info(path: str, expected_image_type: str = None, allow_raw_image=False) -> dict: +async def read_image_info(path: str, expected_image_type: str | None = None, allow_raw_image=False) -> dict: try: return await asyncio.to_thread(inspect_image_file, path, expected_image_type, allow_raw_image) except OSError as e: @@ -376,7 +376,7 @@ async def write_image( break if len(prefix) < 7: raise InvalidImageError("The image content is empty or too small to be valid") - image_type = check_valid_image_header(image_path, prefix, allow_raw_image or not check_image_header) + image_type = check_valid_image_header(image_path, bytes(prefix), allow_raw_image or not check_image_header) if not image_dir: image_path = os.path.abspath(os.path.join(default_images_directory(image_type), image_name)) root = os.path.realpath(os.path.expanduser(Config.instance().settings.Server.images_path)) @@ -399,8 +399,8 @@ async def write_image( raise InvalidImageError( f"File '{image_path}' already exists, please choose a different name or remove the existing image" ) - checksum = checksum.hexdigest() - duplicate_image = await images_repo.get_image_by_checksum(checksum, os.path.dirname(image_path)) + checksum_str: str = checksum.hexdigest() + duplicate_image = await images_repo.get_image_by_checksum(checksum_str, os.path.dirname(image_path)) if duplicate_image: raise InvalidImageError( f"Image '{duplicate_image.filename}' with the same checksum " @@ -410,17 +410,20 @@ async def write_image( publish_image(tmp_path, image_path) # Complete files survive a database failure so the next scan can # recover them. Never compensate by unlinking a published image. - return await images_repo.save_verified_image( + image = await images_repo.save_verified_image( dict( image_name=image_name, image_type=image_type, image_size=image_size, path=image_path, - checksum=checksum, + checksum=checksum_str, checksum_algorithm="md5", file_fingerprint=fingerprint(image_path), ) ) + if image is None: + raise InvalidImageError(f"Failed to save image '{image_name}' to database") + return image finally: try: if os.path.exists(tmp_path): diff --git a/pyproject.toml b/pyproject.toml index 9c76a1d74..07c46dc7a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -277,9 +277,7 @@ extend-immutable-calls = [ "tests/test_utils.py" = ["F403", "F405"] "tests/utils/test_asyncio.py" = ["B017"] "tests/utils/test_packet_filter_validation.py" = ["W292"] -"tests/tests/services/test_image_reconciliation.py" = ["S324"] - - +"tests/services/test_image_reconciliation.py" = ["S324"] [tool.mypy] files = ["gns3server"]