From 9b39bfb845438c6fe29771b943beabcc646506cf Mon Sep 17 00:00:00 2001 From: grossmj Date: Sun, 20 Mar 2022 16:20:17 +1000 Subject: [PATCH] Detect image type instead of requesting it from user --- gns3server/api/routes/controller/images.py | 23 +++++---- gns3server/controller/appliance_manager.py | 6 +-- gns3server/schemas/controller/images.py | 3 +- gns3server/utils/images.py | 54 +++++++++++----------- tests/api/routes/controller/test_images.py | 13 ++---- 5 files changed, 51 insertions(+), 48 deletions(-) diff --git a/gns3server/api/routes/controller/images.py b/gns3server/api/routes/controller/images.py index ccd15b76..0ad05921 100644 --- a/gns3server/api/routes/controller/images.py +++ b/gns3server/api/routes/controller/images.py @@ -23,11 +23,13 @@ import logging import urllib.parse from fastapi import APIRouter, Request, Response, Depends, status +from starlette.requests import ClientDisconnect from sqlalchemy.orm.exc import MultipleResultsFound from typing import List, Optional from gns3server import schemas -from gns3server.utils.images import InvalidImageError, default_images_directory, write_image +from gns3server.config import Config +from gns3server.utils.images import InvalidImageError, write_image from gns3server.db.repositories.images import ImagesRepository from gns3server.db.repositories.templates import TemplatesRepository from gns3server.db.repositories.rbac import RbacRepository @@ -62,7 +64,6 @@ async def get_images( async def upload_image( image_path: str, request: Request, - image_type: schemas.ImageType = schemas.ImageType.qemu, images_repo: ImagesRepository = Depends(get_repository(ImagesRepository)), templates_repo: TemplatesRepository = Depends(get_repository(TemplatesRepository)), current_user: schemas.User = Depends(get_current_active_user), @@ -72,24 +73,26 @@ async def upload_image( """ Upload an image. - Example: curl -X POST http://host:port/v3/images/upload/my_image_name.qcow2?image_type=qemu \ + Example: curl -X POST http://host:port/v3/images/upload/my_image_name.qcow2 \ -H 'Authorization: Bearer ' --data-binary @"/path/to/image.qcow2" """ image_path = urllib.parse.unquote(image_path) image_dir, image_name = os.path.split(image_path) - directory = default_images_directory(image_type) - full_path = os.path.abspath(os.path.join(directory, image_dir, image_name)) - if os.path.commonprefix([directory, full_path]) != directory: + # check if the path is within the default images directory + base_images_directory = os.path.expanduser(Config.instance().settings.Server.images_path) + full_path = os.path.abspath(os.path.join(base_images_directory, image_dir, image_name)) + if os.path.commonprefix([base_images_directory, full_path]) != base_images_directory: raise ControllerForbiddenError(f"Cannot write image, '{image_path}' is forbidden") + print(image_path) if await images_repo.get_image(image_path): raise ControllerBadRequestError(f"Image '{image_path}' already exists") try: - image = await write_image(image_name, image_type, full_path, request.stream(), images_repo) - except (OSError, InvalidImageError) as e: - raise ControllerError(f"Could not save {image_type} image '{image_path}': {e}") + image = await write_image(image_path, full_path, request.stream(), images_repo) + except (OSError, InvalidImageError, ClientDisconnect) as e: + raise ControllerError(f"Could not save image '{image_path}': {e}") if install_appliances: # attempt to automatically create templates based on image checksum @@ -100,7 +103,7 @@ async def upload_image( templates_repo, rbac_repo, current_user, - directory + os.path.dirname(image.path) ) return image diff --git a/gns3server/controller/appliance_manager.py b/gns3server/controller/appliance_manager.py index eb968f0d..b7874deb 100644 --- a/gns3server/controller/appliance_manager.py +++ b/gns3server/controller/appliance_manager.py @@ -123,7 +123,7 @@ class ApplianceManager: async with HTTPClient.get(image_url) as response: if response.status != 200: raise ControllerError(f"Could not download '{image_name}' due to HTTP error code {response.status}") - await write_image(image_name, image_type, image_path, response.content.iter_any(), images_repo) + await write_image(image_name, image_path, response.content.iter_any(), images_repo) except (OSError, InvalidImageError) as e: raise ControllerError(f"Could not save {image_type} image '{image_path}': {e}") except ClientError as e: @@ -156,7 +156,7 @@ class ApplianceManager: image_path = os.path.join(image_dir, appliance_file) if os.path.exists(image_path) and await wait_run_in_executor(md5sum, image_path) == image_checksum: async with aiofiles.open(image_path, "rb") as f: - await write_image(appliance_file, appliance.type, image_path, f, images_repo) + await write_image(appliance_file, image_path, f, images_repo) else: # download the image if there is a direct download URL direct_download_url = image.get("direct_download_url") @@ -217,7 +217,7 @@ class ApplianceManager: try: schemas.Appliance.parse_obj(appliance.asdict()) except ValidationError as e: - log.warning(message=f"Could not validate appliance '{appliance.id}': {e}") + log.warning(f"Could not validate appliance '{appliance.id}': {e}") if appliance.versions: for version in appliance.versions: if version.get("name") == image_version: diff --git a/gns3server/schemas/controller/images.py b/gns3server/schemas/controller/images.py index bee6621e..64efcdbf 100644 --- a/gns3server/schemas/controller/images.py +++ b/gns3server/schemas/controller/images.py @@ -32,7 +32,8 @@ class ImageBase(BaseModel): Common image properties. """ - filename: str = Field(..., description="Image name") + filename: str = Field(..., description="Image filename") + path: str = Field(..., description="Image path") image_type: ImageType = Field(..., description="Image type") image_size: int = Field(..., description="Image size in bytes") checksum: str = Field(..., description="Checksum value") diff --git a/gns3server/utils/images.py b/gns3server/utils/images.py index 1c4c2e5b..b9f2ff22 100644 --- a/gns3server/utils/images.py +++ b/gns3server/utils/images.py @@ -225,45 +225,43 @@ class InvalidImageError(Exception): return self._message -def check_valid_image_header(data: bytes, image_type: str, header_magic_len: int) -> None: - - if image_type == "ios": - # file must start with the ELF magic number, be 32-bit, big endian and have an ELF version of 1 - if data[:header_magic_len] != b'\x7fELF\x01\x02\x01': - raise InvalidImageError("Invalid IOS file detected") - elif image_type == "iou": - # file must start with the ELF magic number, be 32-bit or 64-bit, little endian and have an ELF version of 1 - # (normal IOS images are big endian!) - if data[:header_magic_len] != b'\x7fELF\x01\x01\x01' and data[:7] != b'\x7fELF\x02\x01\x01': - raise InvalidImageError("Invalid IOU file detected") - elif image_type == "qemu": - if data[:header_magic_len] != b'QFI\xfb' and data[:header_magic_len] != b'KDMV': - raise InvalidImageError("Invalid Qemu file detected (must be qcow2 or VDMK format)") +def check_valid_image_header(data: bytes) -> str: + + if data[:7] == b'\x7fELF\x01\x02\x01': + # for IOS images: file must start with the ELF magic number, be 32-bit, big endian and have an ELF version of 1 + return "ios" + elif data[:7] == b'\x7fELF\x01\x01\x01' or data[:7] == b'\x7fELF\x02\x01\x01': + # for IOU images file must start with the ELF magic number, be 32-bit or 64-bit, little endian and + # have an ELF version of 1 (normal IOS images are big endian!) + return "iou" + elif data[:4] != b'QFI\xfb' or data[:4] != b'KDMV': + return "qemu" + else: + raise InvalidImageError("Could not detect image type, please make sure it is a valid image") async def write_image( - image_name: str, - image_type: str, - path: str, + image_filename: str, + image_path: str, stream: AsyncGenerator[bytes, None], images_repo: ImagesRepository, check_image_header=True ) -> models.Image: - log.info(f"Writing image file to '{path}'") + image_dir, image_name = os.path.split(image_filename) + log.info(f"Writing image file to '{image_path}'") # Store the file under its final name only when the upload is completed - tmp_path = path + ".tmp" - os.makedirs(os.path.dirname(path), exist_ok=True) + tmp_path = image_path + ".tmp" + os.makedirs(os.path.dirname(image_path), exist_ok=True) checksum = hashlib.md5() header_magic_len = 7 - if image_type == "qemu": - header_magic_len = 4 + image_type = None try: async with aiofiles.open(tmp_path, "wb") as f: async for chunk in stream: if check_image_header and len(chunk) >= header_magic_len: check_image_header = False - check_valid_image_header(chunk, image_type, header_magic_len) + image_type = check_valid_image_header(chunk) await f.write(chunk) checksum.update(chunk) @@ -273,12 +271,16 @@ async def write_image( checksum = checksum.hexdigest() duplicate_image = await images_repo.get_image_by_checksum(checksum) - if duplicate_image and os.path.dirname(duplicate_image.path) == os.path.dirname(path): + if duplicate_image and os.path.dirname(duplicate_image.path) == os.path.dirname(image_path): raise InvalidImageError(f"Image {duplicate_image.filename} with " f"same checksum already exists in the same directory") except InvalidImageError: os.remove(tmp_path) raise os.chmod(tmp_path, stat.S_IWRITE | stat.S_IREAD | stat.S_IEXEC) - shutil.move(tmp_path, path) - return await images_repo.add_image(image_name, image_type, image_size, path, checksum, checksum_algorithm="md5") + if not image_dir: + directory = default_images_directory(image_type) + os.makedirs(directory, exist_ok=True) + image_path = os.path.abspath(os.path.join(directory, image_filename)) + shutil.move(tmp_path, image_path) + return await images_repo.add_image(image_name, image_type, image_size, image_path, checksum, checksum_algorithm="md5") diff --git a/tests/api/routes/controller/test_images.py b/tests/api/routes/controller/test_images.py index 61d3d33f..4ca09389 100644 --- a/tests/api/routes/controller/test_images.py +++ b/tests/api/routes/controller/test_images.py @@ -60,7 +60,7 @@ def ios_image(tmpdir) -> str: Create a fake IOS image on disk """ - path = os.path.join(tmpdir, "ios.bin") + path = os.path.join(tmpdir, "ios_image.bin") with open(path, "wb+") as f: f.write(b'\x7fELF\x01\x02\x01') return path @@ -74,7 +74,7 @@ def qcow2_image(tmpdir) -> str: path = os.path.join(tmpdir, "image.qcow2") with open(path, "wb+") as f: - f.write(b'QFI\xfb') + f.write(b'QFI\xfb\x00\x00\x00') return path @@ -137,7 +137,6 @@ class TestImageRoutes: response = await client.post( app.url_path_for("upload_image", image_path=image_name), - params={"image_type": image_type}, content=image_data) if valid_request: @@ -168,7 +167,6 @@ class TestImageRoutes: image_data = f.read() response = await client.post( app.url_path_for("upload_image", image_path=image_name), - params={"image_type": "qemu"}, content=image_data) assert response.status_code == status.HTTP_400_BAD_REQUEST @@ -191,7 +189,6 @@ class TestImageRoutes: image_data = f.read() response = await client.post( app.url_path_for("upload_image", image_path=image_name), - params={"image_type": "qemu"}, content=image_data) assert response.status_code == status.HTTP_201_CREATED @@ -214,7 +211,8 @@ class TestImageRoutes: images_dir: str, qcow2_image: str, subdir: str, - expected_result: int + expected_result: int, + db_session: AsyncSession ) -> None: image_name = os.path.basename(qcow2_image) @@ -223,7 +221,6 @@ class TestImageRoutes: image_path = os.path.join(subdir, image_name) response = await client.post( app.url_path_for("upload_image", image_path=image_path), - params={"image_type": "qemu"}, content=image_data) assert response.status_code == expected_result @@ -273,7 +270,7 @@ class TestImageRoutes: image_data = f.read() response = await client.post( app.url_path_for("upload_image", image_path=image_name), - params={"image_type": "qemu", "install_appliances": "true"}, + params={"install_appliances": "true"}, content=image_data) assert response.status_code == status.HTTP_201_CREATED