Detect image type instead of requesting it from user

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 <token>' --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

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:

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")

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")

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