<s> """<STR_LIT>""" <EOL> import datetime <EOL> import os . path <EOL> import time <EOL> import uuid <EOL> from apiclient import discovery <EOL> import httplib2 <EOL> import oauth2client . client <EOL> import oauth2client . file <EOL> import oauth2client . tools <EOL> from agents . base_agent import AgentConfigurationException <EOL> from agents . base_agent import AgentRuntimeException <EOL> from agents . base_agent import BaseAgent <EOL> from utils import utils <EOL> class GCEAgent ( BaseAgent ) : <EOL> """<STR_LIT>""" <EOL> MAX_VM_CREATION_TIME = <NUM_LIT> <EOL> SLEEP_TIME = <NUM_LIT:20> <EOL> PARAM_GROUP = '<STR_LIT>' <EOL> PARAM_IMAGE_ID = '<STR_LIT>' <EOL> PARAM_INSTANCE_IDS = '<STR_LIT>' <EOL> PARAM_INSTANCE_TYPE = '<STR_LIT>' <EOL> PARAM_KEYNAME = '<STR_LIT>' <EOL> PARAM_PROJECT = '<STR_LIT>' <EOL> PARAM_SECRETS = '<STR_LIT>' <EOL> PARAM_VERBOSE = '<STR_LIT>' <EOL> PARAM_ZONE = '<STR_LIT>' <EOL> REQUIRED_CREDENTIALS = ( <EOL> PARAM_GROUP , <EOL> PARAM_IMAGE_ID , <EOL> PARAM_KEYNAME , <EOL> PARAM_PROJECT , <EOL> PARAM_ZONE <EOL> ) <EOL> GCE_SCOPE = '<STR_LIT>' <EOL> API_VERSION = '<STR_LIT>' <EOL> GCE_URL = '<STR_LIT>' . format ( API_VERSION ) <EOL> DEFAULT_SERVICE_EMAIL = '<STR_LIT:default>' <EOL> GCE_PRIVATE_SSH_KEY = os . path . expanduser ( "<STR_LIT>" ) <EOL> GCE_PUBLIC_SSH_KEY = GCE_PRIVATE_SSH_KEY + "<STR_LIT>" <EOL> CLIENT_SECRETS_LOCATION = '<STR_LIT>' <EOL> OAUTH2_STORAGE_LOCATION = '<STR_LIT>' <EOL> def configure_instance_security ( self , parameters ) : <EOL> """<STR_LIT>""" <EOL> return True <EOL> def assert_required_parameters ( self , parameters , operation ) : <EOL> """<STR_LIT>""" <EOL> for param in self . REQUIRED_CREDENTIALS : <EOL> if param not in parameters : <EOL> raise AgentConfigurationException ( '<STR_LIT>' '<STR_LIT>' . format ( param ) ) <EOL> if not os . path . exists ( self . OAUTH2_STORAGE_LOCATION ) : <EOL> raise AgentConfigurationException ( '<STR_LIT>' '<STR_LIT>' . format ( self . OAUTH2_STORAGE_LOCATION ) ) <EOL> def describe_instances ( self , parameters ) : <EOL> """<STR_LIT>""" <EOL> gce_service , credentials = self . open_connection ( parameters ) <EOL> http = httplib2 . Http ( ) <EOL> auth_http = credentials . authorize ( http ) <EOL> request = gce_service . instances ( ) . list ( <EOL> project = parameters [ self . PARAM_PROJECT ] , <EOL> filter = "<STR_LIT>" . format ( parameters [ self . PARAM_GROUP ] ) , <EOL> zone = parameters [ self . PARAM_ZONE ] <EOL> ) <EOL> response = request . execute ( auth_http ) <EOL> utils . log ( str ( response ) ) <EOL> instance_ids = [ ] <EOL> public_ips = [ ] <EOL> private_ips = [ ] <EOL> if response and '<STR_LIT>' in response : <EOL> instances = response [ '<STR_LIT>' ] <EOL> for instance in instances : <EOL> if instance [ '<STR_LIT:status>' ] == "<STR_LIT>" : <EOL> instance_ids . append ( instance [ '<STR_LIT:name>' ] ) <EOL> public_ips . append ( instance [ '<STR_LIT>' ] [ <NUM_LIT:0> ] [ '<STR_LIT>' ] [ <NUM_LIT:0> ] [ '<STR_LIT>' ] ) <EOL> private_ips . append ( instance [ '<STR_LIT>' ] [ <NUM_LIT:0> ] [ '<STR_LIT>' ] ) <EOL> return public_ips , private_ips , instance_ids <EOL> def generate_disk_name ( self , parameters ) : <EOL> """<STR_LIT>""" <EOL> return "<STR_LIT>" . format ( parameters [ self . PARAM_GROUP ] , <EOL> str ( int ( time . time ( ) * <NUM_LIT:1000> ) ) ) <EOL> def create_scratch_disk ( self , parameters ) : <EOL> """<STR_LIT>""" <EOL> gce_service , credentials = self . open_connection ( parameters ) <EOL> http = httplib2 . Http ( ) <EOL> auth_http = credentials . authorize ( http ) <EOL> disk_name = self . generate_disk_name ( parameters ) <EOL> project_name = parameters [ self . PARAM_PROJECT ] <EOL> project_url = '<STR_LIT>' . format ( self . GCE_URL , <EOL> parameters [ self . PARAM_PROJECT ] ) <EOL> source_image_url = '<STR_LIT>' . format ( self . GCE_URL , <EOL> project_name , parameters [ self . PARAM_IMAGE_ID ] ) <EOL> request = gce_service . disks ( ) . insert ( <EOL> project = project_name , <EOL> zone = parameters [ self . PARAM_ZONE ] , <EOL> body = { <EOL> '<STR_LIT:name>' : disk_name <EOL> } , <EOL> sourceImage = source_image_url <EOL> ) <EOL> response = request . execute ( http = auth_http ) <EOL> utils . log ( str ( response ) ) <EOL> self . ensure_operation_succeeds ( gce_service , auth_http , response , <EOL> parameters [ self . PARAM_PROJECT ] ) <EOL> disk_url = "<STR_LIT>" . format ( <EOL> project_url , parameters [ self . PARAM_ZONE ] , disk_name ) <EOL> return disk_url <EOL> def run_instances ( self , count , parameters , security_configured ) : <EOL> """<STR_LIT>""" <EOL> project_id = parameters [ self . PARAM_PROJECT ] <EOL> image_id = parameters [ self . PARAM_IMAGE_ID ] <EOL> instance_type = parameters [ self . PARAM_INSTANCE_TYPE ] <EOL> keyname = parameters [ self . PARAM_KEYNAME ] <EOL> group = parameters [ self . PARAM_GROUP ] <EOL> zone = parameters [ self . PARAM_ZONE ] <EOL> utils . log ( "<STR_LIT>" "<STR_LIT>" . format ( count , image_id , instance_type , keyname , group , zone ) ) <EOL> start_time = datetime . datetime . now ( ) <EOL> active_public_ips , active_private_ips , active_instances = self . describe_instances ( parameters ) <EOL> image_url = '<STR_LIT>' . format ( self . GCE_URL , project_id , image_id ) <EOL> project_url = '<STR_LIT>' . format ( self . GCE_URL , project_id ) <EOL> machine_type_url = '<STR_LIT>' . format ( project_url , <EOL> zone , instance_type ) <EOL> zone_url = '<STR_LIT>' . format ( project_url , zone ) <EOL> network_url = '<STR_LIT>' . format ( project_url , group ) <EOL> public_key = utils . get_public_key ( keyname ) <EOL> for index in range ( count ) : <EOL> disk_url = self . create_scratch_disk ( parameters ) <EOL> instances = { <EOL> '<STR_LIT:name>' : "<STR_LIT>" . format ( group , uuid . uuid4 ( ) ) [ : <NUM_LIT> ] , <EOL> '<STR_LIT>' : machine_type_url , <EOL> '<STR_LIT>' : [ { <EOL> '<STR_LIT:source>' : disk_url , <EOL> '<STR_LIT>' : '<STR_LIT:true>' , <EOL> '<STR_LIT:type>' : '<STR_LIT>' <EOL> } ] , <EOL> '<STR_LIT:image>' : image_url , <EOL> '<STR_LIT>' : [ { <EOL> '<STR_LIT>' : [ { <EOL> '<STR_LIT:type>' : '<STR_LIT>' , <EOL> '<STR_LIT:name>' : '<STR_LIT>' <EOL> } ] , <EOL> '<STR_LIT>' : network_url <EOL> } ] , <EOL> '<STR_LIT>' : [ { <EOL> '<STR_LIT:email>' : self . DEFAULT_SERVICE_EMAIL , <EOL> '<STR_LIT>' : [ self . GCE_SCOPE ] <EOL> } ] , <EOL> '<STR_LIT>' : { <EOL> '<STR_LIT>' : [ { <EOL> '<STR_LIT:key>' : '<STR_LIT>' , <EOL> '<STR_LIT:value>' : '<STR_LIT>' . format ( public_key ) <EOL> } ] <EOL> } <EOL> } <EOL> gce_service , credentials = self . open_connection ( parameters ) <EOL> http = httplib2 . Http ( ) <EOL> auth_http = credentials . authorize ( http ) <EOL> request = gce_service . instances ( ) . insert ( <EOL> project = project_id , body = instances , zone = zone ) <EOL> response = request . execute ( auth_http ) <EOL> utils . log ( str ( response ) ) <EOL> self . ensure_operation_succeeds ( gce_service , auth_http , response , <EOL> parameters [ self . PARAM_PROJECT ] ) <EOL> instance_ids = [ ] <EOL> public_ips = [ ] <EOL> private_ips = [ ] <EOL> end_time = datetime . datetime . now ( ) + datetime . timedelta ( <NUM_LIT:0> , <EOL> self . MAX_VM_CREATION_TIME ) <EOL> now = datetime . datetime . now ( ) <EOL> while now < end_time : <EOL> time_left = ( end_time - now ) . seconds <EOL> utils . log ( "<STR_LIT>" ) <EOL> instance_info = self . describe_instances ( parameters ) <EOL> public_ips = instance_info [ <NUM_LIT:0> ] <EOL> private_ips = instance_info [ <NUM_LIT:1> ] <EOL> instance_ids = instance_info [ <NUM_LIT:2> ] <EOL> public_ips = utils . diff ( public_ips , active_public_ips ) <EOL> private_ips = utils . diff ( private_ips , active_private_ips ) <EOL> instance_ids = utils . diff ( instance_ids , active_instances ) <EOL> if count == len ( public_ips ) : <EOL> break <EOL> time . sleep ( self . SLEEP_TIME ) <EOL> now = datetime . datetime . now ( ) <EOL> if not public_ips : <EOL> self . handle_failure ( '<STR_LIT>' <EOL> '<STR_LIT>' ) <EOL> if len ( public_ips ) != count : <EOL> for index in range ( <NUM_LIT:0> , len ( public_ips ) ) : <EOL> if public_ips [ index ] == '<STR_LIT>' : <EOL> instance_to_term = instance_ids [ index ] <EOL> utils . log ( '<STR_LIT>' '<STR_LIT>' . format ( instance_to_term ) ) <EOL> self . terminate_instances ( [ instance_to_term ] ) <EOL> end_time = datetime . datetime . now ( ) <EOL> total_time = end_time - start_time <EOL> utils . log ( "<STR_LIT>" . format ( count , total_time . seconds ) ) <EOL> return instance_ids , public_ips , private_ips <EOL> def open_connection ( self , parameters ) : <EOL> """<STR_LIT>""" <EOL> storage = oauth2client . file . Storage ( self . OAUTH2_STORAGE_LOCATION ) <EOL> credentials = storage . get ( ) <EOL> return discovery . build ( '<STR_LIT>' , self . API_VERSION ) , credentials <EOL> def terminate_instances ( self , parameters ) : <EOL> """<STR_LIT>""" <EOL> instance_ids = parameters [ self . PARAM_INSTANCE_IDS ] <EOL> for instance_id in instance_ids : <EOL> gce_service , credentials = self . open_connection ( parameters ) <EOL> http = httplib2 . Http ( ) <EOL> auth_http = credentials . authorize ( http ) <EOL> request = gce_service . instances ( ) . delete ( <EOL> project = parameters [ self . PARAM_PROJECT ] , <EOL> zone = parameters [ self . PARAM_ZONE ] , <EOL> instance = instance_id <EOL> ) <EOL> response = request . execute ( auth_http ) <EOL> utils . log ( str ( response ) ) <EOL> self . ensure_operation_succeeds ( gce_service , auth_http , response , <EOL> parameters [ self . PARAM_PROJECT ] ) <EOL> def attach_disk ( self , parameters , disk_name , instance_id ) : <EOL> """<STR_LIT>""" <EOL> gce_service , credentials = self . open_connection ( parameters ) <EOL> http = httplib2 . Http ( ) <EOL> auth_http = credentials . authorize ( http ) <EOL> project = parameters [ self . PARAM_PROJECT ] <EOL> zone = parameters [ self . PARAM_ZONE ] <EOL> request = gce_service . instances ( ) . attachDisk ( <EOL> project = project , <EOL> zone = zone , <EOL> instance = instance_id , <EOL> body = { <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT:type>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT:source>' : "<STR_LIT>" "<STR_LIT>" . format ( self . API_VERSION , project , <EOL> zone , disk_name ) , <EOL> '<STR_LIT>' : '<STR_LIT>' <EOL> } <EOL> ) <EOL> response = request . execute ( auth_http ) <EOL> utils . log ( str ( response ) ) <EOL> self . ensure_operation_succeeds ( gce_service , auth_http , response , <EOL> parameters [ self . PARAM_PROJECT ] ) <EOL> return '<STR_LIT>' <EOL> def ensure_operation_succeeds ( self , gce_service , auth_http , response , project_id ) : <EOL> """<STR_LIT>""" <EOL> status = response [ '<STR_LIT:status>' ] <EOL> while status != '<STR_LIT>' and response : <EOL> operation_id = response [ '<STR_LIT:name>' ] <EOL> if '<STR_LIT>' in response : <EOL> zone_name = response [ '<STR_LIT>' ] . split ( '<STR_LIT:/>' ) [ - <NUM_LIT:1> ] <EOL> request = gce_service . zoneOperations ( ) . get ( <EOL> project = project_id , <EOL> operation = operation_id , <EOL> zone = zone_name ) <EOL> else : <EOL> request = gce_service . globalOperations ( ) . get ( <EOL> project = project_id , operation = operation_id ) <EOL> response = request . execute ( auth_http ) <EOL> if response : <EOL> status = response [ '<STR_LIT:status>' ] <EOL> if '<STR_LIT:error>' in response : <EOL> message = "<STR_LIT:\n>" . join ( [ errors [ '<STR_LIT:message>' ] for errors in <EOL> response [ '<STR_LIT:error>' ] [ '<STR_LIT>' ] ] ) <EOL> raise AgentRuntimeException ( str ( message ) ) </s>
