Detect image type instead of requesting it from user

This commit is contained in:
grossmj 2022-03-20 16:20:17 +10:00
parent 2a5a4b5f77
commit 9b39bfb845
5 changed files with 50 additions and 47 deletions

View File

@ -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

View File

@ -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:

View File

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

View File

@ -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:
def check_valid_image_header(data: bytes) -> str:
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)")
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")

View File

@ -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