Try to get the VM environment from the instance tags

Signed-off-by: Pedro Algarvio <palgarvio@vmware.com>

.github/workflows/test-action.yml

@@ -137,11 +137,18 @@ jobs:
             echo NOX_SESSION=${{ inputs.nox-session }}-tcp >> "$GITHUB_ENV"
           fi
 
+      - name: Get Salt Project GitHub Actions Bot Environment
+        if: steps.nox-dependencies-cache.outputs.cache-hit != 'true'
+        run: |
+          TOKEN=$(curl -sS -f -X PUT "http://169.254.169.254/latest/api/token" -H "X-aws-ec2-metadata-token-ttl-seconds: 30")
+          SPB_ENVIRONMENT=$(curl -sS -f -H "X-aws-ec2-metadata-token: $TOKEN" -v http://169.254.169.254/latest/meta-data/tags/instance/spb:environment)
+          echo "SPB_ENVIRONMENT=$SPB_ENVIRONMENT" >> "$GITHUB_ENV"
+
       - name: Start VM
         if: steps.nox-dependencies-cache.outputs.cache-hit != 'true'
         id: spin-up-vm
         run: |
-          tools --timestamps vm create --retries=2 ${{ inputs.distro-slug }}
+          tools --timestamps vm create --environment "${SPB_ENVIRONMENT}" --retries=2 ${{ inputs.distro-slug }}
 
       - name: List Free Space
         if: steps.nox-dependencies-cache.outputs.cache-hit != 'true'
@@ -241,12 +248,18 @@ jobs:
         with:
           name: testrun-changed-files.txt
 
+      - name: Get Salt Project GitHub Actions Bot Environment
+        run: |
+          TOKEN=$(curl -sS -f -X PUT "http://169.254.169.254/latest/api/token" -H "X-aws-ec2-metadata-token-ttl-seconds: 30")
+          SPB_ENVIRONMENT=$(curl -sS -f -H "X-aws-ec2-metadata-token: $TOKEN" -v http://169.254.169.254/latest/meta-data/tags/instance/spb:environment)
+          echo "SPB_ENVIRONMENT=$SPB_ENVIRONMENT" >> "$GITHUB_ENV"
+
       - name: Start VM
         id: spin-up-vm
         env:
           TESTS_CHUNK: ${{ matrix.tests-chunk }}
         run: |
-          tools --timestamps vm create --retries=2 ${{ inputs.distro-slug }}
+          tools --timestamps vm create --environment "${SPB_ENVIRONMENT}" --retries=2 ${{ inputs.distro-slug }}
 
       - name: List Free Space
         run: |

.github/workflows/test-packages-action.yml

@@ -110,11 +110,18 @@ jobs:
         if: steps.nox-dependencies-cache.outputs.cache-hit != 'true'
         uses: ./.github/actions/setup-python-tools-scripts
 
+      - name: Get Salt Project GitHub Actions Bot Environment
+        if: steps.nox-dependencies-cache.outputs.cache-hit != 'true'
+        run: |
+          TOKEN=$(curl -sS -f -X PUT "http://169.254.169.254/latest/api/token" -H "X-aws-ec2-metadata-token-ttl-seconds: 30")
+          SPB_ENVIRONMENT=$(curl -sS -f -H "X-aws-ec2-metadata-token: $TOKEN" -v http://169.254.169.254/latest/meta-data/tags/instance/spb:environment)
+          echo "SPB_ENVIRONMENT=$SPB_ENVIRONMENT" >> "$GITHUB_ENV"
+
       - name: Start VM
         if: steps.nox-dependencies-cache.outputs.cache-hit != 'true'
         id: spin-up-vm
         run: |
-          tools --timestamps vm create --retries=2 ${{ inputs.distro-slug }}
+          tools --timestamps vm create --environment "${SPB_ENVIRONMENT}" --retries=2 ${{ inputs.distro-slug }}
 
       - name: List Free Space
         if: steps.nox-dependencies-cache.outputs.cache-hit != 'true'
@@ -216,10 +223,16 @@ jobs:
       - name: Setup Python Tools Scripts
         uses: ./.github/actions/setup-python-tools-scripts
 
+      - name: Get Salt Project GitHub Actions Bot Environment
+        run: |
+          TOKEN=$(curl -sS -f -X PUT "http://169.254.169.254/latest/api/token" -H "X-aws-ec2-metadata-token-ttl-seconds: 30")
+          SPB_ENVIRONMENT=$(curl -sS -f -H "X-aws-ec2-metadata-token: $TOKEN" -v http://169.254.169.254/latest/meta-data/tags/instance/spb:environment)
+          echo "SPB_ENVIRONMENT=$SPB_ENVIRONMENT" >> "$GITHUB_ENV"
+
       - name: Start VM
         id: spin-up-vm
         run: |
-          tools --timestamps vm create --retries=2 ${{ inputs.distro-slug }}
+          tools --timestamps vm create --environment "${SPB_ENVIRONMENT}" --retries=2 ${{ inputs.distro-slug }}
 
       - name: List Free Space
         run: |

.pre-commit-config.yaml

@@ -1330,6 +1330,7 @@ repos:
           - rich
           - types-attrs
           - types-pyyaml
+          - types-requests
 
   - repo: https://github.com/saltstack/mirrors-nox
     rev: v2021.6.12

tools/vm.py

@@ -21,6 +21,7 @@ from functools import lru_cache
 from typing import TYPE_CHECKING, cast
 
 from ptscripts import Context, command_group
+from requests.exceptions import ConnectTimeout
 
 import tools.utils
 
@@ -95,8 +96,13 @@ vm.add_argument("--region", help="The AWS region.", default=AWS_REGION)
             "help": "How many times to retry creating and connecting to a vm",
         },
         "environment": {
-            "help": "The AWS environment to use.",
+            "help": (
-            "choices": ("prod", "test"),
+                "The AWS environment to use. When the value is auto, an "
+                "attempt will be made to get the right environment from the "
+                "AWS instance metadata endpoint. This only works for bastion "
+                "VMs."
+            ),
+            "choices": ("prod", "test", "auto"),
         },
     }
 )
@@ -117,6 +123,28 @@ def create(
         ctx.exit(1, "We need a key name to spin a VM")
     if not retries:
         retries = 1
+    if environment == "auto":
+        # Lets get the environment from the instance profile if we're on a bastion VM
+        with ctx.web as web:
+            try:
+                ret = web.put(
+                    "http://169.254.169.254/latest/api/token",
+                    headers={"X-aws-ec2-metadata-token-ttl-seconds": "10"},
+                    timeout=1,
+                )
+                token = ret.text.strip()
+                ret = web.get(
+                    "http://169.254.169.254/latest/meta-data/tags/instance/spb:environment",
+                    headers={"X-aws-ec2-metadata-token": token},
+                )
+                spb_environment = ret.text.strip()
+                if spb_environment:
+                    ctx.info(f"Discovered VM environment: {spb_environment}")
+                    environment = spb_environment
+            except ConnectTimeout:
+                # We're apparently not in bastion VM
+                environment = None
+
     attempts = 0
     while True:
         attempts += 1