From 846d27285060b7e8de4c44f3dd6ba94faaf7fa03 Mon Sep 17 00:00:00 2001
From: Scott Stephens <stephens.js@gmail.com>
Date: Fri, 8 Aug 2025 12:12:50 -0400
Subject: [PATCH 1/2] Refactor out ec2.register_image arguments

---
 upload-ami/src/upload_ami/upload_ami.py | 32 +++++++++++++------------
 1 file changed, 17 insertions(+), 15 deletions(-)

diff --git a/upload-ami/src/upload_ami/upload_ami.py b/upload-ami/src/upload_ami/upload_ami.py
index 85676c8..17aa563 100644
--- a/upload-ami/src/upload_ami/upload_ami.py
+++ b/upload-ami/src/upload_ami/upload_ami.py
@@ -12,7 +12,7 @@
 
 from mypy_boto3_ec2.client import EC2Client
 from mypy_boto3_ec2.literals import BootModeValuesType
-from mypy_boto3_ec2.type_defs import RegionTypeDef
+from mypy_boto3_ec2.type_defs import RegionTypeDef, RegisterImageRequestTypeDef
 from mypy_boto3_s3.client import S3Client
 
 from concurrent.futures import ThreadPoolExecutor
@@ -150,19 +150,17 @@ def register_image_if_not_exists(
         else:
             raise Exception("Unknown system: " + image_info["system"])
 
-        logging.info(f"Registering image {image_name} with snapshot {snapshot_id}")
-
         # TODO(arianvp): Not all instance types support TPM 2.0 yet. We should
         # upload two images, one with and one without TPM 2.0 support.
 
         # if architecture == "x86_64" and image_info["boot_mode"] == "uefi":
         #    tpmsupport['TpmSupport'] = "v2.0"
 
-        register_image = ec2.register_image(
-            Name=image_name,
-            Architecture=architecture,
-            BootMode=image_info["boot_mode"],
-            BlockDeviceMappings=[
+        register_image_kwargs: RegisterImageRequestTypeDef = {
+            "Name": image_name,
+            "Architecture": architecture,
+            "BootMode": image_info["boot_mode"],
+            "BlockDeviceMappings": [
                 {
                     "DeviceName": "/dev/xvda",
                     "Ebs": {
@@ -171,12 +169,12 @@ def register_image_if_not_exists(
                     },
                 }
             ],
-            RootDeviceName="/dev/xvda",
-            VirtualizationType="hvm",
-            EnaSupport=True,
-            ImdsSupport="v2.0",
-            SriovNetSupport="simple",
-            TagSpecifications=[
+            "RootDeviceName": "/dev/xvda",
+            "VirtualizationType": "hvm",
+            "EnaSupport": True,
+            "ImdsSupport": "v2.0",
+            "SriovNetSupport": "simple",
+            "TagSpecifications": [
                 {
                     "ResourceType": "image",
                     "Tags": [
@@ -185,7 +183,11 @@ def register_image_if_not_exists(
                     ],
                 }
             ],
-        )
+        }
+
+        logging.info(f"Registering image {image_name} with snapshot {snapshot_id}")
+
+        register_image = ec2.register_image(**register_image_kwargs)
         image_id = register_image["ImageId"]
 
     ec2.get_waiter("image_available").wait(ImageIds=[image_id])

From 1c6f0e80d15597ad985fd8bd03db32cd21f439f8 Mon Sep 17 00:00:00 2001
From: Scott Stephens <stephens.js@gmail.com>
Date: Fri, 8 Aug 2025 12:15:19 -0400
Subject: [PATCH 2/2] Add command line arg to enable TPM

---
 upload-ami/src/upload_ami/upload_ami.py | 24 +++++++++++++++++-------
 1 file changed, 17 insertions(+), 7 deletions(-)

diff --git a/upload-ami/src/upload_ami/upload_ami.py b/upload-ami/src/upload_ami/upload_ami.py
index 17aa563..c366638 100644
--- a/upload-ami/src/upload_ami/upload_ami.py
+++ b/upload-ami/src/upload_ami/upload_ami.py
@@ -127,6 +127,7 @@ def register_image_if_not_exists(
     image_info: ImageInfo,
     snapshot_id: str,
     public: bool,
+    enable_tpm: bool,
 ) -> str:
     """
     Register image if it doesn't exist yet
@@ -150,12 +151,6 @@ def register_image_if_not_exists(
         else:
             raise Exception("Unknown system: " + image_info["system"])
 
-        # TODO(arianvp): Not all instance types support TPM 2.0 yet. We should
-        # upload two images, one with and one without TPM 2.0 support.
-
-        # if architecture == "x86_64" and image_info["boot_mode"] == "uefi":
-        #    tpmsupport['TpmSupport'] = "v2.0"
-
         register_image_kwargs: RegisterImageRequestTypeDef = {
             "Name": image_name,
             "Architecture": architecture,
@@ -185,6 +180,13 @@ def register_image_if_not_exists(
             ],
         }
 
+        if (
+            enable_tpm
+            and architecture == "x86_64"
+            and image_info["boot_mode"] == "uefi"
+        ):
+            register_image_kwargs["TpmSupport"] = "v2.0"
+
         logging.info(f"Registering image {image_name} with snapshot {snapshot_id}")
 
         register_image = ec2.register_image(**register_image_kwargs)
@@ -305,6 +307,7 @@ def upload_ami(
     run_id: str,
     public: bool,
     dest_regions: list[str],
+    enable_tpm: bool,
 ) -> dict[str, str]:
     """
     Upload NixOS AMI to AWS and return the image ids for each region
@@ -326,7 +329,7 @@ def upload_ami(
     )
 
     image_id = register_image_if_not_exists(
-        ec2, image_name, image_info, snapshot_id, public
+        ec2, image_name, image_info, snapshot_id, public, enable_tpm
     )
 
     regions = filter(
@@ -368,6 +371,12 @@ def main() -> None:
         action="append",
         default=[],
     )
+    parser.add_argument(
+        "--enable-tpm",
+        action="store_true",
+        default=False,
+        help="Enable TPM 2.0 support for UEFI x86_64 images",
+    )
 
     args = parser.parse_args()
 
@@ -386,6 +395,7 @@ def main() -> None:
         args.run_id,
         args.public,
         args.dest_region,
+        args.enable_tpm,
     )
     print(json.dumps(image_ids))