1#!/usr/bin/env python3 2# Copyright (C) 2024 The Android Open Source Project 3# 4# Licensed under the Apache License, Version 2.0 (the 'License'); 5# you may not use this file except in compliance with the License. 6# You may obtain a copy of the License at 7# 8# http://www.apache.org/licenses/LICENSE-2.0 9# 10# Unless required by applicable law or agreed to in writing, software 11# distributed under the License is distributed on an 'AS IS' BASIS, 12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13# See the License for the specific language governing permissions and 14# limitations under the License. 15 16import argparse 17import os 18import subprocess 19import sys 20 21 22def list_instances(project_id): 23 try: 24 result = subprocess.run([ 25 'gcloud', 'compute', 'instances', 'list', '--project', project_id, 26 '--format', 'table(name,zone)' 27 ], 28 check=True, 29 stdout=subprocess.PIPE, 30 stderr=subprocess.PIPE, 31 text=True) 32 lines = result.stdout.strip().split('\n') 33 instances = [tuple(line.split()) for line in lines[1:]] # Skip the header 34 return instances 35 except subprocess.CalledProcessError as e: 36 print(f'Error retrieving instances: {e.stderr}') 37 sys.exit(1) 38 39 40def main(): 41 DEFAULT_PROJECT_ID = 'perfetto-ci' 42 # project_id = os.getenv('CLOUDSDK_CORE_PROJECT', DEFAULT_PROJECT_ID) 43 44 parser = argparse.ArgumentParser() 45 parser.add_argument( 46 '-p', 47 '--project-id', 48 metavar='PROJECT_ID', 49 required=False, 50 help='The Cloud project id. Defaults to CLOUDSDK_CORE_PROJECT', 51 default=os.getenv('CLOUDSDK_CORE_PROJECT', DEFAULT_PROJECT_ID)) 52 args = parser.parse_args() 53 project_id = args.project_id 54 55 print('Using Cloud project: %s' % project_id) 56 print('If this script fail ensure that:') 57 print(' - The cloud project has been configured as per go/gce-beyondcorp-ssh') 58 print(' - Register your key as per "Ensure that you are registered with OS') 59 60 instances = list_instances(project_id) 61 if not instances: 62 print('No GCE instances found.') 63 sys.exit(0) 64 65 print('Available VMs:') 66 for idx, (name, zone) in enumerate(instances, start=1): 67 print(f'{idx}. {name} ({zone})') 68 69 try: 70 vm_number = int(input('Enter the number of the VM you want to ssh into: ')) 71 if vm_number < 1 or vm_number > len(instances): 72 raise ValueError 73 except ValueError: 74 print('Invalid selection. Please run the script again.') 75 sys.exit(1) 76 77 # Get the selected VM's name and zone 78 selected_instance = instances[vm_number - 1] 79 vm_name, vm_zone = selected_instance 80 user = os.getenv('USER', 'username') 81 ssh_arg = '%s_google_com@nic0.%s.%s.c.%s.internal.gcpnode.com' % ( 82 user, vm_name, vm_zone, project_id) 83 print('ssh ' + ssh_arg) 84 os.execvp('ssh', ['ssh', ssh_arg]) 85 86 87if __name__ == '__main__': 88 main() 89