| Joseph Lanoux | b3e1f87 | 2015-01-30 11:13:07 +0000 | [diff] [blame] | 1 | # Copyright (c) 2015 Hewlett-Packard Development Company, L.P. | 
|  | 2 | # All Rights Reserved. | 
|  | 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 |  | 
| Clark Boylan | 844180e | 2017-03-15 15:24:58 -0700 | [diff] [blame] | 16 | import base64 | 
| Markus Zoeller | ae36ce8 | 2017-03-20 16:27:26 +0100 | [diff] [blame] | 17 | import socket | 
| xxj | 8eb9098 | 2017-04-10 21:18:39 +0800 | [diff] [blame] | 18 | import ssl | 
| Markus Zoeller | ae36ce8 | 2017-03-20 16:27:26 +0100 | [diff] [blame] | 19 | import struct | 
| Clark Boylan | 844180e | 2017-03-15 15:24:58 -0700 | [diff] [blame] | 20 | import textwrap | 
|  | 21 |  | 
| Markus Zoeller | ae36ce8 | 2017-03-20 16:27:26 +0100 | [diff] [blame] | 22 | import six | 
|  | 23 | from six.moves.urllib import parse as urlparse | 
|  | 24 |  | 
| Joseph Lanoux | b3e1f87 | 2015-01-30 11:13:07 +0000 | [diff] [blame] | 25 | from oslo_log import log as logging | 
|  | 26 | from oslo_utils import excutils | 
| Joseph Lanoux | b3e1f87 | 2015-01-30 11:13:07 +0000 | [diff] [blame] | 27 |  | 
| Ken'ichi Ohmichi | 0eb153c | 2015-07-13 02:18:25 +0000 | [diff] [blame] | 28 | from tempest.common import waiters | 
| Joseph Lanoux | b3e1f87 | 2015-01-30 11:13:07 +0000 | [diff] [blame] | 29 | from tempest import config | 
| Matthew Treinish | b19c55d | 2017-07-17 12:38:35 -0400 | [diff] [blame] | 30 | from tempest.lib.common import fixed_network | 
| Ken'ichi Ohmichi | 5403052 | 2016-03-02 11:01:34 -0800 | [diff] [blame] | 31 | from tempest.lib.common import rest_client | 
| Ken'ichi Ohmichi | 757833a | 2017-03-10 10:30:30 -0800 | [diff] [blame] | 32 | from tempest.lib.common.utils import data_utils | 
| Joseph Lanoux | b3e1f87 | 2015-01-30 11:13:07 +0000 | [diff] [blame] | 33 |  | 
| Markus Zoeller | ae36ce8 | 2017-03-20 16:27:26 +0100 | [diff] [blame] | 34 | if six.PY2: | 
|  | 35 | ord_func = ord | 
|  | 36 | else: | 
|  | 37 | ord_func = int | 
|  | 38 |  | 
| Joseph Lanoux | b3e1f87 | 2015-01-30 11:13:07 +0000 | [diff] [blame] | 39 | CONF = config.CONF | 
|  | 40 |  | 
|  | 41 | LOG = logging.getLogger(__name__) | 
|  | 42 |  | 
|  | 43 |  | 
| Andrea Frittoli | 88eb677 | 2017-08-07 21:06:27 +0100 | [diff] [blame] | 44 | def is_scheduler_filter_enabled(filter_name): | 
|  | 45 | """Check the list of enabled compute scheduler filters from config. | 
|  | 46 |  | 
|  | 47 | This function checks whether the given compute scheduler filter is | 
|  | 48 | available and configured in the config file. If the | 
|  | 49 | scheduler_available_filters option is set to 'all' (Default value. which | 
|  | 50 | means default filters are configured in nova) in tempest.conf then, this | 
|  | 51 | function returns True with assumption that requested filter 'filter_name' | 
|  | 52 | is one of available filter in nova ("nova.scheduler.filters.all_filters"). | 
|  | 53 | """ | 
|  | 54 |  | 
|  | 55 | filters = CONF.compute_feature_enabled.scheduler_available_filters | 
|  | 56 | if not filters: | 
|  | 57 | return False | 
|  | 58 | if 'all' in filters: | 
|  | 59 | return True | 
|  | 60 | if filter_name in filters: | 
|  | 61 | return True | 
|  | 62 | return False | 
|  | 63 |  | 
|  | 64 |  | 
| Andrea Frittoli (andreaf) | 476e919 | 2015-08-14 23:59:58 +0100 | [diff] [blame] | 65 | def create_test_server(clients, validatable=False, validation_resources=None, | 
| Joe Gordon | 8843f0f | 2015-03-17 15:07:34 -0700 | [diff] [blame] | 66 | tenant_network=None, wait_until=None, | 
| Anusha Ramineni | 9aaef8b | 2016-01-19 10:56:40 +0530 | [diff] [blame] | 67 | volume_backed=False, name=None, flavor=None, | 
| ghanshyam | 61db96e | 2016-12-16 12:49:25 +0900 | [diff] [blame] | 68 | image_id=None, **kwargs): | 
| Joseph Lanoux | b3e1f87 | 2015-01-30 11:13:07 +0000 | [diff] [blame] | 69 | """Common wrapper utility returning a test server. | 
|  | 70 |  | 
|  | 71 | This method is a common wrapper returning a test server that can be | 
|  | 72 | pingable or sshable. | 
|  | 73 |  | 
| Takashi NATSUME | 6d5a2b4 | 2015-09-08 11:27:49 +0900 | [diff] [blame] | 74 | :param clients: Client manager which provides OpenStack Tempest clients. | 
| Joseph Lanoux | b3e1f87 | 2015-01-30 11:13:07 +0000 | [diff] [blame] | 75 | :param validatable: Whether the server will be pingable or sshable. | 
|  | 76 | :param validation_resources: Resources created for the connection to the | 
| Andrea Frittoli (andreaf) | 9df3a52 | 2016-07-06 14:09:48 +0100 | [diff] [blame] | 77 | server. Include a keypair, a security group and an IP. | 
| Ken'ichi Ohmichi | d5bc31a | 2015-09-02 01:45:28 +0000 | [diff] [blame] | 78 | :param tenant_network: Tenant network to be used for creating a server. | 
| Ken'ichi Ohmichi | fc25e69 | 2015-09-02 01:48:06 +0000 | [diff] [blame] | 79 | :param wait_until: Server status to wait for the server to reach after | 
| Andrea Frittoli (andreaf) | 9df3a52 | 2016-07-06 14:09:48 +0100 | [diff] [blame] | 80 | its creation. | 
| ghanshyam | 61db96e | 2016-12-16 12:49:25 +0900 | [diff] [blame] | 81 | :param volume_backed: Whether the server is volume backed or not. | 
|  | 82 | If this is true, a volume will be created and | 
|  | 83 | create server will be requested with | 
|  | 84 | 'block_device_mapping_v2' populated with below | 
|  | 85 | values: | 
|  | 86 | -------------------------------------------- | 
|  | 87 | bd_map_v2 = [{ | 
|  | 88 | 'uuid': volume['volume']['id'], | 
|  | 89 | 'source_type': 'volume', | 
|  | 90 | 'destination_type': 'volume', | 
|  | 91 | 'boot_index': 0, | 
|  | 92 | 'delete_on_termination': True}] | 
|  | 93 | kwargs['block_device_mapping_v2'] = bd_map_v2 | 
|  | 94 | --------------------------------------------- | 
|  | 95 | If server needs to be booted from volume with other | 
|  | 96 | combination of bdm inputs than mentioned above, then | 
|  | 97 | pass the bdm inputs explicitly as kwargs and image_id | 
|  | 98 | as empty string (''). | 
| Andrea Frittoli (andreaf) | 9df3a52 | 2016-07-06 14:09:48 +0100 | [diff] [blame] | 99 | :param name: Name of the server to be provisioned. If not defined a random | 
|  | 100 | string ending with '-instance' will be generated. | 
|  | 101 | :param flavor: Flavor of the server to be provisioned. If not defined, | 
|  | 102 | CONF.compute.flavor_ref will be used instead. | 
|  | 103 | :param image_id: ID of the image to be used to provision the server. If not | 
|  | 104 | defined, CONF.compute.image_ref will be used instead. | 
| lei zhang | dd552b2 | 2015-11-25 20:41:48 +0800 | [diff] [blame] | 105 | :returns: a tuple | 
| Joseph Lanoux | b3e1f87 | 2015-01-30 11:13:07 +0000 | [diff] [blame] | 106 | """ | 
|  | 107 |  | 
|  | 108 | # TODO(jlanoux) add support of wait_until PINGABLE/SSHABLE | 
|  | 109 |  | 
| Anusha Ramineni | 9aaef8b | 2016-01-19 10:56:40 +0530 | [diff] [blame] | 110 | if name is None: | 
|  | 111 | name = data_utils.rand_name(__name__ + "-instance") | 
|  | 112 | if flavor is None: | 
|  | 113 | flavor = CONF.compute.flavor_ref | 
|  | 114 | if image_id is None: | 
|  | 115 | image_id = CONF.compute.image_ref | 
| Joseph Lanoux | b3e1f87 | 2015-01-30 11:13:07 +0000 | [diff] [blame] | 116 |  | 
|  | 117 | kwargs = fixed_network.set_networks_kwarg( | 
|  | 118 | tenant_network, kwargs) or {} | 
|  | 119 |  | 
| Ghanshyam | 4de44ae | 2015-12-25 10:34:00 +0900 | [diff] [blame] | 120 | multiple_create_request = (max(kwargs.get('min_count', 0), | 
|  | 121 | kwargs.get('max_count', 0)) > 1) | 
|  | 122 |  | 
| Joseph Lanoux | b3e1f87 | 2015-01-30 11:13:07 +0000 | [diff] [blame] | 123 | if CONF.validation.run_validation and validatable: | 
|  | 124 | # As a first implementation, multiple pingable or sshable servers will | 
|  | 125 | # not be supported | 
| Ghanshyam | 4de44ae | 2015-12-25 10:34:00 +0900 | [diff] [blame] | 126 | if multiple_create_request: | 
| Joseph Lanoux | b3e1f87 | 2015-01-30 11:13:07 +0000 | [diff] [blame] | 127 | msg = ("Multiple pingable or sshable servers not supported at " | 
|  | 128 | "this stage.") | 
|  | 129 | raise ValueError(msg) | 
|  | 130 |  | 
| Andrea Frittoli | 9f416dd | 2017-08-10 15:38:00 +0100 | [diff] [blame] | 131 | LOG.debug("Provisioning test server with validation resources %s", | 
|  | 132 | validation_resources) | 
| Joseph Lanoux | b3e1f87 | 2015-01-30 11:13:07 +0000 | [diff] [blame] | 133 | if 'security_groups' in kwargs: | 
|  | 134 | kwargs['security_groups'].append( | 
|  | 135 | {'name': validation_resources['security_group']['name']}) | 
|  | 136 | else: | 
|  | 137 | try: | 
|  | 138 | kwargs['security_groups'] = [ | 
|  | 139 | {'name': validation_resources['security_group']['name']}] | 
|  | 140 | except KeyError: | 
|  | 141 | LOG.debug("No security group provided.") | 
|  | 142 |  | 
|  | 143 | if 'key_name' not in kwargs: | 
|  | 144 | try: | 
|  | 145 | kwargs['key_name'] = validation_resources['keypair']['name'] | 
|  | 146 | except KeyError: | 
|  | 147 | LOG.debug("No key provided.") | 
|  | 148 |  | 
|  | 149 | if CONF.validation.connect_method == 'floating': | 
| Ken'ichi Ohmichi | fc25e69 | 2015-09-02 01:48:06 +0000 | [diff] [blame] | 150 | if wait_until is None: | 
|  | 151 | wait_until = 'ACTIVE' | 
| Joseph Lanoux | b3e1f87 | 2015-01-30 11:13:07 +0000 | [diff] [blame] | 152 |  | 
| Clark Boylan | 844180e | 2017-03-15 15:24:58 -0700 | [diff] [blame] | 153 | if 'user_data' not in kwargs: | 
|  | 154 | # If nothing overrides the default user data script then run | 
|  | 155 | # a simple script on the host to print networking info. This is | 
|  | 156 | # to aid in debugging ssh failures. | 
|  | 157 | script = ''' | 
|  | 158 | #!/bin/sh | 
|  | 159 | echo "Printing {user} user authorized keys" | 
|  | 160 | cat ~{user}/.ssh/authorized_keys || true | 
|  | 161 | '''.format(user=CONF.validation.image_ssh_user) | 
|  | 162 | script_clean = textwrap.dedent(script).lstrip().encode('utf8') | 
|  | 163 | script_b64 = base64.b64encode(script_clean) | 
|  | 164 | kwargs['user_data'] = script_b64 | 
|  | 165 |  | 
| Joe Gordon | 8843f0f | 2015-03-17 15:07:34 -0700 | [diff] [blame] | 166 | if volume_backed: | 
| zhufl | c6ce539 | 2016-08-17 14:34:37 +0800 | [diff] [blame] | 167 | volume_name = data_utils.rand_name(__name__ + '-volume') | 
| John Griffith | bc678ad | 2015-09-29 09:38:39 -0600 | [diff] [blame] | 168 | volumes_client = clients.volumes_v2_client | 
| ghanshyam | 3bd0d2b | 2017-03-23 01:57:28 +0000 | [diff] [blame] | 169 | params = {'name': volume_name, | 
| ghanshyam | 61db96e | 2016-12-16 12:49:25 +0900 | [diff] [blame] | 170 | 'imageRef': image_id, | 
|  | 171 | 'size': CONF.volume.volume_size} | 
|  | 172 | volume = volumes_client.create_volume(**params) | 
| lkuchlan | 52d7b0d | 2016-11-07 20:53:19 +0200 | [diff] [blame] | 173 | waiters.wait_for_volume_resource_status(volumes_client, | 
|  | 174 | volume['volume']['id'], | 
|  | 175 | 'available') | 
| Joe Gordon | 8843f0f | 2015-03-17 15:07:34 -0700 | [diff] [blame] | 176 |  | 
|  | 177 | bd_map_v2 = [{ | 
|  | 178 | 'uuid': volume['volume']['id'], | 
|  | 179 | 'source_type': 'volume', | 
|  | 180 | 'destination_type': 'volume', | 
|  | 181 | 'boot_index': 0, | 
| ghanshyam | 61db96e | 2016-12-16 12:49:25 +0900 | [diff] [blame] | 182 | 'delete_on_termination': True}] | 
| Joe Gordon | 8843f0f | 2015-03-17 15:07:34 -0700 | [diff] [blame] | 183 | kwargs['block_device_mapping_v2'] = bd_map_v2 | 
|  | 184 |  | 
|  | 185 | # Since this is boot from volume an image does not need | 
|  | 186 | # to be specified. | 
|  | 187 | image_id = '' | 
|  | 188 |  | 
| Ken'ichi Ohmichi | f2d436e | 2015-09-03 01:13:16 +0000 | [diff] [blame] | 189 | body = clients.servers_client.create_server(name=name, imageRef=image_id, | 
|  | 190 | flavorRef=flavor, | 
| Joseph Lanoux | b3e1f87 | 2015-01-30 11:13:07 +0000 | [diff] [blame] | 191 | **kwargs) | 
|  | 192 |  | 
|  | 193 | # handle the case of multiple servers | 
| Ghanshyam | 4de44ae | 2015-12-25 10:34:00 +0900 | [diff] [blame] | 194 | if multiple_create_request: | 
| Joseph Lanoux | b3e1f87 | 2015-01-30 11:13:07 +0000 | [diff] [blame] | 195 | # Get servers created which name match with name param. | 
|  | 196 | body_servers = clients.servers_client.list_servers() | 
|  | 197 | servers = \ | 
|  | 198 | [s for s in body_servers['servers'] if s['name'].startswith(name)] | 
| ghanshyam | 0f82525 | 2015-08-25 16:02:50 +0900 | [diff] [blame] | 199 | else: | 
| Ken'ichi Ohmichi | 5403052 | 2016-03-02 11:01:34 -0800 | [diff] [blame] | 200 | body = rest_client.ResponseBody(body.response, body['server']) | 
| ghanshyam | 0f82525 | 2015-08-25 16:02:50 +0900 | [diff] [blame] | 201 | servers = [body] | 
| Joseph Lanoux | b3e1f87 | 2015-01-30 11:13:07 +0000 | [diff] [blame] | 202 |  | 
| Artom Lifshitz | 70d7a11 | 2017-05-10 17:25:54 +0000 | [diff] [blame] | 203 | def _setup_validation_fip(): | 
|  | 204 | if CONF.service_available.neutron: | 
|  | 205 | ifaces = clients.interfaces_client.list_interfaces(server['id']) | 
|  | 206 | validation_port = None | 
|  | 207 | for iface in ifaces['interfaceAttachments']: | 
|  | 208 | if iface['net_id'] == tenant_network['id']: | 
|  | 209 | validation_port = iface['port_id'] | 
|  | 210 | break | 
|  | 211 | if not validation_port: | 
|  | 212 | # NOTE(artom) This will get caught by the catch-all clause in | 
|  | 213 | # the wait_until loop below | 
|  | 214 | raise ValueError('Unable to setup floating IP for validation: ' | 
|  | 215 | 'port not found on tenant network') | 
|  | 216 | clients.floating_ips_client.update_floatingip( | 
|  | 217 | validation_resources['floating_ip']['id'], | 
|  | 218 | port_id=validation_port) | 
|  | 219 | else: | 
|  | 220 | fip_client = clients.compute_floating_ips_client | 
|  | 221 | fip_client.associate_floating_ip_to_server( | 
|  | 222 | floating_ip=validation_resources['floating_ip']['ip'], | 
|  | 223 | server_id=servers[0]['id']) | 
| Joseph Lanoux | b3e1f87 | 2015-01-30 11:13:07 +0000 | [diff] [blame] | 224 |  | 
| Ken'ichi Ohmichi | fc25e69 | 2015-09-02 01:48:06 +0000 | [diff] [blame] | 225 | if wait_until: | 
| Joseph Lanoux | b3e1f87 | 2015-01-30 11:13:07 +0000 | [diff] [blame] | 226 | for server in servers: | 
|  | 227 | try: | 
| Ken'ichi Ohmichi | 0eb153c | 2015-07-13 02:18:25 +0000 | [diff] [blame] | 228 | waiters.wait_for_server_status( | 
| Ken'ichi Ohmichi | fc25e69 | 2015-09-02 01:48:06 +0000 | [diff] [blame] | 229 | clients.servers_client, server['id'], wait_until) | 
| Joseph Lanoux | b3e1f87 | 2015-01-30 11:13:07 +0000 | [diff] [blame] | 230 |  | 
|  | 231 | # Multiple validatable servers are not supported for now. Their | 
|  | 232 | # creation will fail with the condition above (l.58). | 
|  | 233 | if CONF.validation.run_validation and validatable: | 
|  | 234 | if CONF.validation.connect_method == 'floating': | 
| Artom Lifshitz | 70d7a11 | 2017-05-10 17:25:54 +0000 | [diff] [blame] | 235 | _setup_validation_fip() | 
| Joseph Lanoux | b3e1f87 | 2015-01-30 11:13:07 +0000 | [diff] [blame] | 236 |  | 
|  | 237 | except Exception: | 
|  | 238 | with excutils.save_and_reraise_exception(): | 
| Jordan Pittier | 87ba287 | 2016-03-08 11:43:11 +0100 | [diff] [blame] | 239 | for server in servers: | 
|  | 240 | try: | 
|  | 241 | clients.servers_client.delete_server( | 
|  | 242 | server['id']) | 
|  | 243 | except Exception: | 
| Jordan Pittier | 525ec71 | 2016-12-07 17:51:26 +0100 | [diff] [blame] | 244 | LOG.exception('Deleting server %s failed', | 
|  | 245 | server['id']) | 
| Artom Lifshitz | 9b3f42b | 2017-06-19 05:46:32 +0000 | [diff] [blame] | 246 | for server in servers: | 
|  | 247 | # NOTE(artom) If the servers were booted with volumes | 
|  | 248 | # and with delete_on_termination=False we need to wait | 
|  | 249 | # for the servers to go away before proceeding with | 
|  | 250 | # cleanup, otherwise we'll attempt to delete the | 
|  | 251 | # volumes while they're still attached to servers that | 
|  | 252 | # are in the process of being deleted. | 
|  | 253 | try: | 
|  | 254 | waiters.wait_for_server_termination( | 
|  | 255 | clients.servers_client, server['id']) | 
|  | 256 | except Exception: | 
|  | 257 | LOG.exception('Server %s failed to delete in time', | 
|  | 258 | server['id']) | 
| Joseph Lanoux | b3e1f87 | 2015-01-30 11:13:07 +0000 | [diff] [blame] | 259 |  | 
|  | 260 | return body, servers | 
| ghanshyam | 017b5fe | 2016-04-15 18:49:26 +0900 | [diff] [blame] | 261 |  | 
|  | 262 |  | 
| ghanshyam | 4c1391c | 2016-12-01 13:13:06 +0900 | [diff] [blame] | 263 | def shelve_server(servers_client, server_id, force_shelve_offload=False): | 
| ghanshyam | 017b5fe | 2016-04-15 18:49:26 +0900 | [diff] [blame] | 264 | """Common wrapper utility to shelve server. | 
|  | 265 |  | 
|  | 266 | This method is a common wrapper to make server in 'SHELVED' | 
|  | 267 | or 'SHELVED_OFFLOADED' state. | 
|  | 268 |  | 
| ghanshyam | 4c1391c | 2016-12-01 13:13:06 +0900 | [diff] [blame] | 269 | :param servers_clients: Compute servers client instance. | 
| ghanshyam | 017b5fe | 2016-04-15 18:49:26 +0900 | [diff] [blame] | 270 | :param server_id: Server to make in shelve state | 
|  | 271 | :param force_shelve_offload: Forcefully offload shelve server if it | 
|  | 272 | is configured not to offload server | 
|  | 273 | automatically after offload time. | 
|  | 274 | """ | 
| ghanshyam | 4c1391c | 2016-12-01 13:13:06 +0900 | [diff] [blame] | 275 | servers_client.shelve_server(server_id) | 
| ghanshyam | 017b5fe | 2016-04-15 18:49:26 +0900 | [diff] [blame] | 276 |  | 
|  | 277 | offload_time = CONF.compute.shelved_offload_time | 
|  | 278 | if offload_time >= 0: | 
| ghanshyam | 4c1391c | 2016-12-01 13:13:06 +0900 | [diff] [blame] | 279 | waiters.wait_for_server_status(servers_client, server_id, | 
| ghanshyam | 017b5fe | 2016-04-15 18:49:26 +0900 | [diff] [blame] | 280 | 'SHELVED_OFFLOADED', | 
|  | 281 | extra_timeout=offload_time) | 
|  | 282 | else: | 
| ghanshyam | 4c1391c | 2016-12-01 13:13:06 +0900 | [diff] [blame] | 283 | waiters.wait_for_server_status(servers_client, server_id, 'SHELVED') | 
| ghanshyam | 017b5fe | 2016-04-15 18:49:26 +0900 | [diff] [blame] | 284 | if force_shelve_offload: | 
| ghanshyam | 4c1391c | 2016-12-01 13:13:06 +0900 | [diff] [blame] | 285 | servers_client.shelve_offload_server(server_id) | 
|  | 286 | waiters.wait_for_server_status(servers_client, server_id, | 
| ghanshyam | 017b5fe | 2016-04-15 18:49:26 +0900 | [diff] [blame] | 287 | 'SHELVED_OFFLOADED') | 
| Markus Zoeller | ae36ce8 | 2017-03-20 16:27:26 +0100 | [diff] [blame] | 288 |  | 
|  | 289 |  | 
|  | 290 | def create_websocket(url): | 
|  | 291 | url = urlparse.urlparse(url) | 
| Jens Harbott | 6bc422d | 2017-09-27 10:29:34 +0000 | [diff] [blame] | 292 | for res in socket.getaddrinfo(url.hostname, url.port, | 
|  | 293 | socket.AF_UNSPEC, socket.SOCK_STREAM): | 
|  | 294 | af, socktype, proto, _, sa = res | 
|  | 295 | client_socket = socket.socket(af, socktype, proto) | 
|  | 296 | if url.scheme == 'https': | 
|  | 297 | client_socket = ssl.wrap_socket(client_socket) | 
|  | 298 | client_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) | 
|  | 299 | try: | 
|  | 300 | client_socket.connect(sa) | 
|  | 301 | except socket.error: | 
|  | 302 | client_socket.close() | 
|  | 303 | continue | 
|  | 304 | break | 
| xxj | 8eb9098 | 2017-04-10 21:18:39 +0800 | [diff] [blame] | 305 | else: | 
| Jens Harbott | 6bc422d | 2017-09-27 10:29:34 +0000 | [diff] [blame] | 306 | raise socket.error('WebSocket creation failed') | 
| Markus Zoeller | ae36ce8 | 2017-03-20 16:27:26 +0100 | [diff] [blame] | 307 | # Turn the Socket into a WebSocket to do the communication | 
|  | 308 | return _WebSocket(client_socket, url) | 
|  | 309 |  | 
|  | 310 |  | 
|  | 311 | class _WebSocket(object): | 
|  | 312 | def __init__(self, client_socket, url): | 
|  | 313 | """Contructor for the WebSocket wrapper to the socket.""" | 
|  | 314 | self._socket = client_socket | 
| jianghua wang | d22514a | 2017-05-08 08:05:04 +0100 | [diff] [blame] | 315 | # cached stream for early frames. | 
|  | 316 | self.cached_stream = b'' | 
| Markus Zoeller | ae36ce8 | 2017-03-20 16:27:26 +0100 | [diff] [blame] | 317 | # Upgrade the HTTP connection to a WebSocket | 
|  | 318 | self._upgrade(url) | 
|  | 319 |  | 
| jianghua wang | d22514a | 2017-05-08 08:05:04 +0100 | [diff] [blame] | 320 | def _recv(self, recv_size): | 
|  | 321 | """Wrapper to receive data from the cached stream or socket.""" | 
|  | 322 | if recv_size <= 0: | 
|  | 323 | return None | 
|  | 324 |  | 
|  | 325 | data_from_cached = b'' | 
|  | 326 | data_from_socket = b'' | 
|  | 327 | if len(self.cached_stream) > 0: | 
|  | 328 | read_from_cached = min(len(self.cached_stream), recv_size) | 
|  | 329 | data_from_cached += self.cached_stream[:read_from_cached] | 
|  | 330 | self.cached_stream = self.cached_stream[read_from_cached:] | 
|  | 331 | recv_size -= read_from_cached | 
|  | 332 | if recv_size > 0: | 
|  | 333 | data_from_socket = self._socket.recv(recv_size) | 
|  | 334 | return data_from_cached + data_from_socket | 
|  | 335 |  | 
| Markus Zoeller | ae36ce8 | 2017-03-20 16:27:26 +0100 | [diff] [blame] | 336 | def receive_frame(self): | 
|  | 337 | """Wrapper for receiving data to parse the WebSocket frame format""" | 
|  | 338 | # We need to loop until we either get some bytes back in the frame | 
|  | 339 | # or no data was received (meaning the socket was closed).  This is | 
|  | 340 | # done to handle the case where we get back some empty frames | 
|  | 341 | while True: | 
| jianghua wang | d22514a | 2017-05-08 08:05:04 +0100 | [diff] [blame] | 342 | header = self._recv(2) | 
| Markus Zoeller | ae36ce8 | 2017-03-20 16:27:26 +0100 | [diff] [blame] | 343 | # If we didn't receive any data, just return None | 
| Masayuki Igawa | 0c0f014 | 2017-04-10 17:22:02 +0900 | [diff] [blame] | 344 | if not header: | 
| Markus Zoeller | ae36ce8 | 2017-03-20 16:27:26 +0100 | [diff] [blame] | 345 | return None | 
|  | 346 | # We will make the assumption that we are only dealing with | 
|  | 347 | # frames less than 125 bytes here (for the negotiation) and | 
|  | 348 | # that only the 2nd byte contains the length, and since the | 
|  | 349 | # server doesn't do masking, we can just read the data length | 
|  | 350 | if ord_func(header[1]) & 127 > 0: | 
| jianghua wang | d22514a | 2017-05-08 08:05:04 +0100 | [diff] [blame] | 351 | return self._recv(ord_func(header[1]) & 127) | 
| Markus Zoeller | ae36ce8 | 2017-03-20 16:27:26 +0100 | [diff] [blame] | 352 |  | 
|  | 353 | def send_frame(self, data): | 
|  | 354 | """Wrapper for sending data to add in the WebSocket frame format.""" | 
|  | 355 | frame_bytes = list() | 
|  | 356 | # For the first byte, want to say we are sending binary data (130) | 
|  | 357 | frame_bytes.append(130) | 
|  | 358 | # Only sending negotiation data so don't need to worry about > 125 | 
|  | 359 | # We do need to add the bit that says we are masking the data | 
|  | 360 | frame_bytes.append(len(data) | 128) | 
|  | 361 | # We don't really care about providing a random mask for security | 
|  | 362 | # So we will just hard-code a value since a test program | 
|  | 363 | mask = [7, 2, 1, 9] | 
|  | 364 | for i in range(len(mask)): | 
|  | 365 | frame_bytes.append(mask[i]) | 
|  | 366 | # Mask each of the actual data bytes that we are going to send | 
|  | 367 | for i in range(len(data)): | 
|  | 368 | frame_bytes.append(ord_func(data[i]) ^ mask[i % 4]) | 
|  | 369 | # Convert our integer list to a binary array of bytes | 
|  | 370 | frame_bytes = struct.pack('!%iB' % len(frame_bytes), * frame_bytes) | 
|  | 371 | self._socket.sendall(frame_bytes) | 
|  | 372 |  | 
|  | 373 | def close(self): | 
|  | 374 | """Helper method to close the connection.""" | 
|  | 375 | # Close down the real socket connection and exit the test program | 
|  | 376 | if self._socket is not None: | 
|  | 377 | self._socket.shutdown(1) | 
|  | 378 | self._socket.close() | 
|  | 379 | self._socket = None | 
|  | 380 |  | 
|  | 381 | def _upgrade(self, url): | 
|  | 382 | """Upgrade the HTTP connection to a WebSocket and verify.""" | 
|  | 383 | # The real request goes to the /websockify URI always | 
|  | 384 | reqdata = 'GET /websockify HTTP/1.1\r\n' | 
|  | 385 | reqdata += 'Host: %s:%s\r\n' % (url.hostname, url.port) | 
|  | 386 | # Tell the HTTP Server to Upgrade the connection to a WebSocket | 
|  | 387 | reqdata += 'Upgrade: websocket\r\nConnection: Upgrade\r\n' | 
|  | 388 | # The token=xxx is sent as a Cookie not in the URI | 
|  | 389 | reqdata += 'Cookie: %s\r\n' % url.query | 
|  | 390 | # Use a hard-coded WebSocket key since a test program | 
|  | 391 | reqdata += 'Sec-WebSocket-Key: x3JJHMbDL1EzLkh9GBhXDw==\r\n' | 
|  | 392 | reqdata += 'Sec-WebSocket-Version: 13\r\n' | 
|  | 393 | # We are choosing to use binary even though browser may do Base64 | 
|  | 394 | reqdata += 'Sec-WebSocket-Protocol: binary\r\n\r\n' | 
|  | 395 | # Send the HTTP GET request and get the response back | 
|  | 396 | self._socket.sendall(reqdata.encode('utf8')) | 
|  | 397 | self.response = data = self._socket.recv(4096) | 
|  | 398 | # Loop through & concatenate all of the data in the response body | 
| jianghua wang | d22514a | 2017-05-08 08:05:04 +0100 | [diff] [blame] | 399 | end_loc = self.response.find(b'\r\n\r\n') | 
|  | 400 | while data and end_loc < 0: | 
| Markus Zoeller | ae36ce8 | 2017-03-20 16:27:26 +0100 | [diff] [blame] | 401 | data = self._socket.recv(4096) | 
|  | 402 | self.response += data | 
| jianghua wang | d22514a | 2017-05-08 08:05:04 +0100 | [diff] [blame] | 403 | end_loc = self.response.find(b'\r\n\r\n') | 
|  | 404 |  | 
|  | 405 | if len(self.response) > end_loc + 4: | 
|  | 406 | # In case some frames (e.g. the first RFP negotiation) have | 
|  | 407 | # arrived, cache it for next reading. | 
|  | 408 | self.cached_stream = self.response[end_loc + 4:] | 
|  | 409 | # ensure response ends with '\r\n\r\n'. | 
|  | 410 | self.response = self.response[:end_loc + 4] |