Enable regions (#45)

* Add enable-regions command This enables all regions that are disabled. This is useful when you want to upload an AMI to all regions, but some regions are disabled. * Expose all pyproject scripts as apps * quota increase only when change detected
NixOS · Jan 8, 2024 · a9f28b5 · a9f28b5
1 parent 5357737
commit a9f28b5
Show file tree

Hide file tree

Showing 5 changed files with 43 additions and 28 deletions.
diff --git a/flake.nix b/flake.nix
@@ -72,20 +72,13 @@
         });
 
       apps = lib.genAttrs self.lib.supportedSystems (system:
-        let pkgs = nixpkgs.legacyPackages.${system}; in {
-          nuke = {
-            type = "app";
-            program = "${self.packages.${system}.upload-ami}/bin/nuke";
-          };
-          smoke-test = {
-            type = "app";
-            program = "${self.packages.${system}.upload-ami}/bin/smoke-test";
-          };
-          disable-image-block-public-access = {
-            type = "app";
-            program = "${self.packages.${system}.upload-ami}/bin/disable-image-block-public-access";  
-          };
-        });
+        let
+          pkgs = nixpkgs.legacyPackages.${system};
+          upload-ami = self.packages.${system}.upload-ami;
+          mkApp = name: _: { type = "app"; program = "${upload-ami}/bin/${name}"; };
+        in
+          lib.mapAttrs mkApp self.packages.${system}.upload-ami.passthru.pyproject.project.scripts
+        );
 
 
       # TODO: unfortunately I don't have access to a aarch64-linux hardware with virtualisation support

diff --git a/upload-ami/default.nix b/upload-ami/default.nix
@@ -1,5 +1,6 @@
 { buildPythonApplication
 , python3Packages
+, lib
 }:
 
 let pyproject = builtins.fromTOML (builtins.readFile ./pyproject.toml);
@@ -16,4 +17,5 @@ buildPythonApplication {
   propagatedBuildInputs =
     map (name: python3Packages.${name}) pyproject.project.dependencies;
 
+  passthru.pyproject = pyproject;
 }
diff --git a/upload-ami/pyproject.toml b/upload-ami/pyproject.toml
@@ -12,4 +12,6 @@ dependencies = [
 upload-ami = "upload_ami.cli:main"
 nuke = "upload_ami.nuke:main"
 smoke-test = "upload_ami.smoke_test:main"
-disable-image-block-public-access = "upload_ami.disable_image_block_public_access:main"
+disable-image-block-public-access = "upload_ami.disable_image_block_public_access:main"
+enable-regions = "upload_ami.enable_regions:main"
+request-public-ami-quota-increase  = "upload_ami.request_public_ami_quota_increase:main"
diff --git a/upload-ami/src/upload_ami/enable_regions.py b/upload-ami/src/upload_ami/enable_regions.py
@@ -0,0 +1,19 @@
+import boto3
+import logging
+
+
+def main():
+    """
+    Enable all regions that are disabled
+
+    Due to rate limiting, you might  need to run this multiple times.
+    """
+    logging.basicConfig(level=logging.INFO)
+    account = boto3.client("account")
+    pages = account.get_paginator("list_regions").paginate(
+        RegionOptStatusContains=["DISABLED"]
+    )
+    for page in pages:
+        for region in page["Regions"]:
+            logging.info(f"enabling region {region['RegionName']}")
+            account.enable_region(RegionName=region["RegionName"])
diff --git a/upload-ami/src/upload_ami/request_public_ami_quota_increase.py b/upload-ami/src/upload_ami/request_public_ami_quota_increase.py
@@ -3,13 +3,10 @@
 
 
 def get_public_ami_service_quota(servicequotas):
-    pages = servicequotas.get_paginator(
-        'list_service_quotas').paginate(ServiceCode="ec2")
-    for page in pages:
-        for quota in page["Quotas"]:
-            if quota["QuotaName"] == "Public AMIs":
-                return quota
-    raise Exception("No public AMI quota found")
+    return next(servicequotas
+                .get_paginator('list_service_quotas')
+                .paginate(ServiceCode="ec2")
+                .search("Quotas[?QuotaName=='Public AMIs']"))
 
 
 def main():
@@ -21,12 +18,14 @@ def main():
             "service-quotas", region_name=region["RegionName"])
         service_quota = get_public_ami_service_quota(servicequotas)
         try:
-            logging.info(f"Requesting quota increase for {region['RegionName']}")
-            servicequotas.request_service_quota_increase(
-                ServiceCode="ec2",
-                QuotaCode=service_quota['QuotaCode'],
-                DesiredValue=100,
-            )
+            desired_value = 100
+            if service_quota['Value'] >= desired_value:
+                logging.info(
+                    f"Quota for {region['RegionName']} is already {service_quota['Value']}")
+                continue
+            logging.info(
+                f"Requesting quota increase for {region['RegionName']}")
+            servicequotas.request_service_quota_increase( ServiceCode="ec2", QuotaCode=service_quota['QuotaCode'], DesiredValue=100,)
         except Exception as e:
             logging.warn(e)