From 9088ec4dbdbac9d246cf01aa090a2967baf39320 Mon Sep 17 00:00:00 2001
From: Aarsh2001 <ac5973@srmist.edu.in>
Date: Wed, 4 Oct 2023 18:29:25 +0100
Subject: [PATCH] gcp auth

---
 .github/auth/vm_auth.py         | 18 ++++++++++++++----
 .github/workflows/gcp-test.yaml |  2 +-
 2 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/.github/auth/vm_auth.py b/.github/auth/vm_auth.py
index ba0e0c38..fa4cbdbc 100644
--- a/.github/auth/vm_auth.py
+++ b/.github/auth/vm_auth.py
@@ -5,11 +5,12 @@
 from google.auth import compute_engine
 from googleapiclient import discovery
 from google.oauth2.service_account import Credentials
+from google.auth import impersonated_credentials
 
 def authenticate_vm(path):
     credentials = Credentials.from_service_account_file(path)
     return discovery.build('compute', 'v1', credentials=credentials)
-def start_runner(creds, pkey, id = "gpu-insatnce", zone='us-central1-a', instance='demos-tests'):
+def start_runner(creds, user, id = "gpu-insatnce", zone='us-central1-a', instance='demos-tests'):
     compute = authenticate_vm(creds)
     compute.instances().start(project=id, zone=zone, instance=instance).execute()
     request = compute.instances().get(project=id, zone=zone, instance=instance)
@@ -19,9 +20,18 @@ def start_runner(creds, pkey, id = "gpu-insatnce", zone='us-central1-a', instanc
     external_ip = response['networkInterfaces'][0]['accessConfigs'][0]['natIP']
 
     # Establish an SSH connection to the instance
+    credentials, _ = impersonated_credentials.load_credentials_from_file(
+        creds,
+        target_principal=id
+    )
     ssh = paramiko.SSHClient()
     ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
-    ssh.connect(external_ip, pkey=key)
+    ssh.connect(
+        hostname=f'{instance}.{zone}.compute.internal',
+        username=user,  # Typically 'your-username' or 'gce-username'
+        pkey=paramiko.RSAKey(file_obj=credentials.signer.key_file),
+    )
+
 
     # Execute the command on the instance
     stdin, stdout, stderr = ssh.exec_command('cd actions-runner; nohup ./run.sh')
@@ -35,9 +45,9 @@ def start_runner(creds, pkey, id = "gpu-insatnce", zone='us-central1-a', instanc
     return output
 
 if __name__ == "__main__":
-    key = sys.argv[1]
+    user = sys.argv[1]
     # Start the instance
-    start_runner('gcp_auth.json', pkey=key)
+    start_runner('gcp_auth.json', user)
 
 
 
diff --git a/.github/workflows/gcp-test.yaml b/.github/workflows/gcp-test.yaml
index 3ec13748..f6d3dab4 100644
--- a/.github/workflows/gcp-test.yaml
+++ b/.github/workflows/gcp-test.yaml
@@ -28,7 +28,7 @@ jobs:
          run: |
            cd demos/.github/auth
            python3 db_auth.py ${{ secrets.DB_ENDPOINT }} ${{ secrets.DB_OBJ_ID }}
-           python3 vm_auth.py ${{ secrets.SSH_KEY }}
+           python3 vm_auth.py ${{ secrets.SSH_USERNAME}}
 
   run-test:
     needs: activate-vm