Mercurial > public > mercurial-scm > hg-stable

comparison contrib/automation/hgautomation/aws.py @ 43076:2372284d9457

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
formatting: blacken the codebase This is using my patch to black (https://github.com/psf/black/pull/826) so we don't un-wrap collection literals. Done with: hg files 'set:**.py - mercurial/thirdparty/** - "contrib/python-zstandard/**"' | xargs black -S # skip-blame mass-reformatting only # no-check-commit reformats foo_bar functions Differential Revision: https://phab.mercurial-scm.org/D6971
author Augie Fackler <augie@google.com>
date Sun, 06 Oct 2019 09:45:02 -0400
parents d1d919f679f7
children c09e8ac3f61f
comparison
equal deleted inserted replaced
43075:57875cf423c9 43076:2372284d9457
17 import time 17 import time
18 18
19 import boto3 19 import boto3
20 import botocore.exceptions 20 import botocore.exceptions
21 21
22 from .linux import ( 22 from .linux import BOOTSTRAP_DEBIAN
23 BOOTSTRAP_DEBIAN,
24 )
25 from .ssh import ( 23 from .ssh import (
26 exec_command as ssh_exec_command, 24 exec_command as ssh_exec_command,
27 wait_for_ssh, 25 wait_for_ssh,
28 ) 26 )
29 from .winrm import ( 27 from .winrm import (
30 run_powershell, 28 run_powershell,
31 wait_for_winrm, 29 wait_for_winrm,
32 ) 30 )
33 31
34 32
35 SOURCE_ROOT = pathlib.Path(os.path.abspath(__file__)).parent.parent.parent.parent 33 SOURCE_ROOT = pathlib.Path(
36 34 os.path.abspath(__file__)
37 INSTALL_WINDOWS_DEPENDENCIES = (SOURCE_ROOT / 'contrib' / 35 ).parent.parent.parent.parent
38 'install-windows-dependencies.ps1') 36
37 INSTALL_WINDOWS_DEPENDENCIES = (
38 SOURCE_ROOT / 'contrib' / 'install-windows-dependencies.ps1'
39 )
39 40
40 41
41 INSTANCE_TYPES_WITH_STORAGE = { 42 INSTANCE_TYPES_WITH_STORAGE = {
42 'c5d', 43 'c5d',
43 'd2', 44 'd2',
105 { 106 {
106 'CidrIp': '0.0.0.0/0', 107 'CidrIp': '0.0.0.0/0',
107 'Description': 'RDP from entire Internet', 108 'Description': 'RDP from entire Internet',
108 }, 109 },
109 ], 110 ],
110
111 }, 111 },
112 { 112 {
113 'FromPort': 5985, 113 'FromPort': 5985,
114 'ToPort': 5986, 114 'ToPort': 5986,
115 'IpProtocol': 'tcp', 115 'IpProtocol': 'tcp',
117 { 117 {
118 'CidrIp': '0.0.0.0/0', 118 'CidrIp': '0.0.0.0/0',
119 'Description': 'PowerShell Remoting (Windows Remote Management)', 119 'Description': 'PowerShell Remoting (Windows Remote Management)',
120 }, 120 },
121 ], 121 ],
122 } 122 },
123 ], 123 ],
124 }, 124 },
125 } 125 }
126 126
127 127
150 } 150 }
151 '''.strip() 151 '''.strip()
152 152
153 153
154 IAM_INSTANCE_PROFILES = { 154 IAM_INSTANCE_PROFILES = {
155 'ephemeral-ec2-1': { 155 'ephemeral-ec2-1': {'roles': ['ephemeral-ec2-role-1',],}
156 'roles': [
157 'ephemeral-ec2-role-1',
158 ],
159 }
160 } 156 }
161 157
162 158
163 # User Data for Windows EC2 instance. Mainly used to set the password 159 # User Data for Windows EC2 instance. Mainly used to set the password
164 # and configure WinRM. 160 # and configure WinRM.
224 220
225 221
226 class AWSConnection: 222 class AWSConnection:
227 """Manages the state of a connection with AWS.""" 223 """Manages the state of a connection with AWS."""
228 224
229 def __init__(self, automation, region: str, ensure_ec2_state: bool=True): 225 def __init__(self, automation, region: str, ensure_ec2_state: bool = True):
230 self.automation = automation 226 self.automation = automation
231 self.local_state_path = automation.state_path 227 self.local_state_path = automation.state_path
232 228
233 self.prefix = 'hg-' 229 self.prefix = 'hg-'
234 230
255 def rsa_key_fingerprint(p: pathlib.Path): 251 def rsa_key_fingerprint(p: pathlib.Path):
256 """Compute the fingerprint of an RSA private key.""" 252 """Compute the fingerprint of an RSA private key."""
257 253
258 # TODO use rsa package. 254 # TODO use rsa package.
259 res = subprocess.run( 255 res = subprocess.run(
260 ['openssl', 'pkcs8', '-in', str(p), '-nocrypt', '-topk8', 256 [
261 '-outform', 'DER'], 257 'openssl',
258 'pkcs8',
259 '-in',
260 str(p),
261 '-nocrypt',
262 '-topk8',
263 '-outform',
264 'DER',
265 ],
262 capture_output=True, 266 capture_output=True,
263 check=True) 267 check=True,
268 )
264 269
265 sha1 = hashlib.sha1(res.stdout).hexdigest() 270 sha1 = hashlib.sha1(res.stdout).hexdigest()
266 return ':'.join(a + b for a, b in zip(sha1[::2], sha1[1::2])) 271 return ':'.join(a + b for a, b in zip(sha1[::2], sha1[1::2]))
267 272
268 273
269 def ensure_key_pairs(state_path: pathlib.Path, ec2resource, prefix='hg-'): 274 def ensure_key_pairs(state_path: pathlib.Path, ec2resource, prefix='hg-'):
270 remote_existing = {} 275 remote_existing = {}
271 276
272 for kpi in ec2resource.key_pairs.all(): 277 for kpi in ec2resource.key_pairs.all():
273 if kpi.name.startswith(prefix): 278 if kpi.name.startswith(prefix):
274 remote_existing[kpi.name[len(prefix):]] = kpi.key_fingerprint 279 remote_existing[kpi.name[len(prefix) :]] = kpi.key_fingerprint
275 280
276 # Validate that we have these keys locally. 281 # Validate that we have these keys locally.
277 key_path = state_path / 'keys' 282 key_path = state_path / 'keys'
278 key_path.mkdir(exist_ok=True, mode=0o700) 283 key_path.mkdir(exist_ok=True, mode=0o700)
279 284
295 300
296 for f in sorted(os.listdir(key_path)): 301 for f in sorted(os.listdir(key_path)):
297 if not f.startswith('keypair-') or not f.endswith('.pub'): 302 if not f.startswith('keypair-') or not f.endswith('.pub'):
298 continue 303 continue
299 304
300 name = f[len('keypair-'):-len('.pub')] 305 name = f[len('keypair-') : -len('.pub')]
301 306
302 pub_full = key_path / f 307 pub_full = key_path / f
303 priv_full = key_path / ('keypair-%s' % name) 308 priv_full = key_path / ('keypair-%s' % name)
304 309
305 with open(pub_full, 'r', encoding='ascii') as fh: 310 with open(pub_full, 'r', encoding='ascii') as fh:
306 data = fh.read() 311 data = fh.read()
307 312
308 if not data.startswith('ssh-rsa '): 313 if not data.startswith('ssh-rsa '):
309 print('unexpected format for key pair file: %s; removing' % 314 print(
310 pub_full) 315 'unexpected format for key pair file: %s; removing' % pub_full
316 )
311 pub_full.unlink() 317 pub_full.unlink()
312 priv_full.unlink() 318 priv_full.unlink()
313 continue 319 continue
314 320
315 local_existing[name] = rsa_key_fingerprint(priv_full) 321 local_existing[name] = rsa_key_fingerprint(priv_full)
325 print('local key %s does not exist remotely' % name) 331 print('local key %s does not exist remotely' % name)
326 remove_local(name) 332 remove_local(name)
327 del local_existing[name] 333 del local_existing[name]
328 334
329 elif remote_existing[name] != local_existing[name]: 335 elif remote_existing[name] != local_existing[name]:
330 print('key fingerprint mismatch for %s; ' 336 print(
331 'removing from local and remote' % name) 337 'key fingerprint mismatch for %s; '
338 'removing from local and remote' % name
339 )
332 remove_local(name) 340 remove_local(name)
333 remove_remote('%s%s' % (prefix, name)) 341 remove_remote('%s%s' % (prefix, name))
334 del local_existing[name] 342 del local_existing[name]
335 del remote_existing[name] 343 del remote_existing[name]
336 344
354 # SSH public key can be extracted via `ssh-keygen`. 362 # SSH public key can be extracted via `ssh-keygen`.
355 with pub_full.open('w', encoding='ascii') as fh: 363 with pub_full.open('w', encoding='ascii') as fh:
356 subprocess.run( 364 subprocess.run(
357 ['ssh-keygen', '-y', '-f', str(priv_full)], 365 ['ssh-keygen', '-y', '-f', str(priv_full)],
358 stdout=fh, 366 stdout=fh,
359 check=True) 367 check=True,
368 )
360 369
361 pub_full.chmod(0o0600) 370 pub_full.chmod(0o0600)
362 371
363 372
364 def delete_instance_profile(profile): 373 def delete_instance_profile(profile):
365 for role in profile.roles: 374 for role in profile.roles:
366 print('removing role %s from instance profile %s' % (role.name, 375 print(
367 profile.name)) 376 'removing role %s from instance profile %s'
377 % (role.name, profile.name)
378 )
368 profile.remove_role(RoleName=role.name) 379 profile.remove_role(RoleName=role.name)
369 380
370 print('deleting instance profile %s' % profile.name) 381 print('deleting instance profile %s' % profile.name)
371 profile.delete() 382 profile.delete()
372 383
376 387
377 remote_profiles = {} 388 remote_profiles = {}
378 389
379 for profile in iamresource.instance_profiles.all(): 390 for profile in iamresource.instance_profiles.all():
380 if profile.name.startswith(prefix): 391 if profile.name.startswith(prefix):
381 remote_profiles[profile.name[len(prefix):]] = profile 392 remote_profiles[profile.name[len(prefix) :]] = profile
382 393
383 for name in sorted(set(remote_profiles) - set(IAM_INSTANCE_PROFILES)): 394 for name in sorted(set(remote_profiles) - set(IAM_INSTANCE_PROFILES)):
384 delete_instance_profile(remote_profiles[name]) 395 delete_instance_profile(remote_profiles[name])
385 del remote_profiles[name] 396 del remote_profiles[name]
386 397
387 remote_roles = {} 398 remote_roles = {}
388 399
389 for role in iamresource.roles.all(): 400 for role in iamresource.roles.all():
390 if role.name.startswith(prefix): 401 if role.name.startswith(prefix):
391 remote_roles[role.name[len(prefix):]] = role 402 remote_roles[role.name[len(prefix) :]] = role
392 403
393 for name in sorted(set(remote_roles) - set(IAM_ROLES)): 404 for name in sorted(set(remote_roles) - set(IAM_ROLES)):
394 role = remote_roles[name] 405 role = remote_roles[name]
395 406
396 print('removing role %s' % role.name) 407 print('removing role %s' % role.name)
402 for name in sorted(set(IAM_INSTANCE_PROFILES) - set(remote_profiles)): 413 for name in sorted(set(IAM_INSTANCE_PROFILES) - set(remote_profiles)):
403 actual = '%s%s' % (prefix, name) 414 actual = '%s%s' % (prefix, name)
404 print('creating IAM instance profile %s' % actual) 415 print('creating IAM instance profile %s' % actual)
405 416
406 profile = iamresource.create_instance_profile( 417 profile = iamresource.create_instance_profile(
407 InstanceProfileName=actual) 418 InstanceProfileName=actual
419 )
408 remote_profiles[name] = profile 420 remote_profiles[name] = profile
409 421
410 waiter = iamclient.get_waiter('instance_profile_exists') 422 waiter = iamclient.get_waiter('instance_profile_exists')
411 waiter.wait(InstanceProfileName=actual) 423 waiter.wait(InstanceProfileName=actual)
412 print('IAM instance profile %s is available' % actual) 424 print('IAM instance profile %s is available' % actual)
451 def find_image(ec2resource, owner_id, name): 463 def find_image(ec2resource, owner_id, name):
452 """Find an AMI by its owner ID and name.""" 464 """Find an AMI by its owner ID and name."""
453 465
454 images = ec2resource.images.filter( 466 images = ec2resource.images.filter(
455 Filters=[ 467 Filters=[
456 { 468 {'Name': 'owner-id', 'Values': [owner_id],},
457 'Name': 'owner-id', 469 {'Name': 'state', 'Values': ['available'],},
458 'Values': [owner_id], 470 {'Name': 'image-type', 'Values': ['machine'],},
459 }, 471 {'Name': 'name', 'Values': [name],},
460 { 472 ]
461 'Name': 'state', 473 )
462 'Values': ['available'],
463 },
464 {
465 'Name': 'image-type',
466 'Values': ['machine'],
467 },
468 {
469 'Name': 'name',
470 'Values': [name],
471 },
472 ])
473 474
474 for image in images: 475 for image in images:
475 return image 476 return image
476 477
477 raise Exception('unable to find image for %s' % name) 478 raise Exception('unable to find image for %s' % name)
485 """ 486 """
486 existing = {} 487 existing = {}
487 488
488 for group in ec2resource.security_groups.all(): 489 for group in ec2resource.security_groups.all():
489 if group.group_name.startswith(prefix): 490 if group.group_name.startswith(prefix):
490 existing[group.group_name[len(prefix):]] = group 491 existing[group.group_name[len(prefix) :]] = group
491 492
492 purge = set(existing) - set(SECURITY_GROUPS) 493 purge = set(existing) - set(SECURITY_GROUPS)
493 494
494 for name in sorted(purge): 495 for name in sorted(purge):
495 group = existing[name] 496 group = existing[name]
505 506
506 actual = '%s%s' % (prefix, name) 507 actual = '%s%s' % (prefix, name)
507 print('adding security group %s' % actual) 508 print('adding security group %s' % actual)
508 509
509 group_res = ec2resource.create_security_group( 510 group_res = ec2resource.create_security_group(
510 Description=group['description'], 511 Description=group['description'], GroupName=actual,
511 GroupName=actual, 512 )
512 ) 513
513 514 group_res.authorize_ingress(IpPermissions=group['ingress'],)
514 group_res.authorize_ingress(
515 IpPermissions=group['ingress'],
516 )
517 515
518 security_groups[name] = group_res 516 security_groups[name] = group_res
519 517
520 return security_groups 518 return security_groups
521 519
575 if not instance.public_ip_address: 573 if not instance.public_ip_address:
576 time.sleep(2) 574 time.sleep(2)
577 instance.reload() 575 instance.reload()
578 continue 576 continue
579 577
580 print('public IP address for %s: %s' % ( 578 print(
581 instance.id, instance.public_ip_address)) 579 'public IP address for %s: %s'
580 % (instance.id, instance.public_ip_address)
581 )
582 break 582 break
583 583
584 584
585 def remove_ami(ec2resource, image): 585 def remove_ami(ec2resource, image):
586 """Remove an AMI and its underlying snapshots.""" 586 """Remove an AMI and its underlying snapshots."""
601 def wait_for_ssm(ssmclient, instances): 601 def wait_for_ssm(ssmclient, instances):
602 """Wait for SSM to come online for an iterable of instance IDs.""" 602 """Wait for SSM to come online for an iterable of instance IDs."""
603 while True: 603 while True:
604 res = ssmclient.describe_instance_information( 604 res = ssmclient.describe_instance_information(
605 Filters=[ 605 Filters=[
606 { 606 {'Key': 'InstanceIds', 'Values': [i.id for i in instances],},
607 'Key': 'InstanceIds',
608 'Values': [i.id for i in instances],
609 },
610 ], 607 ],
611 ) 608 )
612 609
613 available = len(res['InstanceInformationList']) 610 available = len(res['InstanceInformationList'])
614 wanted = len(instances) 611 wanted = len(instances)
626 623
627 res = ssmclient.send_command( 624 res = ssmclient.send_command(
628 InstanceIds=[i.id for i in instances], 625 InstanceIds=[i.id for i in instances],
629 DocumentName=document_name, 626 DocumentName=document_name,
630 Parameters=parameters, 627 Parameters=parameters,
631 CloudWatchOutputConfig={ 628 CloudWatchOutputConfig={'CloudWatchOutputEnabled': True,},
632 'CloudWatchOutputEnabled': True,
633 },
634 ) 629 )
635 630
636 command_id = res['Command']['CommandId'] 631 command_id = res['Command']['CommandId']
637 632
638 for instance in instances: 633 for instance in instances:
639 while True: 634 while True:
640 try: 635 try:
641 res = ssmclient.get_command_invocation( 636 res = ssmclient.get_command_invocation(
642 CommandId=command_id, 637 CommandId=command_id, InstanceId=instance.id,
643 InstanceId=instance.id,
644 ) 638 )
645 except botocore.exceptions.ClientError as e: 639 except botocore.exceptions.ClientError as e:
646 if e.response['Error']['Code'] == 'InvocationDoesNotExist': 640 if e.response['Error']['Code'] == 'InvocationDoesNotExist':
647 print('could not find SSM command invocation; waiting') 641 print('could not find SSM command invocation; waiting')
648 time.sleep(1) 642 time.sleep(1)
653 if res['Status'] == 'Success': 647 if res['Status'] == 'Success':
654 break 648 break
655 elif res['Status'] in ('Pending', 'InProgress', 'Delayed'): 649 elif res['Status'] in ('Pending', 'InProgress', 'Delayed'):
656 time.sleep(2) 650 time.sleep(2)
657 else: 651 else:
658 raise Exception('command failed on %s: %s' % ( 652 raise Exception(
659 instance.id, res['Status'])) 653 'command failed on %s: %s' % (instance.id, res['Status'])
654 )
660 655
661 656
662 @contextlib.contextmanager 657 @contextlib.contextmanager
663 def temporary_ec2_instances(ec2resource, config): 658 def temporary_ec2_instances(ec2resource, config):
664 """Create temporary EC2 instances. 659 """Create temporary EC2 instances.
709 704
710 config = copy.deepcopy(config) 705 config = copy.deepcopy(config)
711 config['IamInstanceProfile'] = { 706 config['IamInstanceProfile'] = {
712 'Name': 'hg-ephemeral-ec2-1', 707 'Name': 'hg-ephemeral-ec2-1',
713 } 708 }
714 config.setdefault('TagSpecifications', []).append({ 709 config.setdefault('TagSpecifications', []).append(
715 'ResourceType': 'instance', 710 {
716 'Tags': [{'Key': 'Name', 'Value': 'hg-temp-windows'}], 711 'ResourceType': 'instance',
717 }) 712 'Tags': [{'Key': 'Name', 'Value': 'hg-temp-windows'}],
713 }
714 )
718 config['UserData'] = WINDOWS_USER_DATA % password 715 config['UserData'] = WINDOWS_USER_DATA % password
719 716
720 with temporary_ec2_instances(c.ec2resource, config) as instances: 717 with temporary_ec2_instances(c.ec2resource, config) as instances:
721 wait_for_ip_addresses(instances) 718 wait_for_ip_addresses(instances)
722 719
723 print('waiting for Windows Remote Management service...') 720 print('waiting for Windows Remote Management service...')
724 721
725 for instance in instances: 722 for instance in instances:
726 client = wait_for_winrm(instance.public_ip_address, 'Administrator', password) 723 client = wait_for_winrm(
724 instance.public_ip_address, 'Administrator', password
725 )
727 print('established WinRM connection to %s' % instance.id) 726 print('established WinRM connection to %s' % instance.id)
728 instance.winrm_client = client 727 instance.winrm_client = client
729 728
730 yield instances 729 yield instances
731 730
746 """ 745 """
747 # Find existing AMIs with this name and delete the ones that are invalid. 746 # Find existing AMIs with this name and delete the ones that are invalid.
748 # Store a reference to a good image so it can be returned one the 747 # Store a reference to a good image so it can be returned one the
749 # image state is reconciled. 748 # image state is reconciled.
750 images = ec2resource.images.filter( 749 images = ec2resource.images.filter(
751 Filters=[{'Name': 'name', 'Values': [name]}]) 750 Filters=[{'Name': 'name', 'Values': [name]}]
751 )
752 752
753 existing_image = None 753 existing_image = None
754 754
755 for image in images: 755 for image in images:
756 if image.tags is None: 756 if image.tags is None:
757 print('image %s for %s lacks required tags; removing' % ( 757 print(
758 image.id, image.name)) 758 'image %s for %s lacks required tags; removing'
759 % (image.id, image.name)
760 )
759 remove_ami(ec2resource, image) 761 remove_ami(ec2resource, image)
760 else: 762 else:
761 tags = {t['Key']: t['Value'] for t in image.tags} 763 tags = {t['Key']: t['Value'] for t in image.tags}
762 764
763 if tags.get('HGIMAGEFINGERPRINT') == fingerprint: 765 if tags.get('HGIMAGEFINGERPRINT') == fingerprint:
764 existing_image = image 766 existing_image = image
765 else: 767 else:
766 print('image %s for %s has wrong fingerprint; removing' % ( 768 print(
767 image.id, image.name)) 769 'image %s for %s has wrong fingerprint; removing'
770 % (image.id, image.name)
771 )
768 remove_ami(ec2resource, image) 772 remove_ami(ec2resource, image)
769 773
770 return existing_image 774 return existing_image
771 775
772 776
773 def create_ami_from_instance(ec2client, instance, name, description, 777 def create_ami_from_instance(
774 fingerprint): 778 ec2client, instance, name, description, fingerprint
779 ):
775 """Create an AMI from a running instance. 780 """Create an AMI from a running instance.
776 781
777 Returns the ``ec2resource.Image`` representing the created AMI. 782 Returns the ``ec2resource.Image`` representing the created AMI.
778 """ 783 """
779 instance.stop() 784 instance.stop()
780 785
781 ec2client.get_waiter('instance_stopped').wait( 786 ec2client.get_waiter('instance_stopped').wait(
782 InstanceIds=[instance.id], 787 InstanceIds=[instance.id], WaiterConfig={'Delay': 5,}
783 WaiterConfig={ 788 )
784 'Delay': 5,
785 })
786 print('%s is stopped' % instance.id) 789 print('%s is stopped' % instance.id)
787 790
788 image = instance.create_image( 791 image = instance.create_image(Name=name, Description=description,)
789 Name=name, 792
790 Description=description, 793 image.create_tags(
794 Tags=[{'Key': 'HGIMAGEFINGERPRINT', 'Value': fingerprint,},]
791 ) 795 )
792 796
793 image.create_tags(Tags=[
794 {
795 'Key': 'HGIMAGEFINGERPRINT',
796 'Value': fingerprint,
797 },
798 ])
799
800 print('waiting for image %s' % image.id) 797 print('waiting for image %s' % image.id)
801 798
802 ec2client.get_waiter('image_available').wait( 799 ec2client.get_waiter('image_available').wait(ImageIds=[image.id],)
803 ImageIds=[image.id],
804 )
805 800
806 print('image %s available as %s' % (image.id, image.name)) 801 print('image %s available as %s' % (image.id, image.name))
807 802
808 return image 803 return image
809 804
825 'debian-stretch-hvm-x86_64-gp2-2019-09-08-17994', 820 'debian-stretch-hvm-x86_64-gp2-2019-09-08-17994',
826 ) 821 )
827 ssh_username = 'admin' 822 ssh_username = 'admin'
828 elif distro == 'debian10': 823 elif distro == 'debian10':
829 image = find_image( 824 image = find_image(
830 ec2resource, 825 ec2resource, DEBIAN_ACCOUNT_ID_2, 'debian-10-amd64-20190909-10',
831 DEBIAN_ACCOUNT_ID_2,
832 'debian-10-amd64-20190909-10',
833 ) 826 )
834 ssh_username = 'admin' 827 ssh_username = 'admin'
835 elif distro == 'ubuntu18.04': 828 elif distro == 'ubuntu18.04':
836 image = find_image( 829 image = find_image(
837 ec2resource, 830 ec2resource,
869 'MaxCount': 1, 862 'MaxCount': 1,
870 'MinCount': 1, 863 'MinCount': 1,
871 'SecurityGroupIds': [c.security_groups['linux-dev-1'].id], 864 'SecurityGroupIds': [c.security_groups['linux-dev-1'].id],
872 } 865 }
873 866
874 requirements2_path = (pathlib.Path(__file__).parent.parent / 867 requirements2_path = (
875 'linux-requirements-py2.txt') 868 pathlib.Path(__file__).parent.parent / 'linux-requirements-py2.txt'
876 requirements3_path = (pathlib.Path(__file__).parent.parent / 869 )
877 'linux-requirements-py3.txt') 870 requirements3_path = (
871 pathlib.Path(__file__).parent.parent / 'linux-requirements-py3.txt'
872 )
878 with requirements2_path.open('r', encoding='utf-8') as fh: 873 with requirements2_path.open('r', encoding='utf-8') as fh:
879 requirements2 = fh.read() 874 requirements2 = fh.read()
880 with requirements3_path.open('r', encoding='utf-8') as fh: 875 with requirements3_path.open('r', encoding='utf-8') as fh:
881 requirements3 = fh.read() 876 requirements3 = fh.read()
882 877
883 # Compute a deterministic fingerprint to determine whether image needs to 878 # Compute a deterministic fingerprint to determine whether image needs to
884 # be regenerated. 879 # be regenerated.
885 fingerprint = resolve_fingerprint({ 880 fingerprint = resolve_fingerprint(
886 'instance_config': config, 881 {
887 'bootstrap_script': BOOTSTRAP_DEBIAN, 882 'instance_config': config,
888 'requirements_py2': requirements2, 883 'bootstrap_script': BOOTSTRAP_DEBIAN,
889 'requirements_py3': requirements3, 884 'requirements_py2': requirements2,
890 }) 885 'requirements_py3': requirements3,
886 }
887 )
891 888
892 existing_image = find_and_reconcile_image(ec2resource, name, fingerprint) 889 existing_image = find_and_reconcile_image(ec2resource, name, fingerprint)
893 890
894 if existing_image: 891 if existing_image:
895 return existing_image 892 return existing_image
900 wait_for_ip_addresses(instances) 897 wait_for_ip_addresses(instances)
901 898
902 instance = instances[0] 899 instance = instances[0]
903 900
904 client = wait_for_ssh( 901 client = wait_for_ssh(
905 instance.public_ip_address, 22, 902 instance.public_ip_address,
903 22,
906 username=ssh_username, 904 username=ssh_username,
907 key_filename=str(c.key_pair_path_private('automation'))) 905 key_filename=str(c.key_pair_path_private('automation')),
906 )
908 907
909 home = '/home/%s' % ssh_username 908 home = '/home/%s' % ssh_username
910 909
911 with client: 910 with client:
912 print('connecting to SSH server') 911 print('connecting to SSH server')
924 with sftp.open('%s/requirements-py3.txt' % home, 'wb') as fh: 923 with sftp.open('%s/requirements-py3.txt' % home, 'wb') as fh:
925 fh.write(requirements3) 924 fh.write(requirements3)
926 fh.chmod(0o0700) 925 fh.chmod(0o0700)
927 926
928 print('executing bootstrap') 927 print('executing bootstrap')
929 chan, stdin, stdout = ssh_exec_command(client, 928 chan, stdin, stdout = ssh_exec_command(
930 '%s/bootstrap' % home) 929 client, '%s/bootstrap' % home
930 )
931 stdin.close() 931 stdin.close()
932 932
933 for line in stdout: 933 for line in stdout:
934 print(line, end='') 934 print(line, end='')
935 935
936 res = chan.recv_exit_status() 936 res = chan.recv_exit_status()
937 if res: 937 if res:
938 raise Exception('non-0 exit from bootstrap: %d' % res) 938 raise Exception('non-0 exit from bootstrap: %d' % res)
939 939
940 print('bootstrap completed; stopping %s to create %s' % ( 940 print(
941 instance.id, name)) 941 'bootstrap completed; stopping %s to create %s'
942 942 % (instance.id, name)
943 return create_ami_from_instance(ec2client, instance, name, 943 )
944 'Mercurial Linux development environment', 944
945 fingerprint) 945 return create_ami_from_instance(
946 ec2client,
947 instance,
948 name,
949 'Mercurial Linux development environment',
950 fingerprint,
951 )
946 952
947 953
948 @contextlib.contextmanager 954 @contextlib.contextmanager
949 def temporary_linux_dev_instances(c: AWSConnection, image, instance_type, 955 def temporary_linux_dev_instances(
950 prefix='hg-', ensure_extra_volume=False): 956 c: AWSConnection,
957 image,
958 instance_type,
959 prefix='hg-',
960 ensure_extra_volume=False,
961 ):
951 """Create temporary Linux development EC2 instances. 962 """Create temporary Linux development EC2 instances.
952 963
953 Context manager resolves to a list of ``ec2.Instance`` that were created 964 Context manager resolves to a list of ``ec2.Instance`` that were created
954 and are running. 965 and are running.
955 966
977 } 988 }
978 ] 989 ]
979 990
980 # This is not an exhaustive list of instance types having instance storage. 991 # This is not an exhaustive list of instance types having instance storage.
981 # But 992 # But
982 if (ensure_extra_volume 993 if ensure_extra_volume and not instance_type.startswith(
983 and not instance_type.startswith(tuple(INSTANCE_TYPES_WITH_STORAGE))): 994 tuple(INSTANCE_TYPES_WITH_STORAGE)
995 ):
984 main_device = block_device_mappings[0]['DeviceName'] 996 main_device = block_device_mappings[0]['DeviceName']
985 997
986 if main_device == 'xvda': 998 if main_device == 'xvda':
987 second_device = 'xvdb' 999 second_device = 'xvdb'
988 elif main_device == '/dev/sda1': 1000 elif main_device == '/dev/sda1':
989 second_device = '/dev/sdb' 1001 second_device = '/dev/sdb'
990 else: 1002 else:
991 raise ValueError('unhandled primary EBS device name: %s' % 1003 raise ValueError(
992 main_device) 1004 'unhandled primary EBS device name: %s' % main_device
993 1005 )
994 block_device_mappings.append({ 1006
995 'DeviceName': second_device, 1007 block_device_mappings.append(
996 'Ebs': { 1008 {
997 'DeleteOnTermination': True, 1009 'DeviceName': second_device,
998 'VolumeSize': 8, 1010 'Ebs': {
999 'VolumeType': 'gp2', 1011 'DeleteOnTermination': True,
1012 'VolumeSize': 8,
1013 'VolumeType': 'gp2',
1014 },
1000 } 1015 }
1001 }) 1016 )
1002 1017
1003 config = { 1018 config = {
1004 'BlockDeviceMappings': block_device_mappings, 1019 'BlockDeviceMappings': block_device_mappings,
1005 'EbsOptimized': True, 1020 'EbsOptimized': True,
1006 'ImageId': image.id, 1021 'ImageId': image.id,
1017 1032
1018 ssh_private_key_path = str(c.key_pair_path_private('automation')) 1033 ssh_private_key_path = str(c.key_pair_path_private('automation'))
1019 1034
1020 for instance in instances: 1035 for instance in instances:
1021 client = wait_for_ssh( 1036 client = wait_for_ssh(
1022 instance.public_ip_address, 22, 1037 instance.public_ip_address,
1038 22,
1023 username='hg', 1039 username='hg',
1024 key_filename=ssh_private_key_path) 1040 key_filename=ssh_private_key_path,
1041 )
1025 1042
1026 instance.ssh_client = client 1043 instance.ssh_client = client
1027 instance.ssh_private_key_path = ssh_private_key_path 1044 instance.ssh_private_key_path = ssh_private_key_path
1028 1045
1029 try: 1046 try:
1031 finally: 1048 finally:
1032 for instance in instances: 1049 for instance in instances:
1033 instance.ssh_client.close() 1050 instance.ssh_client.close()
1034 1051
1035 1052
1036 def ensure_windows_dev_ami(c: AWSConnection, prefix='hg-', 1053 def ensure_windows_dev_ami(
1037 base_image_name=WINDOWS_BASE_IMAGE_NAME): 1054 c: AWSConnection, prefix='hg-', base_image_name=WINDOWS_BASE_IMAGE_NAME
1055 ):
1038 """Ensure Windows Development AMI is available and up-to-date. 1056 """Ensure Windows Development AMI is available and up-to-date.
1039 1057
1040 If necessary, a modern AMI will be built by starting a temporary EC2 1058 If necessary, a modern AMI will be built by starting a temporary EC2
1041 instance and bootstrapping it. 1059 instance and bootstrapping it.
1042 1060
1098 commands.insert(0, 'Set-MpPreference -DisableRealtimeMonitoring $true') 1116 commands.insert(0, 'Set-MpPreference -DisableRealtimeMonitoring $true')
1099 commands.append('Set-MpPreference -DisableRealtimeMonitoring $false') 1117 commands.append('Set-MpPreference -DisableRealtimeMonitoring $false')
1100 1118
1101 # Compute a deterministic fingerprint to determine whether image needs 1119 # Compute a deterministic fingerprint to determine whether image needs
1102 # to be regenerated. 1120 # to be regenerated.
1103 fingerprint = resolve_fingerprint({ 1121 fingerprint = resolve_fingerprint(
1104 'instance_config': config, 1122 {
1105 'user_data': WINDOWS_USER_DATA, 1123 'instance_config': config,
1106 'initial_bootstrap': WINDOWS_BOOTSTRAP_POWERSHELL, 1124 'user_data': WINDOWS_USER_DATA,
1107 'bootstrap_commands': commands, 1125 'initial_bootstrap': WINDOWS_BOOTSTRAP_POWERSHELL,
1108 'base_image_name': base_image_name, 1126 'bootstrap_commands': commands,
1109 }) 1127 'base_image_name': base_image_name,
1128 }
1129 )
1110 1130
1111 existing_image = find_and_reconcile_image(ec2resource, name, fingerprint) 1131 existing_image = find_and_reconcile_image(ec2resource, name, fingerprint)
1112 1132
1113 if existing_image: 1133 if existing_image:
1114 return existing_image 1134 return existing_image
1129 print('installing Windows features...') 1149 print('installing Windows features...')
1130 run_ssm_command( 1150 run_ssm_command(
1131 ssmclient, 1151 ssmclient,
1132 [instance], 1152 [instance],
1133 'AWS-RunPowerShellScript', 1153 'AWS-RunPowerShellScript',
1134 { 1154 {'commands': WINDOWS_BOOTSTRAP_POWERSHELL.split('\n'),},
1135 'commands': WINDOWS_BOOTSTRAP_POWERSHELL.split('\n'),
1136 },
1137 ) 1155 )
1138 1156
1139 # Reboot so all updates are fully applied. 1157 # Reboot so all updates are fully applied.
1140 # 1158 #
1141 # We don't use instance.reboot() here because it is asynchronous and 1159 # We don't use instance.reboot() here because it is asynchronous and
1143 # a while to stop and we may start trying to interact with the instance 1161 # a while to stop and we may start trying to interact with the instance
1144 # before it has rebooted. 1162 # before it has rebooted.
1145 print('rebooting instance %s' % instance.id) 1163 print('rebooting instance %s' % instance.id)
1146 instance.stop() 1164 instance.stop()
1147 ec2client.get_waiter('instance_stopped').wait( 1165 ec2client.get_waiter('instance_stopped').wait(
1148 InstanceIds=[instance.id], 1166 InstanceIds=[instance.id], WaiterConfig={'Delay': 5,}
1149 WaiterConfig={ 1167 )
1150 'Delay': 5,
1151 })
1152 1168
1153 instance.start() 1169 instance.start()
1154 wait_for_ip_addresses([instance]) 1170 wait_for_ip_addresses([instance])
1155 1171
1156 # There is a race condition here between the User Data PS script running 1172 # There is a race condition here between the User Data PS script running
1157 # and us connecting to WinRM. This can manifest as 1173 # and us connecting to WinRM. This can manifest as
1158 # "AuthorizationManager check failed" failures during run_powershell(). 1174 # "AuthorizationManager check failed" failures during run_powershell().
1159 # TODO figure out a workaround. 1175 # TODO figure out a workaround.
1160 1176
1161 print('waiting for Windows Remote Management to come back...') 1177 print('waiting for Windows Remote Management to come back...')
1162 client = wait_for_winrm(instance.public_ip_address, 'Administrator', 1178 client = wait_for_winrm(
1163 c.automation.default_password()) 1179 instance.public_ip_address,
1180 'Administrator',
1181 c.automation.default_password(),
1182 )
1164 print('established WinRM connection to %s' % instance.id) 1183 print('established WinRM connection to %s' % instance.id)
1165 instance.winrm_client = client 1184 instance.winrm_client = client
1166 1185
1167 print('bootstrapping instance...') 1186 print('bootstrapping instance...')
1168 run_powershell(instance.winrm_client, '\n'.join(commands)) 1187 run_powershell(instance.winrm_client, '\n'.join(commands))
1169 1188
1170 print('bootstrap completed; stopping %s to create image' % instance.id) 1189 print('bootstrap completed; stopping %s to create image' % instance.id)
1171 return create_ami_from_instance(ec2client, instance, name, 1190 return create_ami_from_instance(
1172 'Mercurial Windows development environment', 1191 ec2client,
1173 fingerprint) 1192 instance,
1193 name,
1194 'Mercurial Windows development environment',
1195 fingerprint,
1196 )
1174 1197
1175 1198
1176 @contextlib.contextmanager 1199 @contextlib.contextmanager
1177 def temporary_windows_dev_instances(c: AWSConnection, image, instance_type, 1200 def temporary_windows_dev_instances(
1178 prefix='hg-', disable_antivirus=False): 1201 c: AWSConnection,
1202 image,
1203 instance_type,
1204 prefix='hg-',
1205 disable_antivirus=False,
1206 ):
1179 """Create a temporary Windows development EC2 instance. 1207 """Create a temporary Windows development EC2 instance.
1180 1208
1181 Context manager resolves to the list of ``EC2.Instance`` that were created. 1209 Context manager resolves to the list of ``EC2.Instance`` that were created.
1182 """ 1210 """
1183 config = { 1211 config = {
1203 with create_temp_windows_ec2_instances(c, config) as instances: 1231 with create_temp_windows_ec2_instances(c, config) as instances:
1204 if disable_antivirus: 1232 if disable_antivirus:
1205 for instance in instances: 1233 for instance in instances:
1206 run_powershell( 1234 run_powershell(
1207 instance.winrm_client, 1235 instance.winrm_client,
1208 'Set-MpPreference -DisableRealtimeMonitoring $true') 1236 'Set-MpPreference -DisableRealtimeMonitoring $true',
1237 )
1209 1238
1210 yield instances 1239 yield instances