From 3744c343d466850862714df156232608c31f31ca Mon Sep 17 00:00:00 2001
From: "Alexander \"PapaTutuWawa" <papatutuwawa@polynom.me>
Date: Mon, 31 Mar 2025 00:15:24 +0200
Subject: [PATCH] Check image actions against the owning user

---
 src/openec2/actions/deregister_image.py |  6 +++++-
 src/openec2/actions/import_image.py     | 13 ++++++++++---
 src/openec2/db/image.py                 |  3 +++
 3 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/src/openec2/actions/deregister_image.py b/src/openec2/actions/deregister_image.py
index 31510b4..42332cc 100644
--- a/src/openec2/actions/deregister_image.py
+++ b/src/openec2/actions/deregister_image.py
@@ -13,13 +13,17 @@ def deregister_image(
     params: QueryParams,
     config: OpenEC2Config,
     db: DatabaseDep,
-    _: User,
+    user: User,
 ):
     image_id = params["ImageId"]
     ami = db.exec(select(AMI).where(AMI.id == image_id)).one()
     if ami is None:
         raise HTTPException(status_code=404, detail="Unknown AMI")
 
+    # Check if the requester can deregister the image.
+    if ami.owner_id != user.id:
+        raise HTTPException(status_code=403)
+
     # Mark the image as deregistered
     ami.deregistered = True
     db.add(ami)
diff --git a/src/openec2/actions/import_image.py b/src/openec2/actions/import_image.py
index 02ed6b5..ee4d4c6 100644
--- a/src/openec2/actions/import_image.py
+++ b/src/openec2/actions/import_image.py
@@ -4,10 +4,11 @@ from urllib.parse import urlparse
 import uuid
 import shutil
 
+from fastapi import HTTPException
 from fastapi.datastructures import QueryParams
 import requests
 
-from openec2.config import OpenEC2Config
+from openec2.config import OpenEC2Config, ConfigSingleton
 from openec2.db import DatabaseDep
 from openec2.db.user import User
 from openec2.db.image import AMI
@@ -17,7 +18,7 @@ def import_image(
     params: QueryParams,
     config: OpenEC2Config,
     db: DatabaseDep,
-    _: User,
+    user: User,
 ):
     first_disk_image_url = params["DiskContainer.1.Url"]
     url = urlparse(first_disk_image_url)
@@ -35,6 +36,11 @@ def import_image(
                     for chunk in r.iter_content(8196):
                         f.write(chunk)
         case "file":
+            if not ConfigSingleton.of().config.debug:
+                raise HTTPException(
+                    status_code=400,
+                    detail="Unsupported scheme",
+                )
             shutil.copy(
                 url.path,
                 str(dst),
@@ -47,7 +53,8 @@ def import_image(
             id=ami_id,
             description=None,
             originalFilename=filename,
-        )
+            owner_id=user.id,
+        ),
     )
     db.commit()
     return ami_id
diff --git a/src/openec2/db/image.py b/src/openec2/db/image.py
index e433aa3..17374f6 100644
--- a/src/openec2/db/image.py
+++ b/src/openec2/db/image.py
@@ -13,3 +13,6 @@ class AMI(SQLModel, table=True):
 
     # Was the image registered
     deregistered: bool = Field(default=False)
+
+    # Owner of the image who created it
+    owner_id: int = Field(foreign_key="user.id")