client.py (4452B)
1 #!/usr/bin/env python 2 3 import argparse 4 import json 5 import pika 6 import pika.exceptions 7 import sys 8 9 from .config import RABBITMQ_HOST 10 11 12 EXCHANGE_NAME = 'computing' 13 14 15 class ServerCommunication: 16 def __enter__(self): 17 self.connection = pika.BlockingConnection( 18 pika.ConnectionParameters(host=RABBITMQ_HOST)) 19 self.channel = self.connection.channel() 20 21 self.channel.basic_consume( 22 queue='amq.rabbitmq.reply-to', 23 on_message_callback=self.on_response, 24 auto_ack=True) 25 26 return self 27 28 def on_response(self, ch, method, properties, body): 29 self.responses.append(json.loads(body)) 30 31 def on_timeout(self): 32 self.channel.stop_consuming() 33 34 def send_msg(self, data, response_expected=False): 35 self.responses = [] 36 37 reply_to = None 38 if response_expected: 39 reply_to = 'amq.rabbitmq.reply-to' 40 41 self.channel.basic_publish( 42 exchange=EXCHANGE_NAME, 43 routing_key=data['command'], 44 properties=pika.BasicProperties( 45 reply_to=reply_to, 46 content_type='application/json', 47 ), 48 body=json.dumps(data).encode('utf-8')) 49 50 if response_expected: 51 self.connection.call_later(5, self.on_timeout) 52 self.channel.start_consuming() 53 54 return self.responses 55 56 def __exit__(self, exc_type, exc_value, traceback): 57 self.connection.close() 58 59 60 def main(): 61 parser = argparse.ArgumentParser( 62 description='Manage aetherscale instances') 63 subparsers = parser.add_subparsers(dest='subparser_name') 64 65 create_vm_parser = subparsers.add_parser('create-vm') 66 create_vm_parser.add_argument( 67 '--image', help='Name of the image to create a VM from', required=True) 68 create_vm_parser.add_argument( 69 '--init-script', dest='init_script_path', 70 help='Script to execute at first boot of VM', required=False) 71 create_vm_parser.add_argument( 72 '--vpn', help='Name of the VPN to startup/join', required=False) 73 create_vm_parser.add_argument( 74 '--no-public-ip', dest='public_ip', action='store_false', default=True, 75 help='Do not assign a public interface to this VM') 76 start_vm_parser = subparsers.add_parser('start-vm') 77 start_vm_parser.add_argument( 78 '--vm-id', dest='vm_id', help='ID of the VM to start', required=True) 79 stop_vm_parser = subparsers.add_parser('stop-vm') 80 stop_vm_parser.add_argument( 81 '--vm-id', dest='vm_id', help='ID of the VM to stop', required=True) 82 stop_vm_parser.add_argument( 83 '--kill', dest='kill', action='store_true', default=False, 84 help='Kill the VM immediately, no graceful shutdown') 85 delete_vm_parser = subparsers.add_parser('delete-vm') 86 delete_vm_parser.add_argument( 87 '--vm-id', dest='vm_id', help='ID of the VM to delete', required=True) 88 subparsers.add_parser('list-vms') 89 90 args = parser.parse_args() 91 92 if args.subparser_name == 'list-vms': 93 response_expected = True 94 data = { 95 'command': 'list-vms', 96 } 97 elif args.subparser_name == 'create-vm': 98 response_expected = True 99 100 data = { 101 'command': 'create-vm', 102 'options': { 103 'image': args.image, 104 'public-ip': args.public_ip, 105 } 106 } 107 108 if args.vpn: 109 data['options']['vpn'] = args.vpn 110 111 if args.init_script_path: 112 with open(args.init_script_path, 'rt') as f: 113 data['options']['init-script'] = f.read() 114 elif args.subparser_name == 'stop-vm': 115 response_expected = True 116 data = { 117 'command': args.subparser_name, 118 'options': { 119 'vm-id': args.vm_id, 120 'kill': args.kill, 121 } 122 } 123 elif args.subparser_name in ['start-vm', 'delete-vm']: 124 response_expected = True 125 data = { 126 'command': args.subparser_name, 127 'options': { 128 'vm-id': args.vm_id, 129 } 130 } 131 else: 132 parser.print_usage() 133 sys.exit(1) 134 135 try: 136 with ServerCommunication() as c: 137 result = c.send_msg(data, response_expected) 138 print(json.dumps(result)) 139 except pika.exceptions.AMQPConnectionError: 140 print('Could not connect to AMQP broker. Is it running?', 141 file=sys.stderr)